bayinx 0.3.19__tar.gz → 0.4.0__tar.gz

bayinx-0.4.0/PKG-INFO

@@ -0,0 +1,47 @@
+Metadata-Version: 2.4
+Name: bayinx
+Version: 0.4.0
+Summary: Bayesian Inference with JAX
+Author: Todd McCready
+Maintainer: Todd McCready
+License-File: LICENSE
+Requires-Python: >=3.12
+Requires-Dist: equinox>=0.11.12
+Requires-Dist: jax>=0.4.38
+Requires-Dist: jaxtyping>=0.2.36
+Requires-Dist: optax>=0.2.4
+Description-Content-Type: text/markdown
+
+# `Bayinx`: <ins>Bay</ins>esian <ins>In</ins>ference with JA<ins>X</ins>
+
+The original aim of this project was to build a PPL in Python that is similar in feel to `Stan` or `Nimble`(where there is a nice declarative syntax for defining the model) and allows for arbitrary models(e.g., ones with discrete parameters that may not be just integers); most of this goal has been moved to [baycian](https://github.com/toddmccready/baycian) for the foreseeable future.
+
+Part of the reason for this move is that Rust's ability to embed a "nice" DSL is comparitively easier due to [Rust macros](https://doc.rust-lang.org/rust-by-example/macros/dsl.html); I can define syntax similar to `Stan` and parse it to valid Rust code. Additionally, the current state of `bayinx` is relatively functional(plus/minus a few things to clean-up and documentation) and it offers enough functionality for one of my other projects: [disize](https://github.com/toddmccready/disize)! I plan to rewrite `disize` in Python with JAX, and `bayinx` makes it easy to handle constraining transformations, filtering for parameters for gradient calculations, etc.
+
+Instead, this project is narrowing on implementing much of `Stan`'s functionality(restricted to continuously parameterized models, point estimation + vi + mcmc, etc) without most of the nice syntax, at least for versions `0.4.#`. Therefore, people will work with `target` directly and return the density like below:
+
+```py
+class NormalDist(Model):
+    x: Parameter[Array] = define(shape = (2,))
+
+    def eval(self, data: Dict[str, Array]):
+        # Constrain parameters
+        self, target = self.constrain_params() # this does nothing for the current model
+
+        # Evaluate x ~ Normal(10.0, 1.0)
+        target += normal.logprob(self.x(), 10.0, 1.0).sum()
+
+        return target
+```
+
+I have ideas for using a context manager and implementing `Node`: `Observed`/`Stochastic` classes that will try and replicate what `baycian` is trying to do, but that is for the future and versions `0.4.#` will retain the functionality needed for `disize`.
+
+
+# TODO
+- For optimization and variational methods offer a way for users to have custom stopping conditions(perhaps stop if a single parameter has converged, etc).
+- Control variates for meanfield VI? Look at https://proceedings.mlr.press/v33/ranganath14.html more closely.
+- Low-rank affine flow?
+- https://arxiv.org/pdf/1803.05649 implement sylvester flows.
+- Learn how to generate documentation.
+- Figure out how to make transform_pars for flows such that there is no performance loss. Noticing some weird behaviour when adding constraints.
+- Look into adaptively tuning ADAM hyperparameters for VI.

bayinx-0.4.0/README.md

@@ -0,0 +1,33 @@
+# `Bayinx`: <ins>Bay</ins>esian <ins>In</ins>ference with JA<ins>X</ins>
+
+The original aim of this project was to build a PPL in Python that is similar in feel to `Stan` or `Nimble`(where there is a nice declarative syntax for defining the model) and allows for arbitrary models(e.g., ones with discrete parameters that may not be just integers); most of this goal has been moved to [baycian](https://github.com/toddmccready/baycian) for the foreseeable future.
+
+Part of the reason for this move is that Rust's ability to embed a "nice" DSL is comparitively easier due to [Rust macros](https://doc.rust-lang.org/rust-by-example/macros/dsl.html); I can define syntax similar to `Stan` and parse it to valid Rust code. Additionally, the current state of `bayinx` is relatively functional(plus/minus a few things to clean-up and documentation) and it offers enough functionality for one of my other projects: [disize](https://github.com/toddmccready/disize)! I plan to rewrite `disize` in Python with JAX, and `bayinx` makes it easy to handle constraining transformations, filtering for parameters for gradient calculations, etc.
+
+Instead, this project is narrowing on implementing much of `Stan`'s functionality(restricted to continuously parameterized models, point estimation + vi + mcmc, etc) without most of the nice syntax, at least for versions `0.4.#`. Therefore, people will work with `target` directly and return the density like below:
+
+```py
+class NormalDist(Model):
+    x: Parameter[Array] = define(shape = (2,))
+
+    def eval(self, data: Dict[str, Array]):
+        # Constrain parameters
+        self, target = self.constrain_params() # this does nothing for the current model
+
+        # Evaluate x ~ Normal(10.0, 1.0)
+        target += normal.logprob(self.x(), 10.0, 1.0).sum()
+
+        return target
+```
+
+I have ideas for using a context manager and implementing `Node`: `Observed`/`Stochastic` classes that will try and replicate what `baycian` is trying to do, but that is for the future and versions `0.4.#` will retain the functionality needed for `disize`.
+
+
+# TODO
+- For optimization and variational methods offer a way for users to have custom stopping conditions(perhaps stop if a single parameter has converged, etc).
+- Control variates for meanfield VI? Look at https://proceedings.mlr.press/v33/ranganath14.html more closely.
+- Low-rank affine flow?
+- https://arxiv.org/pdf/1803.05649 implement sylvester flows.
+- Learn how to generate documentation.
+- Figure out how to make transform_pars for flows such that there is no performance loss. Noticing some weird behaviour when adding constraints.
+- Look into adaptively tuning ADAM hyperparameters for VI.

{bayinx-0.3.19 → bayinx-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.3.19"
+version = "0.4.0"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -10,6 +10,8 @@ dependencies = [
     "jaxtyping>=0.2.36",
     "optax>=0.2.4",
 ]
+authors = [{ name = "Todd McCready" }]
+maintainers = [{ name = "Todd McCready" }]
 
 [build-system]
 requires = ["hatchling"]
@@ -19,7 +21,7 @@ build-backend = "hatchling.build"
 addopts = "-q --benchmark-min-rounds=30 --benchmark-columns=rounds,mean,median,stddev --benchmark-group-by=func"
 
 [tool.bumpversion]
-current_version = "0.3.19"
+current_version = "0.4.0"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

bayinx-0.4.0/src/bayinx/__init__.py

@@ -0,0 +1,3 @@
+from bayinx.core import Model, Parameter, define
+
+__all__ = ["Model", "Parameter", "define"]

bayinx-0.3.19/src/bayinx/constraints/lower.py → bayinx-0.4.0/src/bayinx/constraints.py

@@ -3,9 +3,10 @@ from typing import Tuple
 import equinox as eqx
 import jax.numpy as jnp
 import jax.tree as jt
-from jaxtyping import PyTree, Scalar, ScalarLike
+from jaxtyping import Scalar, ScalarLike
 
 from bayinx.core import Constraint, Parameter
+from bayinx.core._parameter import T
 
 
 class Lower(Constraint):
@@ -18,33 +19,28 @@ class Lower(Constraint):
     def __init__(self, lb: ScalarLike):
         self.lb = jnp.array(lb)
 
-    @eqx.filter_jit
-    def constrain(self, x: Parameter) -> Tuple[Parameter, Scalar]:
+    def constrain(self, param: Parameter[T]) -> Tuple[Parameter[T], Scalar]:
         """
         Enforces a lower bound on the parameter and adjusts the posterior density.
 
         # Parameters
-        - `x`: The unconstrained `Parameter`.
+        - `param`: The unconstrained `Parameter`.
 
         # Parameters
         A tuple containing:
             - A modified `Parameter` with relevant leaves satisfying the constraint.
             - A scalar Array representing the log-absolute-Jacobian of the transformation.
         """
-        # Extract relevant filter specification
-        filter_spec = x.filter_spec
-
         # Extract relevant parameters(all Array)
-        dyn_params, static_params = eqx.partition(x, filter_spec)
+        dyn, static = eqx.partition(param, param.filter_spec)
 
         # Compute density adjustment
-        laj: PyTree = jt.map(jnp.sum, dyn_params)  # pyright: ignore
-        laj: Scalar = jt.reduce(lambda a, b: a + b, laj)
+        laj: Scalar = jt.reduce(lambda a, b: a + b, jt.map(jnp.sum, dyn))
 
         # Compute transformation
-        dyn_params = jt.map(lambda v: jnp.exp(v) + self.lb, dyn_params)
+        dyn = jt.map(lambda v: jnp.exp(v) + self.lb, dyn)
 
         # Combine into full parameter object
-        x = eqx.combine(dyn_params, static_params)
+        param = eqx.combine(dyn, static)
 
-        return x, laj
+        return param, laj

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/core/__init__.py

@@ -1,6 +1,6 @@
 from ._constraint import Constraint
 from ._flow import Flow
-from ._model import Model, constrain
+from ._model import Model, define
 from ._optimization import optimize_model
 from ._parameter import Parameter
 from ._variational import Variational
@@ -9,7 +9,7 @@ __all__ = [
     "Constraint",
     "Flow",
     "Model",
-    "constrain",
+    "define",
     "optimize_model",
     "Parameter",
     "Variational",

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/core/_flow.py

@@ -37,7 +37,6 @@ class Flow(eqx.Module):
 
     # Default filter specification
     @property
-    @eqx.filter_jit
     def filter_spec(self):
         """
         Generates a filter specification to subset relevant parameters for the flow.
@@ -54,7 +53,6 @@ class Flow(eqx.Module):
 
         return filter_spec
 
-    @eqx.filter_jit
     def constrain_params(self: Self):
         """
         Constrain `params` to the appropriate domain.
@@ -69,7 +67,6 @@ class Flow(eqx.Module):
 
         return t_params
 
-    @eqx.filter_jit
     def transform_params(self: Self) -> Dict[str, Array]:
         """
         Apply a custom transformation to `params` if needed.

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/core/_model.py

@@ -1,19 +1,34 @@
 from abc import abstractmethod
 from dataclasses import field, fields
-from typing import Any, Self, Tuple
+from typing import Any, Dict, Optional, Self, Tuple
 
 import equinox as eqx
 import jax.numpy as jnp
 import jax.tree as jt
-from jaxtyping import Scalar
+from jaxtyping import PyTree, Scalar
 
 from ._constraint import Constraint
 from ._parameter import Parameter
 
 
-def constrain(constraint: Constraint):
-    """Defines constraint metadata."""
-    return field(metadata={"constraint": constraint})
+def define(
+    shape: Optional[Tuple[int, ...]] = None,
+    init: Optional[PyTree] = None,
+    constraint: Optional[Constraint] = None
+):
+    """Define a parameter."""
+    metadata: Dict = {}
+    if constraint is not None:
+            metadata["constraint"] = constraint
+
+    if isinstance(shape, Tuple):
+        metadata["shape"] = shape
+    elif isinstance(init, PyTree):
+        metadata["init"] = init
+    else:
+        raise TypeError("Neither 'shape' nor 'init' were given as proper arguments.")
+
+    return field(metadata = metadata)
 
 
 class Model(eqx.Module):
@@ -22,16 +37,32 @@ class Model(eqx.Module):
 
     Annotate parameter attributes with `Parameter`.
 
-    Include constraints by setting them equal to `constrain(Constraint)`.
+    Include constraints by setting them equal to `define(Constraint)`.
     """
 
+    def __new__(cls, *args, **kwargs):
+        obj = super().__new__(cls)
+
+        # Auto-initialize parameters based on `define` metadata
+        for f in fields(cls):
+            if "shape" in f.metadata:
+                # Construct jax Array with correct dimensions
+                setattr(obj, f.name, Parameter(jnp.zeros(f.metadata["shape"])))
+            elif "init" in f.metadata:
+                # Slot in given 'init' object
+                setattr(obj, f.name, Parameter(f.metadata["init"]))
+
+        return obj
+
+    def __init__(self):
+        return self
+
     @abstractmethod
     def eval(self, data: Any) -> Scalar:
         pass
 
     # Default filter specification
     @property
-    @eqx.filter_jit
     def filter_spec(self) -> Self:
         """
         Generates a filter specification to subset relevant parameters for the model.
@@ -49,12 +80,11 @@ class Model(eqx.Module):
                 filter_spec = eqx.tree_at(
                     lambda model: getattr(model, f.name),
                     filter_spec,
-                    replace=attr.filter_spec,
+                    replace=attr.filter_spec
                 )
 
         return filter_spec
 
-    @eqx.filter_jit
     def constrain_params(self) -> Tuple[Self, Scalar]:
         """
         Constrain parameters to the appropriate domain.
@@ -70,14 +100,14 @@ class Model(eqx.Module):
             attr = getattr(self, f.name)
 
             # Check if constrained parameter
-            if isinstance(attr, Parameter) and "constraint" in f.metadata:
+            if isinstance(attr, Parameter) and ("constraint" in f.metadata):
                 param = attr
                 constraint = f.metadata["constraint"]
 
                 # Apply constraint
                 param, laj = constraint.constrain(param)
 
-                # Update parameters for constrained model
+                # Update parameters for constrained model at same node
                 constrained = eqx.tree_at(
                     lambda model: getattr(model, f.name), constrained, replace=param
                 )
@@ -87,7 +117,6 @@ class Model(eqx.Module):
 
         return constrained, target
 
-    @eqx.filter_jit
     def transform_params(self) -> Tuple[Self, Scalar]:
         """
         Apply a custom transformation to parameters if needed(defaults to constrained parameters).

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/core/_optimization.py

@@ -10,6 +10,8 @@ from optax import GradientTransformation, OptState, Schedule
 from ._model import Model
 
 M = TypeVar("M", bound=Model)
+
+
 @eqx.filter_jit
 def optimize_model(
     model: M,
@@ -39,6 +41,7 @@ def optimize_model(
 
         # Evaluate posterior
         return model.eval(data)
+
     eval_grad: Callable[[M], M] = eqx.filter_jit(eqx.filter_grad(eval))
 
     # Construct scheduler

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/core/_parameter.py

@@ -5,6 +5,8 @@ import jax.tree as jt
 from jaxtyping import PyTree
 
 T = TypeVar("T", bound=PyTree)
+
+
 class Parameter(eqx.Module, Generic[T]):
     """
     A container for a parameter of a `Model`.
@@ -26,19 +28,18 @@ class Parameter(eqx.Module, Generic[T]):
 
     # Default filter specification
     @property
-    @eqx.filter_jit
     def filter_spec(self) -> Self:
         """
-        Generates a filter specification to filter out static parameters.
+        Generates a filter specification to filter for dynamic parameters.
         """
         # Generate empty specification
-        filter_spec = jt.map(lambda _: False, self)
+        filter_spec: Self = jt.map(lambda _: False, self)
 
-        # Specify Array leaves
+        # Specify Array-like leaves
         filter_spec = eqx.tree_at(
-            lambda params: params.vals,
+            lambda param: param.vals,
             filter_spec,
-            replace=jt.map(eqx.is_array_like, self.vals),
+            replace=jt.map(eqx.is_inexact_array_like, self.vals),
         )
 
         return filter_spec

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/core/_variational.py

@@ -22,11 +22,11 @@ class Variational(eqx.Module, Generic[M]):
 
     # Attributes
     - `_unflatten`: A function to transform draws from the variational distribution back to a `Model`.
-    - `_constraints`: The static component of a partitioned `Model` used to initialize the `Variational` object.
+    - `_static`: The static component of a partitioned `Model` used to initialize the `Variational` object.
     """
 
     _unflatten: Callable[[Array], M]
-    _constraints: M
+    _static: M
 
     @abstractmethod
     def filter_spec(self):
@@ -69,7 +69,7 @@ class Variational(eqx.Module, Generic[M]):
         model: M = self._unflatten(draw)
 
         # Combine with constraints
-        model: M = eqx.combine(model, self._constraints)
+        model: M = eqx.combine(model, self._static)
 
         return model
 
@@ -89,7 +89,6 @@ class Variational(eqx.Module, Generic[M]):
         # Evaluate posterior density
         return model.eval(data)
 
-    # TODO: get rid of this and put it all in each vari's methods, forgot abt discrete parameters :V
     @eqx.filter_jit
     def fit(
         self,
@@ -116,11 +115,9 @@ class Variational(eqx.Module, Generic[M]):
         dyn, static = eqx.partition(self, self.filter_spec)
 
         # Construct scheduler
-        schedule: Schedule = opx.warmup_cosine_decay_schedule(
-            init_value=1e-16,
-            peak_value=learning_rate,
-            warmup_steps=int(max_iters / 10),
-            decay_steps=max_iters - int(max_iters / 10),
+        schedule: Schedule = opx.cosine_decay_schedule(
+            init_value=learning_rate,
+            decay_steps=max_iters,
         )
 
         # Initialize optimizer
@@ -175,7 +172,7 @@ class Variational(eqx.Module, Generic[M]):
         return eqx.combine(dyn, static)
 
     @eqx.filter_jit
-    def posterior_predictive(
+    def _posterior_predictive(
         self,
         func: Callable[[M, Any], Array],
         n: int,

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/dists/censored/posnormal/r.py

@@ -111,6 +111,11 @@ def sample(
 
     # Construct draws
     draws = jr.uniform(key, shape)
-    draws = mu + sigma * ndtri(normal.cdf(-mu/sigma, 0.0, 1.0) + draws * normal.cdf(mu/sigma, 0.0, 1.0))
+    draws = mu + sigma * ndtri(
+        normal.cdf(-mu / sigma, 0.0, 1.0) + draws * normal.cdf(mu / sigma, 0.0, 1.0)
+    )
+
+    # Censor draws
+    draws.at[censor <= draws].set(censor)
 
     return draws

bayinx-0.4.0/src/bayinx/dists/negbinom3.py

@@ -0,0 +1,113 @@
+import jax.numpy as jnp
+from jax.scipy.special import gammaln
+from jaxtyping import Array, ArrayLike, Float, UInt
+
+__PI = 3.141592653589793
+
+
+def __binom(x, y):
+    """
+    Helper function for the Binomial coefficient.
+    """
+    return jnp.exp(gammaln(x + 1) - gammaln(y + 1) - gammaln(x - y + 1))
+
+
+def prob(
+    x: UInt[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    phi: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The probability mass function (PMF) for a (mean-inverse overdispersion parameterized) Negatvie Binomial distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the PMF.
+    - `mu`: The mean.
+    - `phi`: The inverse overdispersion.
+
+    # Returns
+    The PMF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `phi`.
+    """
+    # Cast to Array
+    x, mu, phi = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(phi)
+
+    return jnp.exp(logprob(x, mu, phi))
+
+
+def logprob(
+    x: UInt[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    phi: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    """
+    The log-transformed probability mass function (PMF) for a (mean-inverse overdispersion parameterized) Negatvie Binomial distribution.
+
+    # Parameters
+    - `x`: Where to evaluate the log PMF.
+    - `mu`: The mean.
+    - `phi`: The inverse overdispersion.
+
+    # Returns
+    The log PMF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `phi`.
+    """
+    # Cast to Array
+    x, mu, phi = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(phi)
+
+    # Evaluate log PMF
+    evals: Array = jnp.where(
+        x < 0,
+        -jnp.inf,
+        (
+            gammaln(x + phi)
+            - gammaln(x + 1)
+            - gammaln(phi)
+            + x * (jnp.log(mu) - jnp.log(mu + phi))
+            + phi * (jnp.log(phi) - jnp.log(mu + phi))
+        ),
+    )
+
+    return evals
+
+
+def cdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    return jnp.array(1.0)
+
+
+def logcdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    return jnp.array(1.0)
+
+
+def ccdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    return jnp.array(1.0)
+
+
+def logccdf(
+    x: Float[ArrayLike, "..."],
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+) -> Float[Array, "..."]:
+    # Cast to Array
+    x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
+
+    return jnp.array(1.0)

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/dists/posnormal.py

@@ -1,5 +1,7 @@
 import jax.numpy as jnp
-from jaxtyping import Array, ArrayLike, Float
+import jax.random as jr
+from jax.scipy.special import ndtri
+from jaxtyping import Array, ArrayLike, Float, Key
 
 from bayinx.dists import normal
 
@@ -251,3 +253,38 @@ def logccdf(
     evals = jnp.where(non_negative, A - B, jnp.asarray(0.0))
 
     return evals
+
+
+def sample(
+    n: int,
+    mu: Float[ArrayLike, "..."],
+    sigma: Float[ArrayLike, "..."],
+    key: Key = jr.PRNGKey(0),
+) -> Float[Array, "..."]:
+    """
+    Sample from a positive Normal distribution.
+
+    # Parameters
+    - `n`: Number of draws to sample per-parameter.
+    - `mu`: The mean.
+    - `sigma`: The standard deviation.
+
+    # Returns
+    Draws from a positive Normal distribution. The output will have the shape of (n,) + the broadcasted shapes of `mu` and `sigma`.
+    """
+    # Cast to Array
+    mu, sigma = (
+        jnp.asarray(mu),
+        jnp.asarray(sigma),
+    )
+
+    # Derive shape
+    shape = (n,) + jnp.broadcast_shapes(mu.shape, sigma.shape)
+
+    # Construct draws
+    draws = jr.uniform(key, shape)
+    draws = mu + sigma * ndtri(
+        normal.cdf(-mu / sigma, 0.0, 1.0) + draws * normal.cdf(mu / sigma, 0.0, 1.0)
+    )
+
+    return draws

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/dists/uniform.py

@@ -83,8 +83,12 @@ def ulogprob(
 
     return jnp.zeros(jnp.broadcast_shapes(x.shape, lb.shape, ub.shape))
 
+
 def sample(
-    n: int, lb: Float[ArrayLike, "..."], ub: Float[ArrayLike, "..."], key: Key = jr.PRNGKey(0),
+    n: int,
+    lb: Float[ArrayLike, "..."],
+    ub: Float[ArrayLike, "..."],
+    key: Key = jr.PRNGKey(0),
 ) -> Float[Array, "..."]:
     """
     Sample from a Uniform distribution.
@@ -104,6 +108,6 @@ def sample(
     shape = (n,) + jnp.broadcast_shapes(lb.shape, ub.shape)
 
     # Construct draws
-    draws = jr.uniform(key, shape, minval = lb, maxval = ub)
+    draws = jr.uniform(key, shape, minval=lb, maxval=ub)
 
     return draws

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/mhx/vi/flows/fullaffine.py

@@ -33,8 +33,6 @@ class FullAffine(Flow):
         if dim == 1:
             self.constraints = {}
         else:
-
-            @eqx.filter_jit
             def constrain_scale(scale: Array):
                 # Extract diagonal and apply exponential
                 diag: Array = jnp.exp(jnp.diag(scale))

{bayinx-0.3.19 → bayinx-0.4.0}/src/bayinx/mhx/vi/meanfield.py

@@ -34,7 +34,7 @@ class MeanField(Variational, Generic[M]):
         - `init_log_std`: The initial log-transformed standard deviation of the Gaussian approximation.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, model.filter_spec)
+        params, self._static = eqx.partition(model, model.filter_spec)
 
         # Flatten params component
         params, self._unflatten = ravel_pytree(params)
@@ -44,7 +44,6 @@ class MeanField(Variational, Generic[M]):
         self.log_std = jnp.full(params.size, init_log_std, params.dtype)
 
     @property
-    @eqx.filter_jit
     def filter_spec(self):
         # Generate empty specification
         filter_spec = jtu.tree_map(lambda _: False, self)
@@ -67,8 +66,7 @@ class MeanField(Variational, Generic[M]):
     def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:
         # Sample variational draws
         draws: Array = (
-            jr.normal(key=key, shape=(n, self.mean.size))
-            * jnp.exp(self.log_std)
+            jr.normal(key=key, shape=(n, self.mean.size)) * jnp.exp(self.log_std)
             + self.mean
         )
 
@@ -108,10 +106,9 @@ class MeanField(Variational, Generic[M]):
     def elbo_grad(self, n: int, key: Key, data: Any = None) -> Self:
         dyn, static = eqx.partition(self, self.filter_spec)
 
-        @eqx.filter_grad
         @eqx.filter_jit
+        @eqx.filter_grad
         def elbo_grad(dyn: Self, n: int, key: Key, data: Any = None):
-            # Combine
             vari = eqx.combine(dyn, static)
 
             # Sample draws from variational distribution