Merge pull request #325 from google-research/siriuz42-2.0-pr

Manually trigger jitting in Flax.

src/timesfm/timesfm_2p5/timesfm_2p5_flax.py

@@ -493,6 +493,8 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
 
   def compile(self, forecast_config: configs.ForecastConfig, **kwargs):
     # Acrobym used during validation.
+    print("Compiling model...")
+
     fc = forecast_config
     if fc.max_context % self.model.p != 0:
       logging.info(
@@ -581,3 +583,14 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
     self.compiled_decode = functools.partial(
       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
+      ),
+    )
+    print("Compiling done.")