From f7f6fa5d999c9624e355436ff722a45187429081 Mon Sep 17 00:00:00 2001 From: Yichen Zhou Date: Fri, 5 Dec 2025 09:46:19 -0800 Subject: [PATCH] Update timesfm_2p5_flax.py Allow compile to delay the jitting run. --- src/timesfm/timesfm_2p5/timesfm_2p5_flax.py | 26 +++++++++++++-------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py index 8da19b8..0f3b150 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py @@ -491,7 +491,12 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5): instance.model = nnx.merge(graph, state) return instance - def compile(self, forecast_config: configs.ForecastConfig, **kwargs): + def compile( + self, + forecast_config: configs.ForecastConfig, + dryrun: bool = True, + **kwargs + ): # Acrobym used during validation. print("Compiling model...") @@ -584,13 +589,14 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5): compiled_decode_kernel, self.forecast_config ) - _ = self.compiled_decode( - self.forecast_config.max_horizon, - jnp.zeros( - (self.global_batch_size, self.forecast_config.max_context), dtype=jnp.float32 - ), - jnp.zeros( - (self.global_batch_size, self.forecast_config.max_context), dtype=jnp.bool - ), - ) + if dryrun: + _ = self.compiled_decode( + self.forecast_config.max_horizon, + jnp.zeros( + (self.global_batch_size, self.forecast_config.max_context), dtype=jnp.float32 + ), + jnp.zeros( + (self.global_batch_size, self.forecast_config.max_context), dtype=jnp.bool + ), + ) print("Compiling done.")