Merge pull request #345 from google-research/siriuz42-patch-2

Update timesfm_2p5_flax.py

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,6 +589,7 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
       compiled_decode_kernel, self.forecast_config
     )
 
+    if dryrun:
       _ = self.compiled_decode(
         self.forecast_config.max_horizon,
         jnp.zeros(