major fix forecast

Author: Kleyt0n

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dynaris"
-version = "0.1.0"
+version = "0.1.01"
 description = "A JAX-powered Python library for Dynamic Linear Models (DLMs)"
 readme = "README.md"
 license = { text = "MIT" }

src/dynaris/__init__.py

@@ -20,7 +20,7 @@
 from dynaris.filters import KalmanFilter, kalman_filter
 from dynaris.smoothers import RTSSmoother, rts_smooth
 
-__version__ = "0.1.0"
+__version__ = "0.1.01"
 
 __all__ = [
     "DLM",

src/dynaris/forecast/forecast.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import functools
 from typing import NamedTuple
 
 import jax
@@ -41,7 +42,7 @@ class ForecastResult(NamedTuple):
 # ---------------------------------------------------------------------------
 
 
-@jax.jit(static_argnums=(2,))
+@functools.partial(jax.jit, static_argnums=(2,))
 def forecast(
     model: StateSpaceModel,
     last_state: GaussianState,