Merge pull request #270 from theislab/perf/ema

joschif · web-flow · commit 9b26bf691e87 · 2025-08-05T15:08:11.000+02:00
perf: jit ema update
diff --git a/src/cellflow/solvers/utils.py b/src/cellflow/solvers/utils.py
@@ -1,6 +1,7 @@
 import jax
 
 
+@jax.jit
 def ema_update(current_model_params: dict, new_model_params: dict, ema: float) -> dict:
     """
     Update parameters using exponential moving average.
diff --git a/tests/solver/test_solver.py b/tests/solver/test_solver.py
@@ -87,7 +87,7 @@ def test_predict_batch(self, dataloader, solver_class):
             atol=1e-1,
             rtol=1e-2,
         )
-        assert diff_nonbatched - diff_batched > 2
+        assert diff_nonbatched - diff_batched > 0.5
 
     @pytest.mark.parametrize("ema", [0.5, 1.0])
     def test_EMA(self, dataloader, ema):

Original file line number	Diff line number	Diff line change
`@@ -87,7 +87,7 @@ def test_predict_batch(self, dataloader, solver_class):`
`87`	`87`	`atol=1e-1,`
`88`	`88`	`rtol=1e-2,`
`89`	`89`	`)`
`90`		`- assert diff_nonbatched - diff_batched > 2`
	`90`	`+ assert diff_nonbatched - diff_batched > 0.5`
`91`	`91`
`92`	`92`	`@pytest.mark.parametrize("ema", [0.5, 1.0])`
`93`	`93`	`def test_EMA(self, dataloader, ema):`