bayinx 0.2.6__tar.gz → 0.2.9__tar.gz

{bayinx-0.2.6 → bayinx-0.2.9}/PKG-INFO

@@ -1,14 +1,12 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.2.6
+Version: 0.2.9
 Summary: Bayesian Inference with JAX
 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
-Requires-Dist: pytest-benchmark>=5.1.0
-Requires-Dist: pytest>=8.3.5
 Description-Content-Type: text/markdown
 
 # <ins>Bay</ins>esian <ins>In</ins>ference with JA<ins>X</ins>
@@ -23,7 +21,7 @@ In the short-term, I'm going to focus on:
 In the long-term, I'm going to focus on:
 1) How to get `Stan`-like declarative syntax in Python with minimal syntactic overhead(to get as close as possible to statements like `X ~ Normal(mu, 1)`), while also allowing users to work with `target` directly when needed(same as `Stan` does).
 2) How to make working with the posterior as easy as possible.
-    - That's a vague goal but practically it means how to easily evaluate statements like $P(\theta \in [-1, 1] | \mathcal{M})$, or set up contrasts and evaluate $P(\mu_1 - \mu_2 > 0 | \mathcal{M})$, or simulate the posterior predictive to generate plots, etc.
+    - That's a vague goal but practically it means how to easily evaluate statements like $P(\theta \in [-1, 1] | \mathcal{D}, \mathcal{M})$, or set up contrasts and evaluate $P(\mu_1 - \mu_2 > 0 | \mathcal{D}, \mathcal{M})$, or simulate the posterior predictive to generate plots, etc.
 
 Although this is somewhat separate from the goals of the project, if this does pan out how I'm invisioning it I'd like an R formula-like syntax to shorten model construction in scenarios where the model is just a GLMM or similar(think `brms`).
 

{bayinx-0.2.6 → bayinx-0.2.9}/README.md

@@ -10,7 +10,7 @@ In the short-term, I'm going to focus on:
 In the long-term, I'm going to focus on:
 1) How to get `Stan`-like declarative syntax in Python with minimal syntactic overhead(to get as close as possible to statements like `X ~ Normal(mu, 1)`), while also allowing users to work with `target` directly when needed(same as `Stan` does).
 2) How to make working with the posterior as easy as possible.
-    - That's a vague goal but practically it means how to easily evaluate statements like $P(\theta \in [-1, 1] | \mathcal{M})$, or set up contrasts and evaluate $P(\mu_1 - \mu_2 > 0 | \mathcal{M})$, or simulate the posterior predictive to generate plots, etc.
+    - That's a vague goal but practically it means how to easily evaluate statements like $P(\theta \in [-1, 1] | \mathcal{D}, \mathcal{M})$, or set up contrasts and evaluate $P(\mu_1 - \mu_2 > 0 | \mathcal{D}, \mathcal{M})$, or simulate the posterior predictive to generate plots, etc.
 
 Although this is somewhat separate from the goals of the project, if this does pan out how I'm invisioning it I'd like an R formula-like syntax to shorten model construction in scenarios where the model is just a GLMM or similar(think `brms`).
 

{bayinx-0.2.6 → bayinx-0.2.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.2.6"
+version = "0.2.9"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -9,8 +9,6 @@ dependencies = [
     "jax>=0.4.38",
     "jaxtyping>=0.2.36",
     "optax>=0.2.4",
-    "pytest>=8.3.5",
-    "pytest-benchmark>=5.1.0",
 ]
 
 [build-system]
@@ -19,3 +17,6 @@ build-backend = "hatchling.build"
 
 [tool.pytest.ini_options]
 addopts = "-q --benchmark-min-rounds=30 --benchmark-columns=rounds,mean,median,stddev --benchmark-group-by=func"
+
+[dependency-groups]
+dev = ["pytest>=8.3.5", "pytest-benchmark>=5.1.0"]

bayinx-0.2.9/src/bayinx/core/flow.py

@@ -0,0 +1,80 @@
+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
+
+from bayinx.core.utils import __MyMeta
+
+
+class Flow(eqx.Module, metaclass=__MyMeta):
+    """
+    A superclass used to define continuously parameterized diffeomorphisms for normalizing flows.
+
+    # Attributes
+    - `pars`: A dictionary of JAX Arrays representing parameters of the diffeomorphism.
+    - `constraints`: A dictionary of 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-determinant of the Jacobian at `draws` and applies the forward transformation.
+
+        # Returns
+        A tuple of JAX Arrays containing the log-absolute-determinant of the Jacobians and transformed draws.
+        """
+        pass
+
+    # Default filter specification
+    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_pars(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_pars(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_pars()

{bayinx-0.2.6 → bayinx-0.2.9}/src/bayinx/core/model.py

@@ -2,6 +2,7 @@ from abc import abstractmethod
 from typing import Any, Callable, Dict
 
 import equinox as eqx
+import jax.tree_util as jtu
 from jaxtyping import Array, Scalar
 
 from bayinx.core.utils import __MyMeta
@@ -23,6 +24,23 @@ class Model(eqx.Module, metaclass=__MyMeta):
     def eval(self, data: Any) -> Scalar:
         pass
 
+    # Default filter specification
+    def filter_spec(self):
+        """
+        Generates a filter specification to subset relevant parameters for the model.
+        """
+        # Generate empty specification
+        filter_spec = jtu.tree_map(lambda _: False, self)
+
+        # Specify JAX Array parameters
+        filter_spec = eqx.tree_at(
+            lambda model: model.params,
+            filter_spec,
+            replace=jtu.tree_map(eqx.is_array, self.params),
+        )
+
+        return filter_spec
+
     def __init_subclass__(cls):
         # Add constrain method
         def constrain_pars(self: Model) -> Dict[str, Array]:
@@ -43,6 +61,7 @@ class Model(eqx.Module, metaclass=__MyMeta):
 
         # Add transform_pars method if not present
         if not callable(getattr(cls, "transform_pars", None)):
+
             def transform_pars(self: Model) -> Dict[str, Array]:
                 """
                 Apply a custom transformation to `params` if needed.

{bayinx-0.2.6 → bayinx-0.2.9}/src/bayinx/core/variational.py

@@ -104,6 +104,9 @@ class Variational(eqx.Module):
             - `var_draws`: Number of variational draws to draw each iteration.
             - `key`: A PRNG key.
             """
+            # Partition variational
+            dyn, static = eqx.partition(self, self.filter_spec())
+
             # Construct scheduler
             schedule: Schedule = opx.cosine_decay_schedule(
                 init_value=learning_rate, decay_steps=max_iters
@@ -113,20 +116,20 @@ class Variational(eqx.Module):
             optim: GradientTransformation = opx.chain(
                 opx.scale(-1.0), opx.nadam(schedule)
             )
-            opt_state: OptState = optim.init(eqx.filter(self, self.filter_spec()))
+            opt_state: OptState = optim.init(dyn)
 
             # Optimization loop helper functions
             @eqx.filter_jit
             def condition(state: Tuple[Self, OptState, Scalar, Key]):
                 # Unpack iteration state
-                self, opt_state, i, key = state
+                dyn, opt_state, i, key = state
 
                 return i < max_iters
 
             @eqx.filter_jit
             def body(state: Tuple[Self, OptState, Scalar, Key]):
                 # Unpack iteration state
-                self, opt_state, i, key = state
+                dyn, opt_state, i, key = state
 
                 # Update iteration
                 i = i + 1
@@ -134,26 +137,30 @@ class Variational(eqx.Module):
                 # Update PRNG key
                 key, _ = jr.split(key)
 
+                # Combine variational
+                vari = eqx.combine(dyn, static)
+
                 # Compute gradient of the ELBO
-                updates: PyTree = self.elbo_grad(var_draws, key, data)
+                updates: PyTree = vari.elbo_grad(var_draws, key, data)
 
                 # Compute updates
                 updates, opt_state = optim.update(
-                    updates, opt_state, eqx.filter(self, self.filter_spec())
+                    updates, opt_state, eqx.filter(dyn, dyn.filter_spec())
                 )
 
                 # Update variational distribution
-                self: Self = eqx.apply_updates(self, updates)
+                dyn = eqx.apply_updates(dyn, updates)
 
-                return self, opt_state, i, key
+                return dyn, opt_state, i, key
 
             # Run optimization loop
-            self = lax.while_loop(
+            dyn = lax.while_loop(
                 cond_fun=condition,
                 body_fun=body,
-                init_val=(self, opt_state, jnp.array(0, jnp.uint32), key),
+                init_val=(dyn, opt_state, jnp.array(0, jnp.uint32), key),
             )[0]
 
-            return self
+            # Return optimized variational
+            return eqx.combine(dyn, static)
 
         cls.fit = eqx.filter_jit(fit)

bayinx-0.2.9/src/bayinx/mhx/__init__.py

@@ -0,0 +1 @@
+

bayinx-0.2.9/src/bayinx/mhx/vi/__init__.py

@@ -0,0 +1,3 @@
+from bayinx.mhx.vi.meanfield import MeanField as MeanField
+from bayinx.mhx.vi.normalizing_flow import NormalizingFlow as NormalizingFlow
+from bayinx.mhx.vi.standard import Standard as Standard

bayinx-0.2.9/src/bayinx/mhx/vi/flows/__init__.py

@@ -0,0 +1,3 @@
+from bayinx.mhx.vi.flows.affine import Affine as Affine
+from bayinx.mhx.vi.flows.planar import Planar as Planar
+from bayinx.mhx.vi.flows.radial import Radial as Radial

{bayinx-0.2.6/src/bayinx/machinery/variational → bayinx-0.2.9/src/bayinx/mhx/vi}/flows/affine.py

@@ -19,9 +19,7 @@ class Affine(Flow):
     """
 
     params: Dict[str, Float[Array, "..."]]
-    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]] = (
-        eqx.field(static=True)
-    )
+    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]]
 
     def __init__(self, dim: int):
         """

{bayinx-0.2.6/src/bayinx/machinery/variational → bayinx-0.2.9/src/bayinx/mhx/vi}/flows/planar.py

@@ -20,9 +20,7 @@ class Planar(Flow):
     """
 
     params: Dict[str, Float[Array, "..."]]
-    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]] = (
-        eqx.field(static=True)
-    )
+    constraints: Dict[str, Callable[[Array], Array]]
 
     def __init__(self, dim: int, key=jr.PRNGKey(0)):
         """

{bayinx-0.2.6/src/bayinx/machinery/variational → bayinx-0.2.9/src/bayinx/mhx/vi}/flows/radial.py

@@ -21,9 +21,7 @@ class Radial(Flow):
     """
 
     params: Dict[str, Float[Array, "..."]]
-    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]] = (
-        eqx.field(static=True)
-    )
+    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]]
 
     def __init__(self, dim: int, key=jr.PRNGKey(0)):
         """
@@ -88,7 +86,8 @@ class Radial(Flow):
         # Compute density adjustment
         ladj = jnp.log(
             jnp.abs(
-                (1.0 + alpha * beta / (alpha + r) ** 2.0) * (1.0 + x) ** (center.size - 1.0)
+                (1.0 + alpha * beta / (alpha + r) ** 2.0)
+                * (1.0 + x) ** (center.size - 1.0)
             )
         )
 

bayinx-0.2.9/src/bayinx/mhx/vi/flows/sylvester.py

@@ -0,0 +1,19 @@
+from typing import Callable, Dict
+
+from jaxtyping import Array, Float
+
+from bayinx.core import Flow
+
+
+# TODO
+class Sylvester(Flow):
+    """
+    A sylvester flow.
+
+    # Attributes
+    - `params`: A dictionary containing the JAX Arrays representing the flow parameters.
+    - `constraints`: A dictionary of constraining transformations.
+    """
+
+    params: Dict[str, Float[Array, "..."]]
+    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]]

{bayinx-0.2.6/src/bayinx/machinery/variational → bayinx-0.2.9/src/bayinx/mhx/vi}/meanfield.py

@@ -31,15 +31,15 @@ class MeanField(Variational):
         - `model`: A probabilistic `Model` object.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, eqx.is_array)
+        params, self._constraints = eqx.partition(model, model.filter_spec())
 
         # Flatten params component
-        flat_params, self._unflatten = ravel_pytree(params)
+        params, self._unflatten = ravel_pytree(params)
 
         # Initialize variational parameters
         self.var_params = {
-            "mean": flat_params,
-            "log_std": jnp.zeros(flat_params.size, dtype=flat_params.dtype),
+            "mean": params,
+            "log_std": jnp.zeros(params.size, dtype=params.dtype),
         }
 
     @eqx.filter_jit

{bayinx-0.2.6/src/bayinx/machinery/variational → bayinx-0.2.9/src/bayinx/mhx/vi}/normalizing_flow.py

@@ -23,8 +23,8 @@ class NormalizingFlow(Variational):
 
     flows: list[Flow]
     base: Variational
-    _unflatten: Callable[[Float[Array, "..."]], Model] = eqx.field(static=True)
-    _constraints: Model = eqx.field(static=True)
+    _unflatten: Callable[[Float[Array, "..."]], Model]
+    _constraints: Model
 
     def __init__(self, base: Variational, flows: list[Flow], model: Model):
         """
@@ -39,7 +39,7 @@ class NormalizingFlow(Variational):
         params, self._constraints = eqx.partition(model, eqx.is_array)
 
         # Flatten params component
-        flat_params, self._unflatten = jfu.ravel_pytree(params)
+        _, self._unflatten = jfu.ravel_pytree(params)
 
         self.base = base
         self.flows = flows
@@ -73,7 +73,7 @@ class NormalizingFlow(Variational):
         return variational_evals
 
     @eqx.filter_jit
-    def _eval(self, draws: Array, data=None) -> Tuple[Scalar, Array]:
+    def __eval(self, draws: Array, data=None) -> Tuple[Array, Array]:
         """
         Evaluate the posterior and variational densities at the transformed
         `draws` to avoid extra compute when requiring variational draws for
@@ -84,7 +84,7 @@ class NormalizingFlow(Variational):
         - `data`: Any data required to evaluate the posterior density.
 
         # Returns
-        The posterior and variational densities.
+        The posterior and variational densities as JAX Arrays.
         """
         # Evaluate base density
         variational_evals: Array = self.base.eval(draws)
@@ -102,30 +102,30 @@ class NormalizingFlow(Variational):
         return posterior_evals, variational_evals
 
     def filter_spec(self):
-        # Only optimize the parameters of the flows
+        # Generate empty specification
         filter_spec = jtu.tree_map(lambda _: False, self)
+
+        # Specify variational parameters based on each flow's filter spec.
         filter_spec = eqx.tree_at(
-            lambda nf: nf.flows,
+            lambda vari: vari.flows,
             filter_spec,
-            replace=True,
+            replace=[flow.filter_spec() for flow in self.flows],
         )
 
         return filter_spec
 
     @eqx.filter_jit
     def elbo(self, n: int, key: Key, data: Any = None) -> Scalar:
-        # Partition
         dyn, static = eqx.partition(self, self.filter_spec())
 
         @eqx.filter_jit
         def elbo(dyn: Self, n: int, key: Key, data: Any = None):
-            # Combine
             self = eqx.combine(dyn, static)
 
             # Sample draws from variational distribution
             draws: Array = self.base.sample(n, key)
 
-            posterior_evals, variational_evals = self._eval(draws, data)
+            posterior_evals, variational_evals = self.__eval(draws, data)
             # Evaluate ELBO
             return jnp.mean(posterior_evals - variational_evals)
 
@@ -133,19 +133,17 @@ class NormalizingFlow(Variational):
 
     @eqx.filter_jit
     def elbo_grad(self, n: int, key: Key, data: Any = None) -> Self:
-        # Partition
         dyn, static = eqx.partition(self, self.filter_spec())
 
         @eqx.filter_grad
         @eqx.filter_jit
         def elbo_grad(dyn: Self, n: int, key: Key, data: Any = None):
-            # Combine
             self = eqx.combine(dyn, static)
 
             # Sample draws from variational distribution
             draws: Array = self.base.sample(n, key)
 
-            posterior_evals, variational_evals = self._eval(draws, data)
+            posterior_evals, variational_evals = self.__eval(draws, data)
             # Evaluate ELBO
             return jnp.mean(posterior_evals - variational_evals)
 

{bayinx-0.2.6/src/bayinx/machinery/variational → bayinx-0.2.9/src/bayinx/mhx/vi}/standard.py

@@ -13,15 +13,15 @@ from bayinx.dists import normal
 
 class Standard(Variational):
     """
-    A standard normal distribution approximation to a posterior distribution.
+    A standard normal approximation to a posterior distribution.
 
     # Attributes
     - `dim`: Dimension of the parameter space.
     """
 
     dim: int = eqx.field(static=True)
-    _unflatten: Callable[[Float[Array, "..."]], Model] = eqx.field(static=True)
-    _constraints: Model = eqx.field(static=True)
+    _unflatten: Callable[[Float[Array, "..."]], Model]
+    _constraints: Model
 
     def __init__(self, model: Model):
         """
@@ -31,13 +31,13 @@ class Standard(Variational):
         - `model`: A probabilistic `Model` object.
         """
         # Partition model
-        _, self._constraints = eqx.partition(model, eqx.is_array)
+        params, self._constraints = eqx.partition(model, model.filter_spec())
 
         # Flatten params component
-        _, self._unflatten = ravel_pytree(_)
+        params, self._unflatten = ravel_pytree(params)
 
         # Store dimension of parameter space
-        self.dim = jnp.size(_)
+        self.dim = jnp.size(params)
 
     @eqx.filter_jit
     def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:

{bayinx-0.2.6 → bayinx-0.2.9}/tests/test_variational.py

@@ -7,16 +7,14 @@ from jaxtyping import Array
 
 from bayinx import Model
 from bayinx.dists import normal
-from bayinx.machinery.variational import MeanField, NormalizingFlow, Standard
-from bayinx.machinery.variational.flows.affine import Affine
-from bayinx.machinery.variational.flows.planar import Planar
-from bayinx.machinery.variational.flows.radial import Radial
+from bayinx.mhx.vi import MeanField, NormalizingFlow, Standard
+from bayinx.mhx.vi.flows import Affine, Planar, Radial
 
 
 # Tests ----
 @pytest.mark.parametrize("var_draws", [1, 10, 100])
 def test_meanfield(benchmark, var_draws):
-    # Construct model
+    # Construct model definition
     class NormalDist(Model):
         params: Dict[str, Array]
         constraints: Dict[str, Callable[[Array], Array]]
@@ -35,6 +33,7 @@ def test_meanfield(benchmark, var_draws):
                 normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0))
             )
 
+    # Construct model
     model = NormalDist()
 
     # Construct meanfield variational
@@ -42,7 +41,7 @@ def test_meanfield(benchmark, var_draws):
 
     # Optimize variational distribution
     def benchmark_fit():
-        vari.fit(10000, var_draws = var_draws)
+        vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
     vari = vari.fit(10000)
@@ -55,7 +54,7 @@ def test_meanfield(benchmark, var_draws):
 
 @pytest.mark.parametrize("var_draws", [1, 10, 100])
 def test_affine(benchmark, var_draws):
-    # Construct model
+    # Construct model definition
     class NormalDist(Model):
         params: Dict[str, Array]
         constraints: Dict[str, Callable[[Array], Array]]
@@ -74,6 +73,7 @@ def test_affine(benchmark, var_draws):
                 normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0))
             )
 
+    # Construct model
     model = NormalDist()
 
     # Construct normalizing flow variational
@@ -81,7 +81,7 @@ def test_affine(benchmark, var_draws):
 
     # Optimize variational distribution
     def benchmark_fit():
-        vari.fit(10000, var_draws = var_draws)
+        vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
     vari = vari.fit(10000)
@@ -94,7 +94,7 @@ def test_affine(benchmark, var_draws):
 
 @pytest.mark.parametrize("var_draws", [1, 10, 100])
 def test_flows(benchmark, var_draws):
-    # Construct model
+    # Construct model definition
     class NormalDist(Model):
         params: Dict[str, Array]
         constraints: Dict[str, Callable[[Array], Array]]
@@ -113,14 +113,17 @@ def test_flows(benchmark, var_draws):
                 normal.logprob(x=params["mu"], mu=jnp.array(10.0), sigma=jnp.array(1.0))
             )
 
+    # Construct model
     model = NormalDist()
 
     # Construct normalizing flow variational
-    vari = NormalizingFlow(Standard(model), [Planar(2), Radial(2), Planar(2), Radial(2), Planar(2)], model)
+    vari = NormalizingFlow(
+        Standard(model), [Planar(2), Radial(2), Planar(2), Radial(2), Planar(2)], model
+    )
 
     # Optimize variational distribution
     def benchmark_fit():
-        vari.fit(10000, var_draws = var_draws)
+        vari.fit(10000, var_draws=var_draws)
 
     benchmark(benchmark_fit)
     vari = vari.fit(10000)

{bayinx-0.2.6 → bayinx-0.2.9}/uv.lock

@@ -17,13 +17,17 @@ wheels = [
 
 [[package]]
 name = "bayinx"
-version = "0.2.1"
+version = "0.2.2"
 source = { editable = "." }
 dependencies = [
     { name = "equinox" },
     { name = "jax" },
     { name = "jaxtyping" },
     { name = "optax" },
+]
+
+[package.dev-dependencies]
+dev = [
     { name = "pytest" },
     { name = "pytest-benchmark" },
 ]
@@ -34,6 +38,10 @@ requires-dist = [
     { name = "jax", specifier = ">=0.4.38" },
     { name = "jaxtyping", specifier = ">=0.2.36" },
     { name = "optax", specifier = ">=0.2.4" },
+]
+
+[package.metadata.requires-dev]
+dev = [
     { name = "pytest", specifier = ">=8.3.5" },
     { name = "pytest-benchmark", specifier = ">=5.1.0" },
 ]
@@ -96,11 +104,11 @@ epy = [
 
 [[package]]
 name = "iniconfig"
-version = "2.0.0"
+version = "2.1.0"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/d7/4b/cbd8e699e64a6f16ca3a8220661b5f83792b3017d0f79807cb8708d33913/iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3", size = 4646 }
+sdist = { url = "https://files.pythonhosted.org/packages/f2/97/ebf4da567aa6827c909642694d71c9fcf53e5b504f2d96afea02718862f3/iniconfig-2.1.0.tar.gz", hash = "sha256:3abbd2e30b36733fee78f9c7f7308f2d0050e88f0087fd25c2645f63c773e1c7", size = 4793 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/ef/a6/62565a6e1cf69e10f5727360368e451d4b7f58beeac6173dc9db836a5b46/iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374", size = 5892 },
+    { url = "https://files.pythonhosted.org/packages/2c/e1/e6716421ea10d38022b952c159d5161ca1193197fb744506875fbb87ea7b/iniconfig-2.1.0-py3-none-any.whl", hash = "sha256:9deba5723312380e77435581c6bf4935c94cbfab9b1ed33ef8d238ea168eb760", size = 6050 },
 ]
 
 [[package]]

bayinx-0.2.6/src/bayinx/core/flow.py

@@ -1,68 +0,0 @@
-from abc import abstractmethod
-from typing import Callable, Dict, Self, Tuple
-
-import equinox as eqx
-from jaxtyping import Array, Float
-
-from bayinx.core.utils import __MyMeta
-
-
-class Flow(eqx.Module, metaclass=__MyMeta):
-    """
-    A superclass used to define continuously parameterized diffeomorphisms for normalizing flows.
-
-    # Attributes
-    - `pars`: A dictionary of JAX Arrays representing parameters of the diffeomorphism.
-    - `constraints`: A dictionary of 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-determinant of the Jacobian at `draws` and applies the forward transformation.
-
-        # Returns
-        A tuple of JAX Arrays containing the log-absolute-determinant of the Jacobians and transformed draws.
-        """
-        pass
-
-    def __init_subclass__(cls):
-        # Add contrain_pars method
-        def constrain_pars(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
-
-        cls.constrain_pars = eqx.filter_jit(constrain_pars)
-
-        # Add transform_pars method if not present
-        if not callable(getattr(cls, "transform_pars", None)):
-            def transform_pars(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_pars()
-
-            cls.transform_pars = eqx.filter_jit(transform_pars)

bayinx-0.2.6/src/bayinx/machinery/variational/__init__.py

@@ -1,5 +0,0 @@
-from bayinx.machinery.variational.meanfield import MeanField as MeanField
-from bayinx.machinery.variational.normalizing_flow import (
-    NormalizingFlow as NormalizingFlow,
-)
-from bayinx.machinery.variational.standard import Standard as Standard

bayinx-0.2.6/src/bayinx/machinery/variational/flows/__init__.py

@@ -1,3 +0,0 @@
-from bayinx.machinery.variational.flows.affine import Affine as Affine
-from bayinx.machinery.variational.flows.planar import Planar as Planar
-from bayinx.machinery.variational.flows.radial import Radial as Radial

bayinx-0.2.6/src/bayinx/machinery/variational/flows/sylvester.py

@@ -1,76 +0,0 @@
-from functools import partial
-from typing import Callable, Dict, Tuple
-
-import equinox as eqx
-import jax
-import jax.numpy as jnp
-import jax.random as jr
-from jaxtyping import Array, Float, Scalar
-
-from bayinx.core import Flow
-
-
-class Sylvester(Flow):
-    """
-    A sylvester flow.
-
-    # Attributes
-    - `params`: A dictionary containing the JAX Arrays representing the flow parameters.
-    - `constraints`: A dictionary of constraining transformations.
-    """
-
-    params: Dict[str, Float[Array, "..."]]
-    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]] = (
-        eqx.field(static=True)
-    )
-
-    def __init__(self, dim: int, key=jr.PRNGKey(0)):
-        """
-        Initializes a planar flow.
-
-        # Parameters
-        - `dim`: The dimension of the parameter space.
-        """
-        self.params = {
-            "u": jr.normal(key, (dim,)),
-            "w": jr.normal(key, (dim,)),
-            "b": jr.normal(key, (1,)),
-        }
-        self.constraints = {}
-
-    @eqx.filter_jit
-    @partial(jax.vmap, in_axes=(None, 0))
-    def forward(self, draws: Array) -> Array:
-        params = self.constrain_pars()
-
-        # Extract parameters
-        w: Array = params["w"]
-        u: Array = params["u"]
-        b: Array = params["b"]
-
-        # Compute forward transformation
-        draws = draws + u * jnp.tanh(draws.dot(w) + b)
-
-        return draws
-
-    @eqx.filter_jit
-    @partial(jax.vmap, in_axes=(None, 0))
-    def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
-        params = self.constrain_pars()
-
-        # Extract parameters
-        w: Array = params["w"]
-        u: Array = params["u"]
-        b: Array = params["b"]
-
-        # Compute shared intermediates
-        x: Array = draws.dot(w) + b
-
-        # Compute forward transformation
-        draws = draws + u * jnp.tanh(x)
-
-        # Compute ladj
-        h_prime: Scalar = 1.0 - jnp.square(jnp.tanh(x))
-        ladj: Scalar = jnp.log(jnp.abs(1.0 + h_prime * u.dot(w)))
-
-        return ladj, draws