Rework jax compilation order

ecmwf-lab · Dec 11, 2024 · 82580b1 · 82580b1
1 parent 91d43c6
commit 82580b1
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/src/ai_models_gencast/model.py b/src/ai_models_gencast/model.py
@@ -217,14 +217,9 @@ def run_forward(
                 LOG.info("Model license: %s", self.ckpt.license)
 
             jax.jit(self._with_configs(run_forward.init))
-            run_forward_jitted = jax.jit(
-                lambda rng, inputs, targets_template, forcings: run_forward.apply(
-                    rng, inputs, targets_template, forcings
-                )
-            )
-            # We also produce a pmapped version for running in parallel.
+
             self.model = xarray_jax.pmap(
-                self._with_params(self._with_configs(self._drop_state(run_forward_jitted))), dim="sample"
+                jax.jit(self._with_params(self._with_configs(self._drop_state(run_forward.apply)))), dim="sample"
             )
 
     def run(self):