bayinx 0.3.10__py3-none-any.whl → 0.5.3__py3-none-any.whl

bayinx/vi/normalizing_flow.py

@@ -0,0 +1,246 @@
+from typing import Callable, Optional, Self, Tuple
+
+import equinox as eqx
+import jax.flatten_util as jfu
+import jax.numpy as jnp
+import jax.random as jr
+import jax.tree_util as jtu
+from jax.lax import scan
+from jaxtyping import Array, PRNGKeyArray, Scalar
+
+from bayinx.core.flow import FlowLayer
+from bayinx.core.variational import M, Variational
+
+
+class NormalizingFlow(Variational[M]):
+    """
+    An ordered collection of diffeomorphisms that map a base distribution to a
+    normalized approximation of a posterior distribution.
+
+    # Attributes
+    - `dim`: The dimension of the support.
+    - `base`: A base variational distribution.
+    - `flows`: An ordered collection of continuously parameterized diffeomorphisms.
+    """
+    flows: list[FlowLayer]
+    base: Variational[M]
+
+    def __init__(
+        self,
+        base: Variational[M],
+        flows: list[FlowLayer],
+        model: Optional[M] = None,
+        _static: Optional[M] = None,
+        _unflatten: Optional[Callable[[Array], M]] = None
+    ):
+        """
+        Constructs an unoptimized normalizing flow posterior approximation.
+
+        # Parameters
+        - `base`: The base variational distribution.
+        - `flows`: A list of flows.
+        - `model`: A probabilistic `Model` object.
+        """
+        if model is not None:
+            # Partition model
+            params, self._static = eqx.partition(model, model.filter_spec)
+
+            # Flatten params component
+            _, self._unflatten = jfu.ravel_pytree(params)
+        elif _static is not None and _unflatten is not None:
+            self._static = _static
+            self._unflatten = _unflatten
+        else:
+            raise ValueError("Either 'model' or '_static' and '_unflatten' must be specified.")
+
+        self.dim = base.dim
+        self.base = base
+        self.flows = flows
+
+    @property
+    def filter_spec(self) -> Self:
+        # Generate empty specification
+        filter_spec: Self = jtu.tree_map(lambda _: False, self)
+
+        # Specify variational parameters based on each flow's filter spec.
+        filter_spec: Self = eqx.tree_at(
+            lambda vari: vari.flows,
+            filter_spec,
+            replace=[flow.filter_spec for flow in self.flows],
+        )
+
+        return filter_spec
+
+    @eqx.filter_jit
+    def sample(self, n: int, key: PRNGKeyArray = jr.PRNGKey(0)) -> Array:
+        # Sample from the base distribution
+        draws: Array = self.base.sample(n, key)
+
+        assert len(draws.shape) == 2
+
+        # Apply forward transformations
+        for map in self.flows:
+            draws = map.forward(draws)
+
+        assert len(draws.shape) == 2
+
+        return draws
+
+    @eqx.filter_jit
+    def eval(self, draws: Array) -> Array:
+        raise RuntimeError("Evaluating the variational density for a normalizing flow requires an analytic inverse to exist, which many useful flows do not have. Therefore, do not use this method.")
+        return jnp.full(draws.shape[0], jnp.nan)
+
+    @eqx.filter_jit
+    def __eval(self, draws: Array) -> Tuple[Array, Array]:
+        """
+        Evaluate the posterior and variational densities together at the
+        transformed `draws` to avoid extra compute.
+
+        # Parameters
+        - `draws`: Draws from the base variational distribution.
+
+        # Returns
+        The posterior and variational densities as JAX Arrays.
+        """
+        # Evaluate base density
+        variational_evals: Array = self.base.eval(draws)
+
+        # Shape checks
+        assert len(variational_evals.shape) == 1
+        assert len(draws.shape) == 2
+
+        for map in self.flows:
+            # Apply transformation
+            draws, ljas = map.forward_and_adjust(draws)
+            assert len(draws.shape) == 2
+            assert len(ljas.shape) == 1
+
+            # Adjust variational density
+            variational_evals = variational_evals - ljas
+
+        # Evaluate posterior at final variational draws
+        posterior_evals = self.eval_model(draws)
+
+        # Shape checks
+        assert len(posterior_evals.shape) == 1
+        assert len(variational_evals.shape) == 1
+        assert posterior_evals.shape == variational_evals.shape
+
+        return posterior_evals, variational_evals
+
+    @eqx.filter_jit
+    def elbo(self, n: int, batch_size: int, key: PRNGKeyArray = jr.PRNGKey(0)) -> Scalar:
+        dyn, static = eqx.partition(self, self.filter_spec)
+
+        def elbo(dyn: Self, n: int, key: PRNGKeyArray) -> Scalar:
+            self = eqx.combine(dyn, static)
+
+            # Split keys
+            keys = jr.split(key, n // batch_size)
+
+            # Split ELBO calculation into batches
+            def batched_elbo(carry: None, key: PRNGKeyArray) -> Tuple[None, Array]:
+                # Draw from variational distribution
+                draws: Array = self.base.sample(batch_size, key)
+
+                # Evaluate posterior and variational densities
+                batched_post_evals, batched_vari_evals = self.__eval(draws)
+
+                # Compute ELBO estimate
+                batched_elbo_evals: Array = batched_post_evals - batched_vari_evals
+
+                return None, batched_elbo_evals
+
+            elbo_evals = scan(
+                batched_elbo,
+                init=None,
+                xs=keys,
+                length=n // batch_size
+            )[1]
+
+            # Compute average of ELBO estimates
+            return jnp.mean(elbo_evals)
+
+        return elbo(dyn, n, key)
+
+    @eqx.filter_jit
+    def elbo_grad(self, n: int, batch_size: int, key: PRNGKeyArray) -> Self:
+        dyn, static = eqx.partition(self, self.filter_spec)
+
+        # Define ELBO function
+        def elbo(dyn: Self, n: int, key: PRNGKeyArray) -> Scalar:
+            self = eqx.combine(dyn, static)
+
+            # Split key
+            keys = jr.split(key, n // batch_size)
+
+            # Split ELBO calculation into batches
+            def batched_elbo(carry: None, key: PRNGKeyArray) -> Tuple[None, Array]:
+                # Draw from variational distribution
+                draws: Array = self.base.sample(batch_size, key)
+
+                # Evaluate posterior and variational densities
+                batched_post_evals, batched_vari_evals = self.__eval(draws)
+
+                # Compute ELBO estimate
+                batched_elbo_evals: Array = batched_post_evals - batched_vari_evals
+
+                return None, batched_elbo_evals
+
+            elbo_evals = scan(
+                batched_elbo,
+                init=None,
+                xs=keys,
+                length=n // batch_size
+            )[1]
+
+            # Compute average of ELBO estimates
+            return jnp.mean(elbo_evals)
+
+        # Map to its gradient
+        elbo_grad: Callable[
+            [Self, int, PRNGKeyArray], Self
+        ] = eqx.filter_grad(elbo)
+
+        return elbo_grad(dyn, n, key)
+
+    @eqx.filter_jit
+    def elbo_and_grad(self, n: int, batch_size: int, key: PRNGKeyArray) -> Tuple[Scalar, Self]:
+        dyn, static = eqx.partition(self, self.filter_spec)
+
+        def elbo(dyn: Self, n: int, key: PRNGKeyArray) -> Scalar:
+            self = eqx.combine(dyn, static)
+
+            # Split keys
+            keys = jr.split(key, n // batch_size)
+
+            # Split ELBO calculation into batches
+            def batched_elbo(carry: None, key: PRNGKeyArray) -> Tuple[None, Array]:
+                # Draw from variational distribution
+                draws: Array = self.base.sample(batch_size, key)
+
+                # Evaluate posterior and variational densities
+                batched_post_evals, batched_vari_evals = self.__eval(draws)
+
+                # Compute ELBO estimate
+                batched_elbo_evals: Array = batched_post_evals - batched_vari_evals
+
+                return None, batched_elbo_evals
+
+            elbo_evals = scan(
+                batched_elbo,
+                init=None,
+                xs=keys,
+                length=n // batch_size
+            )[1]
+
+            # Compute average of ELBO estimates
+            return jnp.mean(elbo_evals)
+
+        # Map to its value & gradient
+        elbo_and_grad: Callable[
+            [Self, int, PRNGKeyArray], Tuple[Scalar, Self]
+        ] = eqx.filter_value_and_grad(elbo)
+
+        return elbo_and_grad(dyn, n, key)

bayinx/vi/standard.py

@@ -0,0 +1,95 @@
+from typing import Self, Tuple
+
+import equinox as eqx
+import jax.numpy as jnp
+import jax.random as jr
+import jax.tree_util as jtu
+from jax.flatten_util import ravel_pytree
+from jaxtyping import Array, PRNGKeyArray, PyTree, Scalar
+
+from bayinx.core.variational import M, Variational
+from bayinx.dists import normal
+
+
+class Standard(Variational[M]):
+    """
+    A standard normal approximation to a posterior distribution.
+
+    # Attributes
+    - `dim`: The dimension of the support.
+    """
+
+    def __init__(self, model: M):
+        """
+        Constructs a standard normal approximation to a posterior distribution.
+
+        # Parameters
+        - `model`: A probabilistic `Model` object.
+        """
+        # Partition model
+        params, self._static = eqx.partition(model, model.filter_spec)
+
+        # Flatten params component
+        params, self._unflatten = ravel_pytree(params)
+
+        # Store dimension of parameter space
+        self.dim = jnp.size(params)
+
+
+    @eqx.filter_jit
+    def sample(self, n: int, key: PRNGKeyArray = jr.PRNGKey(0)) -> Array:
+        # Sample variational draws
+        draws: Array = jr.normal(key=key, shape=(n, self.dim))
+
+        # Shape checks
+        assert len(draws.shape) == 2
+
+        return draws
+
+    @eqx.filter_jit
+    def eval(self, draws: Array) -> Array:
+        return normal.logprob(
+            x=draws,
+            mu=jnp.array(0.0),
+            sigma=jnp.array(1.0),
+        ).sum(axis=1)
+
+    @property
+    def filter_spec(self):
+        filter_spec = jtu.tree_map(lambda _: False, self)
+
+        return filter_spec
+
+    @eqx.filter_jit
+    def elbo(self, n: int, batch_size: int, key: PRNGKeyArray) -> Scalar:
+        dyn, static = eqx.partition(self, self.filter_spec)
+
+        @eqx.filter_jit
+        def elbo(dyn: Self, n: int, key: PRNGKeyArray) -> Scalar:
+            vari = eqx.combine(dyn, static)
+
+            # Sample draws from variational distribution
+            draws: Array = vari.sample(n, key)
+
+            # Evaluate posterior density for each draw
+            posterior_evals: Array = vari.eval_model(draws)
+
+            # Evaluate variational density for each draw
+            variational_evals: Array = vari.eval(draws)
+
+            # Evaluate ELBO
+            return jnp.mean(posterior_evals - variational_evals)
+
+        return elbo(dyn, n, key)
+
+    @eqx.filter_jit
+    def elbo_grad(self, n: int, batch_size: int, key: PRNGKeyArray) -> Self:
+        raise RuntimeError("Do not use the 'elbo_grad' method for a Standard variational approximation. It has no variational parameters.")
+        return self
+
+    def elbo_and_grad(self, n: int, batch_size: int, key: PRNGKeyArray) -> Tuple[Scalar, PyTree]:
+        """
+        Evaluate the ELBO and its gradient.
+        """
+        raise RuntimeError("Do not use the 'elbo_and_grad' method for a Standard variational approximation. It has no variational parameters.")
+        return self.elbo(n, key), self

bayinx-0.5.3.dist-info/METADATA

@@ -0,0 +1,93 @@
+Metadata-Version: 2.4
+Name: bayinx
+Version: 0.5.3
+Summary: Bayesian Inference with JAX
+Author: Todd Pocuca
+Maintainer: Todd Pocuca
+License-File: LICENSE
+Requires-Python: >=3.13
+Requires-Dist: diffrax>=0.7.0
+Requires-Dist: equinox>=0.13.2
+Requires-Dist: jax>=0.8.0
+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>
+
+Bayinx is an embedded probabilistic programming language in Python, powered by
+[JAX](https://mc-stan.org/). It is heavily inspired by and aims to have
+feature parity with [Stan](https://mc-stan.org/), but extends the types of
+objects you can work with and focuses on normalizing flows variational
+inference for sampling.
+
+
+## Coming From Stan
+
+There are a few differences between the syntax of Bayinx and Stan.
+First, as Bayinx is embedded in Python, model definitions are Pythonic and
+rely on you defining a class that inherits from the `Model` base class:
+
+```py
+class MyModel(Model, init=False):
+    # ...
+```
+
+> Note: Users should specify `init=False` to avoid static type checkers from
+raising irrelevant errors, but more importantly it should remind you that
+you should **NOT** implement your own `__init__` method!
+
+The `data` and `parameters` blocks in Stan are then combined into the attribute
+definitions with Bayinx. For example, if we are modelling a simple normal distribution
+with an unknown mean and variance 1, then we might write:
+
+```py
+class MyModel(Model, init=False):
+    mean: Continuous[Array] = define(shape = ()) # a scalar mean parameter
+    x: Observed[Array] = define(shape = 'n_obs') # a vector of observed values
+
+    # ...
+```
+
+The `model` block in Stan is then defined by implementing the `model` method with Bayinx:
+
+```py
+class MyModel(Model, init=False):
+    mean: Continuous[Array] = define(shape = ())
+    x: Observed[Array] = define(shape = 'n_obs')
+
+    def model(self, target):
+        # Equivalent to 'x ~ normal(mean, 1.0)' in Stan
+        self.x << Normal(self.mean, 1.0)
+
+        return target
+```
+
+Notice that the `~` operator in Stan has been replaced with `<<`, and to reference nodes of a model you must work with `self`.
+
+> Note: Bayinx does not currently have something similar to `transformed data` or `transformed parameters`, however that is likely to be included in a future release.
+
+You can then construct the variational approximation to the posterior:
+
+```py
+import bayinx as byx
+from bayinx.flows import DiagAffine
+import jax.numpy as jnp
+
+# Fit variational approximation
+posterior = byx.Posterior(MyModel, n_obs = 3, x = jnp.array([-1.0, 0.0, 1.0]))
+posterior.configure(flowspecs = [DiagAffine()])
+posterior.fit()
+```
+
+This approximation can then be worked with by sampling nodes:
+
+```py
+mean_draws = posterior.sample('mean', 10000)
+print(mean_draws.mean())
+```
+
+
+## Roadmap
+- [ ] Implement OT-Flow: https://arxiv.org/abs/2006.00104
+- [ ]

bayinx-0.5.3.dist-info/RECORD

@@ -0,0 +1,44 @@
+bayinx/__init__.py,sha256=hhpQ8JM9kzJLUqeV_72ZbUf9gxTFH-SFkq-GyaYscvI,126
+bayinx/ops.py,sha256=QsKrk2tMOxcYXuUidQyiA1e5KbV5WH80XUvM-PV8wZc,2666
+bayinx/posterior.py,sha256=ab7Ubx3BDTkejaNciWDR-J9GEgbygRZPUgc23rz7YOg,6760
+bayinx/constraints/__init__.py,sha256=E9WFI5xPAuVOFTzaLKgG2uV8k5Pho9w0mlmsMYrkSSI,154
+bayinx/constraints/identity.py,sha256=IMR2WHB_GL89IOgkY7sOOScMyMIJkkGpMWpbUVkfOUY,629
+bayinx/constraints/interval.py,sha256=OeM_aED8pZPdhpyrxOUQjg7IXWjXX2GiJGOoocaU7WI,1811
+bayinx/constraints/lower.py,sha256=xXP-vrpQwnBSUN_1f1qYSSKApodfHfveI0f86h6go_k,1517
+bayinx/constraints/upper.py,sha256=RWppD7SKP6KciQt_Wrd_w59vIykm_vaUtE82E9UEcBs,1529
+bayinx/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/core/constraint.py,sha256=w5Roomp-YITFyIcfXpd_P18JYEdYo6yVcNFaYk0MiEE,932
+bayinx/core/context.py,sha256=kSvZFHyveUAkEr-gokZjsG6zp7QqrzcXePFcaTtBFqM,1074
+bayinx/core/distribution.py,sha256=d0dMxzh5NxOv9-IYqV9D-amBcGFLsHsxn8D598IRKEI,990
+bayinx/core/flow.py,sha256=SSGxdLAZLuBUQ8g1D6-QwZPjWuVGTK0fzH6X8CAgFxw,2736
+bayinx/core/model.py,sha256=3FiqNPcLI7QTPhwvH-orYgcKoitozqT0kGxMnoa0dV4,7808
+bayinx/core/node.py,sha256=5nRi3YzKGTUhl6-AOPNkYmYX-xb8DjuOt_XLMVX9k8Q,6134
+bayinx/core/types.py,sha256=of52_tUQurdyfbSzdHjQ0EJUWl8DFEW2Pia9Lr_n3Lk,378
+bayinx/core/utils.py,sha256=_2CxYev5Gu85wMqqqydENmnygvYJh2zB76lBD2-s3y0,3519
+bayinx/core/variational.py,sha256=J7vwGKIulVdaZCoqf7XaRt0Ku7LPzNaqa4xOH_fQ-Nw,4800
+bayinx/dists/__init__.py,sha256=7-nWGyK5obf3lxSlAv5JbzEuQVrRxIC1tU_HseMUtXU,218
+bayinx/dists/bernoulli.py,sha256=vHatAx6FCWCjm-epKu-QqlHjrUN_RXkidGWmzTlurOk,5178
+bayinx/dists/binomial.py,sha256=ap6KN4qW_AbGMZNdz5_8VHtlQYRmIZGlhy9-Z6vulsM,6220
+bayinx/dists/exponential.py,sha256=229Ank-WcJsWhzEY93PvclFWz11VDQLaM1uLEAstf7Y,5705
+bayinx/dists/normal.py,sha256=mnAG2GMjyOdHQ6VAdkDrVJ-d3Rzg2MYqLgRq6-LWwBI,5948
+bayinx/dists/poisson.py,sha256=iBcREtJg0xjXhxKSh0lMTmr7tWAfsO0O3FO-Ft8a6vw,5496
+bayinx/flows/__init__.py,sha256=SozrytzAbeTckrcH_zpiS2bjnhkPsrwZc3nYsUhP4YE,299
+bayinx/flows/diagaffine.py,sha256=cMc2QyC9xipy29FQc7Rzwjo0ga8ajHhxA4hwAIhthSg,3402
+bayinx/flows/fullaffine.py,sha256=b6eXUu72ft_Jf2Kx1wUmdMFGZzn-Bz4J76Hpz7ZKoIw,3594
+bayinx/flows/lowrankaffine.py,sha256=heCWs2iZTTf1BXL6M16yMW8p896Bf15TkMhHv-RaObQ,4769
+bayinx/flows/otflow.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/flows/planar.py,sha256=NLtDcXrT6sVJP4YW0dtgdkins_iEOmvLg0sYsLKUndU,4280
+bayinx/flows/radial.py,sha256=lkIotasV1VteGgHOQxkN2UXRH-ja4oUkU4sXYB52JdQ,6
+bayinx/flows/sylvester.py,sha256=K4Dc8QaKE0WQDRNCMhNdyuBjPnfhljQYTIFz2KBm-h8,6523
+bayinx/nodes/__init__.py,sha256=HWb4Yi-wSd_Fr7AywsAed8EezJXemKagZFRlqypu_qY,141
+bayinx/nodes/continuous.py,sha256=k5GkM6MO6ag4r9wXY48c1j3Hjc0KSalq4xUY9nIz07M,1817
+bayinx/nodes/observed.py,sha256=9SiOHAeLVYSAaFJFuLgF6yiwk8cQgSreSXs5ZvLHaRE,943
+bayinx/nodes/stochastic.py,sha256=2G89o4vC0I0Pw_2zLStaEHJ5OXUSAbVA0hyenyo0aIs,614
+bayinx/vi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/vi/meanfield.py,sha256=N9YeItj9Dk6OpdtBQ4mm-OE4nQXJ-xtLDlbK1t9EP5M,3805
+bayinx/vi/normalizing_flow.py,sha256=LsTZgi6ehDZ4LLuHfnLeaGEo2Z0_CZhRePcoNCXJgyo,8226
+bayinx/vi/standard.py,sha256=r4dydWZMQv5QyCPjBLG3eyU3wczxFUu3pqdA2569eoI,2941
+bayinx-0.5.3.dist-info/METADATA,sha256=ptDkp2X80xOz0pkKajne-pJSllGH2fNHLZzwYYwJt38,2937
+bayinx-0.5.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+bayinx-0.5.3.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.5.3.dist-info/RECORD,,

{bayinx-0.3.10.dist-info → bayinx-0.5.3.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.27.0
+Generator: hatchling 1.28.0
 Root-Is-Purelib: true
 Tag: py3-none-any

bayinx/core/_constraint.py

@@ -1,28 +0,0 @@
-from abc import abstractmethod
-from typing import Tuple
-
-import equinox as eqx
-from jaxtyping import Scalar
-
-from bayinx.core._parameter import Parameter
-
-
-class Constraint(eqx.Module):
-    """
-    Abstract base class for defining parameter constraints.
-    """
-
-    @abstractmethod
-    def constrain(self, x: Parameter) -> Tuple[Parameter, Scalar]:
-        """
-        Applies the constraining transformation to a parameter and computes the log-absolute-Jacobian of the transformation.
-
-        # Parameters
-        - `x`: The unconstrained `Parameter`.
-
-        # Returns
-        A tuple containing:
-            - The constrained `Parameter`.
-            - A scalar Array representing the log-absolute-Jacobian of the transformation.
-        """
-        pass

bayinx/core/_flow.py

@@ -1,80 +0,0 @@
-from abc import abstractmethod
-from typing import Callable, Dict, Self, Tuple
-
-import equinox as eqx
-import jax.tree_util as jtu
-from jaxtyping import Array, Float
-
-
-class Flow(eqx.Module):
-    """
-    An abstract base class for a flow(of a normalizing flow).
-
-    # Attributes
-    - `params`: A dictionary of JAX Arrays representing parameters of the diffeomorphism.
-    - `constraints`: A dictionary of simple functions that constrain their corresponding parameter.
-    """
-
-    params: Dict[str, Float[Array, "..."]]
-    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]]
-
-    @abstractmethod
-    def forward(self, draws: Array) -> Array:
-        """
-        Computes the forward transformation of `draws`.
-        """
-        pass
-
-    @abstractmethod
-    def adjust_density(self, draws: Array) -> Tuple[Array, Array]:
-        """
-        Computes the log-absolute-Jacobian at `draws` and applies the forward transformation.
-
-        # Returns
-            A tuple of JAX Arrays containing the transformed draws and log-absolute-Jacobians.
-        """
-        pass
-
-    # Default filter specification
-    @property
-    @eqx.filter_jit
-    def filter_spec(self):
-        """
-        Generates a filter specification to subset relevant parameters for the flow.
-        """
-        # Generate empty specification
-        filter_spec = jtu.tree_map(lambda _: False, self)
-
-        # Specify JAX Array parameters
-        filter_spec = eqx.tree_at(
-            lambda flow: flow.params,
-            filter_spec,
-            replace=jtu.tree_map(eqx.is_array, self.params),
-        )
-
-        return filter_spec
-
-    @eqx.filter_jit
-    def constrain_params(self: Self):
-        """
-        Constrain `params` to the appropriate domain.
-
-        # Returns
-        A dictionary of transformed JAX Arrays representing the constrained parameters.
-        """
-        t_params = self.params
-
-        for par, map in self.constraints.items():
-            t_params[par] = map(t_params[par])
-
-        return t_params
-
-    @eqx.filter_jit
-    def transform_params(self: Self) -> Dict[str, Array]:
-        """
-        Apply a custom transformation to `params` if needed.
-
-        # Returns
-        A dictionary of transformed JAX Arrays representing the transformed parameters.
-        """
-        return self.constrain_params()

bayinx/core/_model.py

@@ -1,98 +0,0 @@
-from abc import abstractmethod
-from dataclasses import field, fields
-from typing import Any, Self, Tuple
-
-import equinox as eqx
-import jax.numpy as jnp
-import jax.tree as jt
-from jaxtyping import Scalar
-
-from ._constraint import Constraint
-from ._parameter import Parameter
-
-
-def constrain(constraint: Constraint):
-    """Defines constraint metadata."""
-    return field(metadata={"constraint": constraint})
-
-
-class Model(eqx.Module):
-    """
-    An abstract base class used to define probabilistic models.
-
-    Annotate parameter attributes with `Parameter`.
-
-    Include constraints by setting them equal to `constrain(Constraint)`.
-    """
-
-    @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.
-        """
-        # Generate empty specification
-        filter_spec: Self = jt.map(lambda _: False, self)
-
-        for f in fields(self):
-            # Extract attribute from field
-            attr = getattr(self, f.name)
-
-            # Check if attribute is a parameter
-            if isinstance(attr, Parameter):
-                # Update filter specification for parameter
-                filter_spec = eqx.tree_at(
-                    lambda model: getattr(model, f.name),
-                    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.
-
-        # Returns
-        A constrained `Model` object and the adjustment to the posterior.
-        """
-        constrained: Self = self
-        target: Scalar = jnp.array(0.0)
-
-        for f in fields(self):
-            # Extract attribute
-            attr = getattr(self, f.name)
-
-            # Check if constrained parameter
-            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
-                constrained = eqx.tree_at(
-                    lambda model: getattr(model, f.name), constrained, replace=param
-                )
-
-                # Adjust posterior density
-                target += laj
-
-        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).
-
-        # Returns
-        A transformed `Model` object and the adjustment to the posterior.
-        """
-        return self.constrain_params()