style: Fix pre-commit issues

Author: aseyboldt

python/nutpie/__init__.py

@@ -1,8 +1,7 @@
+from nutpie import _lib
 from nutpie.compile_pymc import compile_pymc_model
 from nutpie.compile_stan import compile_stan_model
 from nutpie.sample import sample
 
-from nutpie import _lib
-
 __version__: str = _lib.__version__
 __all__ = ["__version__", "sample", "compile_pymc_model", "compile_stan_model"]

python/nutpie/compile_pymc.py

@@ -4,15 +4,15 @@
 from dataclasses import dataclass
 from importlib.util import find_spec
 from math import prod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Literal, Optional
 
 import numpy as np
 import pandas as pd
 from numpy.typing import NDArray
-from nutpie.compiled_pyfunc import from_pyfunc
-from nutpie.sample import CompiledModel
 
 from nutpie import _lib
+from nutpie.compiled_pyfunc import from_pyfunc
+from nutpie.sample import CompiledModel
 
 try:
     from numba.extending import intrinsic
@@ -184,7 +184,7 @@ def _compile_pymc_model_numba(model: "pm.Model", **kwargs) -> CompiledPyMCModel:
     for val in [*logp_fn_pt.get_shared(), *expand_fn_pt.get_shared()]:
         if val.name in shared_data and val not in seen:
             raise ValueError(f"Shared variables must have unique names: {val.name}")
-        shared_data[val.name] = val.get_value().copy()
+        shared_data[val.name] = val.get_value()
         shared_vars[val.name] = val
         seen.add(val)
 
@@ -308,7 +308,7 @@ def logp_fn_jax_grad(x, **shared):
     for val in [*logp_fn_pt.get_shared(), *expand_fn_pt.get_shared()]:
         if val.name in shared_data and val not in seen:
             raise ValueError(f"Shared variables must have unique names: {val.name}")
-        shared_data[val.name] = jax.numpy.asarray(val.get_value().copy())
+        shared_data[val.name] = jax.numpy.asarray(val.get_value())
         shared_vars[val.name] = val
         seen.add(val)
 
@@ -356,8 +356,12 @@ def expand(x, **shared):
 
 
 def compile_pymc_model(
-    model: "pm.Model", *, backend="numba", gradient_backend=None, **kwargs
-) -> CompiledPyMCModel:
+    model: "pm.Model",
+    *,
+    backend: Literal["numba", "jax"] = "numba",
+    gradient_backend: Literal["pytensor", "jax"] | None = None,
+    **kwargs,
+) -> CompiledModel:
     """Compile necessary functions for sampling a pymc model.
 
     Parameters

python/nutpie/compile_stan.py

@@ -8,9 +8,9 @@
 import numpy as np
 import pandas as pd
 from numpy.typing import NDArray
-from nutpie.sample import CompiledModel
 
 from nutpie import _lib
+from nutpie.sample import CompiledModel
 
 
 class _NumpyArrayEncoder(json.JSONEncoder):

python/nutpie/compiled_pyfunc.py

@@ -1,12 +1,12 @@
 import dataclasses
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Callable, List
+from typing import Any, Callable
 
 import numpy as np
-from nutpie.sample import CompiledModel
 
 from nutpie import _lib
+from nutpie.sample import CompiledModel
 
 
 @dataclass(frozen=True)
@@ -15,7 +15,7 @@ class PyFuncModel(CompiledModel):
     _make_expand_func: Callable
     _shared_data: dict[str, Any]
     _n_dim: int
-    _variables: List[_lib.PyVariable]
+    _variables: list[_lib.PyVariable]
     _coords: dict[str, Any]
 
     @property
@@ -66,17 +66,17 @@ def make_expand_func(seed1, seed2, chain):
 
 
 def from_pyfunc(
-    ndim,
-    make_logp_fn,
-    make_expand_fn,
-    expanded_dtypes,
-    expanded_shapes,
-    expanded_names,
+    ndim: int,
+    make_logp_fn: Callable,
+    make_expand_fn: Callable,
+    expanded_dtypes: list[np.dtype],
+    expanded_shapes: list[tuple[int, ...]],
+    expanded_names: list[str],
     *,
-    initial_mean=None,
-    coords=None,
-    dims=None,
-    shared_data=None,
+    initial_mean: np.ndarray | None = None,
+    coords: dict[str, Any] | None = None,
+    dims: dict[str, tuple[str, ...]] | None = None,
+    shared_data: dict[str, Any] | None = None,
 ):
     variables = []
     for name, shape, dtype in zip(
@@ -91,6 +91,13 @@ def from_pyfunc(
             dtype = _lib.ExpandDtype.int64_array(shape)
         variables.append(_lib.PyVariable(name, dtype))
 
+    if coords is None:
+        coords = {}
+    if dims is None:
+        dims = {}
+    if shared_data is None:
+        shared_data = {}
+
     if shared_data is None:
         shared_data = dict()
     return PyFuncModel(

python/nutpie/sample.py

@@ -284,8 +284,11 @@ def in_colab():
         shell = get_ipython().__class__.__name__
         if shell == "ZMQInteractiveShell":  # Jupyter notebook, Spyder or qtconsole
             try:
-                from IPython.display import (HTML, clear_output,  # noqa: F401
-                                             display)
+                from IPython.display import (
+                    HTML,  # noqa: F401
+                    clear_output,  # noqa: F401
+                    display,  # noqa: F401
+                )
 
                 return True
             except ImportError:

tests/test_pymc.py

@@ -5,7 +5,6 @@
 import nutpie
 import nutpie.compile_pymc
 
-
 parameterize_backends = pytest.mark.parametrize(
     "backend, gradient_backend",
     [("numba", None), ("jax", "pytensor"), ("jax", "jax")],
@@ -131,8 +130,8 @@ def test_det(backend, gradient_backend):
 @parameterize_backends
 def test_pymc_model_shared(backend, gradient_backend):
     with pm.Model() as model:
-        mu = pm.MutableData("mu", 0.1)
-        sigma = pm.MutableData("sigma", np.ones(3))
+        mu = pm.Data("mu", 0.1)
+        sigma = pm.Data("sigma", np.ones(3))
         pm.Normal("a", mu=mu, sigma=sigma, shape=3)
 
     compiled = nutpie.compile_pymc_model(
@@ -150,7 +149,26 @@ def test_pymc_model_shared(backend, gradient_backend):
         nutpie.sample(compiled3, chains=1)
 
 
-@parameterize_backends
+@pytest.mark.parametrize(
+    ("backend", "gradient_backend"),
+    [
+        ("numba", None),
+        pytest.param(
+            "jax",
+            "pytensor",
+            marks=pytest.mark.xfail(
+                reason="https://github.com/pymc-devs/pytensor/issues/853"
+            ),
+        ),
+        pytest.param(
+            "jax",
+            "jax",
+            marks=pytest.mark.xfail(
+                reason="https://github.com/pymc-devs/pytensor/issues/853"
+            ),
+        ),
+    ],
+)
 def test_missing(backend, gradient_backend):
     with pm.Model(coords={"obs": range(4)}) as model:
         mu = pm.Normal("mu")