pass rng from cellflow.predict

MUCDK · MUCDK · commit b9bc6c28c922 · 2025-05-26T09:04:25.000+02:00
diff --git a/src/cellflow/model/_cellflow.py b/src/cellflow/model/_cellflow.py
@@ -619,7 +619,7 @@ def predict(
         src = batch["source"]
         condition = batch.get("condition", None)
         out = jax.tree.map(
-            functools.partial(self.solver.predict, **kwargs),
+            functools.partial(self.solver.predict, rng=rng, **kwargs),
             src,
             condition,  # type: ignore[attr-defined]
         )

Original file line number	Diff line number	Diff line change
`@@ -619,7 +619,7 @@ def predict(`
`619`	`619`	`src = batch["source"]`
`620`	`620`	`condition = batch.get("condition", None)`
`621`	`621`	`out = jax.tree.map(`
`622`		`- functools.partial(self.solver.predict, **kwargs),`
	`622`	`+ functools.partial(self.solver.predict, rng=rng, **kwargs),`
`623`	`623`	`src,`
`624`	`624`	`condition, # type: ignore[attr-defined]`
`625`	`625`	`)`