freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Manually trigger jitting in Flax.

Browse Source
This commit is contained in:
siriuz42
2025-10-08 16:32:41 +00:00
parent a21b1d5b25
commit 9080939cb8
1 changed files with 13 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/timesfm/timesfm_2p5/timesfm_2p5_flax.py
+13
View File
@@ -493,6 +493,8 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
def compile(self, forecast_config: configs.ForecastConfig, **kwargs): def compile(self, forecast_config: configs.ForecastConfig, **kwargs):
# Acrobym used during validation. # Acrobym used during validation.
print("Compiling model...")
fc = forecast_config fc = forecast_config
if fc.max_context % self.model.p != 0: if fc.max_context % self.model.p != 0:
logging.info( logging.info(
@@ -581,3 +583,14 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
self.compiled_decode = functools.partial( self.compiled_decode = functools.partial(
compiled_decode_kernel, self.forecast_config 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.")