Skip to content

Commit 29cd996

Browse files
committed
Changed backend arg defaults to strings
1 parent 8b9521c commit 29cd996

File tree

1 file changed

+4
-5
lines changed

1 file changed

+4
-5
lines changed

python/nutpie/compile_pymc.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -354,8 +354,8 @@ def expand(x, **shared):
354354
def compile_pymc_model(
355355
model: "pm.Model",
356356
*,
357-
backend: Literal["numba", "jax"] | None = None,
358-
gradient_backend: Literal["pytensor", "jax"] | None = None,
357+
backend: Literal["numba", "jax"] = "numba",
358+
gradient_backend: Literal["pytensor", "jax"] = "pytensor",
359359
**kwargs,
360360
) -> CompiledModel:
361361
"""Compile necessary functions for sampling a pymc model.
@@ -384,10 +384,9 @@ def compile_pymc_model(
384384
"and restart your kernel in case you are in an interactive session."
385385
)
386386

387-
if backend is None:
388-
backend = "numba"
389-
390387
if backend.lower() == "numba":
388+
if gradient_backend == "jax":
389+
raise ValueError("Gradient backend cannot be jax when using numba backend")
391390
return _compile_pymc_model_numba(model, **kwargs)
392391
elif backend.lower() == "jax":
393392
return _compile_pymc_model_jax(

0 commit comments

Comments
 (0)