bayinx 0.3.19__py3-none-any.whl → 0.4.1__py3-none-any.whl

bayinx/__init__.py

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

bayinx/constraints.py

@@ -0,0 +1,135 @@
+from typing import Tuple
+
+import equinox as eqx
+import jax.nn as jnn
+import jax.numpy as jnp
+import jax.tree as jt
+from jaxtyping import Array, PyTree, Scalar, ScalarLike
+
+from bayinx.core import Constraint, Parameter
+from bayinx.core._parameter import T
+
+
+class Lower(Constraint):
+    """
+    Enforces a lower bound on the parameter.
+    """
+
+    lb: Scalar
+
+    def __init__(self, lb: ScalarLike):
+        # assert greater than 1
+        self.lb = jnp.asarray(lb)
+
+    def constrain(self, param: Parameter[T]) -> Tuple[Parameter[T], Scalar]:
+        """
+        Enforces a lower bound on the parameter and adjusts the posterior density.
+
+        # Parameters
+        - `param`: The unconstrained `Parameter`.
+
+        # Returns
+        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 parameters(all inexact Arrays)
+        dyn, static = eqx.partition(param, param.filter_spec)
+
+        # Compute Jacobian adjustment
+        total_laj: Scalar = jt.reduce(lambda a, b: a + b, jt.map(jnp.sum, dyn))
+
+        # Compute transformation
+        dyn = jt.map(lambda v: jnp.exp(v) + self.lb, dyn)
+
+        # Combine into full parameter object
+        param = eqx.combine(dyn, static)
+
+        return param, total_laj
+
+
+class LogSimplex(Constraint):
+    """
+    Enforces a log-transformed simplex constraint on the parameter.
+
+    # Attributes
+    - `sum`: The total sum of the parameter.
+    """
+
+    sum: Scalar
+
+    def __init__(self, sum_val: ScalarLike = 1.0):
+        """
+        # Parameters
+        - `sum_val`: The target sum for the exponentiated simplex. Defaults to 1.0.
+        """
+        self.sum = jnp.asarray(sum_val)
+
+    def constrain(self, param: Parameter[T]) -> Tuple[Parameter[T], Scalar]:
+        """
+        Enforces a log-transformed simplex constraint on the parameter and adjusts the posterior density.
+
+        # Parameters
+        - `param`: The unconstrained `Parameter`.
+
+        # Returns
+        A tuple containing:
+            - A modified `Parameter` with relevant leaves satisfying the constraint.
+            - A scalar Array representing the log-absolute-Jacobian of the transformation.
+        """
+        # Partition the parameter into dynamic (to be transformed) and static parts
+        dyn, static = eqx.partition(param, param.filter_spec)
+
+        # Map transformation leaf-wise
+        transformed = jt.map(self._transform_leaf, dyn) ## filter spec handles subsetting arrays, is_leaf unnecessary
+
+        # Extract constrained parameters and Jacobian adjustments
+        dyn_constrained: PyTree = jt.map(lambda x: x[0], transformed)
+        lajs: PyTree = jt.map(lambda x: x[1], transformed)
+
+        # Sum to get total Jacobian adjustment
+        total_laj = jt.reduce(lambda a, b: a + b, lajs)
+
+        # Recombine the transformed dynamic parts with the static parts
+        param = eqx.combine(dyn_constrained, static)
+
+        return param, total_laj
+
+    def _transform_leaf(self, x: Array) -> Tuple[Array, Scalar]:
+        """
+        Internal function that applies a log-transformed simplex constraint on a single array.
+        """
+        laj: Scalar = jnp.array(0.0)
+
+        # Save output shape
+        output_shape: tuple[int, ...] = x.shape
+
+        if x.size == 1:
+            return(jnp.full(output_shape, jnp.log(self.sum)), laj)
+        else:
+            # Flatten x
+            x = x.flatten()
+
+            # Subset first K - 1 elements
+            x = x[:-1]
+
+            # Compute shifted cumulative sum
+            zeta: Array = jnp.concat([jnp.zeros(1), x.cumsum()[:-1]])
+
+            # Compute intermediate proportions vector
+            eta: Array = jnn.sigmoid(x - zeta)
+
+            # Compute Jacobian adjustment
+            laj += jnp.sum(jnp.log(eta) + jnp.log(1 - eta)) # TODO: check for correctness
+
+            # Compute log-transformed simplex weights
+            w: Array = jnp.log(eta) + jnp.concatenate([jnp.array([0.0]), jnp.log(jnp.cumprod((1-eta)[:-1]))])
+            w = jnp.concatenate([w, jnp.log(jnp.prod(1 - eta, keepdims=True))])
+
+            # Scale unit simplex on log-scale
+            w = w + jnp.log(self.sum)
+
+            # Reshape for output
+            w = w.reshape(output_shape)
+
+            return (w, laj)

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/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/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/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/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/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/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/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/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/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/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/mhx/vi/flows/radial.py

@@ -37,8 +37,8 @@ class Radial(Flow):
         }
         self.constraints = {"beta": jnp.exp}
 
-    @partial(jax.vmap, in_axes=(None, 0))
     @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0))
     def forward(self, draws: Array) -> Array:
         """
         Applies the forward radial transformation for each draw.

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

bayinx/mhx/vi/normalizing_flow.py

@@ -31,11 +31,11 @@ class NormalizingFlow(Variational, Generic[M]):
 
         # Parameters
         - `base`: The base variational distribution.
-        - `flows`: A list of diffeomorphisms.
+        - `flows`: A list of flows.
         - `model`: A probabilistic `Model` object.
         """
         # Partition model
-        params, self._constraints = eqx.partition(model, model.filter_spec)
+        params, self._static = eqx.partition(model, model.filter_spec)
 
         # Flatten params component
         _, self._unflatten = jfu.ravel_pytree(params)
@@ -44,7 +44,6 @@ class NormalizingFlow(Variational, Generic[M]):
         self.flows = flows
 
     @property
-    @eqx.filter_jit
     def filter_spec(self):
         # Generate empty specification
         filter_spec = jtu.tree_map(lambda _: False, self)
@@ -78,7 +77,7 @@ class NormalizingFlow(Variational, Generic[M]):
         variational_evals: Array = self.base.eval(draws)
 
         for map in self.flows:
-            # Compute adjustment
+            # Apply transformation
             draws, laj = map.adjust_density(draws)
 
             # Adjust variational density
@@ -103,7 +102,7 @@ class NormalizingFlow(Variational, Generic[M]):
         variational_evals: Array = self.base.eval(draws)
 
         for map in self.flows:
-            # Compute adjustment
+            # Apply transformation
             draws, laj = map.adjust_density(draws)
 
             # Adjust variational density

bayinx/mhx/vi/standard.py

@@ -27,7 +27,7 @@ class Standard(Variational[M]):
         - `model`: A probabilistic `Model` object.
         """
         # 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)
@@ -35,6 +35,7 @@ class Standard(Variational[M]):
         # Store dimension of parameter space
         self.dim = jnp.size(params)
 
+
     @eqx.filter_jit
     def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:
         # Sample variational draws

bayinx-0.4.1.dist-info/METADATA

@@ -0,0 +1,47 @@
+Metadata-Version: 2.4
+Name: bayinx
+Version: 0.4.1
+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 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.1.dist-info/RECORD

@@ -0,0 +1,38 @@
+bayinx/__init__.py,sha256=8etrxEtEGEzSDmKsW0TB4XoUGLiMPt9wpwNR8CGe1gU,93
+bayinx/constraints.py,sha256=2ufHsXR-_bWKR4WKKuR-OTjj3XCc4TkSeHVGWYadwCg,4387
+bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/core/__init__.py,sha256=samkrHp2zYyj8n37k-06tlaVrSqbtcgoa1LO0btAEHc,338
+bayinx/core/_constraint.py,sha256=Gx07ZT66VE2y-qZCmBDm3_y0wO4xQyslZW10Lec1_lM,761
+bayinx/core/_flow.py,sha256=V7uVfo2kBkyD4lWKBdSFFSmla68E_Hqv-uu9akuWBio,2282
+bayinx/core/_model.py,sha256=9uRVG5Zwrt9GHCY1-ULXw4KCZWHBpsG8AirJplDMT78,3848
+bayinx/core/_optimization.py,sha256=Ehp8UqtN-IYrEeX178TFUjCsJEBHxDnqxvgIco15zYU,2654
+bayinx/core/_parameter.py,sha256=O2s0-WxxqWOqwjuqWBh-JU-rRSBKpT_Fzy4ZxXr3cEc,1080
+bayinx/core/_variational.py,sha256=ddjtdq6z8tCJ5RZ4crveiN1HuPILdejYfwcAiZDLNZ4,6164
+bayinx/dists/__init__.py,sha256=BIrypqMnTLWK3a_zw8fYKMyuEMxP_qGsLfLeScias0o,118
+bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
+bayinx/dists/gamma2.py,sha256=HtB60LUQdPj4yDAHme2jsHNmLfrAKWsSZnDYkxAGaOI,1548
+bayinx/dists/negbinom3.py,sha256=u_USHQHxXmdS6hDCW2xmDpcqlxf17SJLaeptpxAI8TQ,2912
+bayinx/dists/normal.py,sha256=Yc2X8F7JoLYwprtK8bA2BPva1tAY7MEs3oSk5pMortI,3822
+bayinx/dists/posnormal.py,sha256=P4ShZPeqw2dr1mnhgaMku8JLEmCeQOe3jLb-BQbqs9o,8087
+bayinx/dists/uniform.py,sha256=FGXEIrvq4UXIhZ3mz4EDmbkpnZxUa2MkyeqEUu0CGJ4,3470
+bayinx/dists/censored/__init__.py,sha256=UVihMbQgAzCoOk_Zt5wrumPv5-acuTzV3TYMB-U1gOc,49
+bayinx/dists/censored/gamma2/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
+bayinx/dists/censored/gamma2/r.py,sha256=dKAOYstufwgDwibQZHrJxA1d2gawj-7K3IkaCRCzNTg,2446
+bayinx/dists/censored/negbinom3/r.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/dists/censored/posnormal/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
+bayinx/dists/censored/posnormal/r.py,sha256=nefEBJvjgekPL8L7L7-UXV5bc9cIWDlcDBwWWrA7YcM,3536
+bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+bayinx/mhx/opt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/mhx/vi/__init__.py,sha256=3T1dEpiiRge4tW-vpS0xBob_RbO1iVFnL3fVCRUawCM,205
+bayinx/mhx/vi/meanfield.py,sha256=GLX6fzXr1v_PEHNyXCcBgTT5_7tGSSEkrTkwCNujplI,3858
+bayinx/mhx/vi/normalizing_flow.py,sha256=xSRxJ6hmrT6Fnx-uSIJ7mqAO5JM7bmbJtj2uljcKERs,4681
+bayinx/mhx/vi/standard.py,sha256=s-Kvw37Y_KzlkcbeS6eVppg6DgPCclZQEcooGgbg3SU,1574
+bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
+bayinx/mhx/vi/flows/fullaffine.py,sha256=7KXuukwzVtMRIa8bSK_4pjnnP-lLIzVJBCAuKVydVgE,1925
+bayinx/mhx/vi/flows/planar.py,sha256=2I2WzIskl8MRpJkK13FQE3vSF-077qo8gRed_EL1Pn8,1920
+bayinx/mhx/vi/flows/radial.py,sha256=AyaqLJCwn871L6E8lBCU4Y8zZBF9UYZu6KIhzV6Z6wo,2503
+bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
+bayinx-0.4.1.dist-info/METADATA,sha256=7Zw-9hVqUVxj3ncyGBfn72FQzTDomdDaXXH2hOsJM60,2989
+bayinx-0.4.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.4.1.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.4.1.dist-info/RECORD,,

bayinx/constraints/__init__.py

@@ -1,3 +0,0 @@
-from bayinx.constraints.lower import Lower
-
-__all__ = ["Lower"]

bayinx/constraints/lower.py

@@ -1,50 +0,0 @@
-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 bayinx.core import Constraint, Parameter
-
-
-class Lower(Constraint):
-    """
-    Enforces a lower bound on the parameter.
-    """
-
-    lb: Scalar
-
-    def __init__(self, lb: ScalarLike):
-        self.lb = jnp.array(lb)
-
-    @eqx.filter_jit
-    def constrain(self, x: Parameter) -> Tuple[Parameter, Scalar]:
-        """
-        Enforces a lower bound on the parameter and adjusts the posterior density.
-
-        # Parameters
-        - `x`: 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)
-
-        # Compute density adjustment
-        laj: PyTree = jt.map(jnp.sum, dyn_params)  # pyright: ignore
-        laj: Scalar = jt.reduce(lambda a, b: a + b, laj)
-
-        # Compute transformation
-        dyn_params = jt.map(lambda v: jnp.exp(v) + self.lb, dyn_params)
-
-        # Combine into full parameter object
-        x = eqx.combine(dyn_params, static_params)
-
-        return x, laj

bayinx-0.3.19.dist-info/METADATA

@@ -1,39 +0,0 @@
-Metadata-Version: 2.4
-Name: bayinx
-Version: 0.3.19
-Summary: Bayesian Inference with JAX
-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 endgoal of this project is to build a Bayesian inference library that is similar in feel to `Stan`(where you can define a probabilistic model with syntax that is similar to how you would write it out on a chalkboard) but allows for arbitrary models(e.g., ones with discrete parameters) and offers a suite of "machinery" to fit the model; this means I want to expand upon `Stan`'s existing toolbox of methods for estimation(point optimization, variational methods, MCMC) while keeping everything performant(hence using `JAX`).
-
-In the short-term, I'm going to focus on:
-1) Implementing as much machinery as I feel is enough.
-2) Figuring out how to design the `Model` superclass to have something like the `transformed pars {}` block but unifies transformations and constraints.
-3) Figuring out how to design the library to automatically recognize what kind of machinery is amenable to a given probabilistic model.
-
-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{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`).
-
-Additionally, when I get around to it I'd like the package documentation to also include theoretical and implementation details for all machinery implemented(with overthinking boxes because I do like that design from McElreath's book).
-
-
-# TODO
-- Find some way to discern between models with all floating-point parameters and weirder models with integer parameters. Useful for restricting variational methods like `MeanField` to `Model`s that only have floating-point parameters.
-- Look into adaptively tuning ADAM hyperparameters.
-- 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 lol.
-- Figure out how to make transform_pars for flows such that there is no performance loss. Noticing some weird behaviour when adding constraints.

bayinx-0.3.19.dist-info/RECORD

@@ -1,37 +0,0 @@
-bayinx/__init__.py,sha256=TM-aoRaPX6jSYtCM7Jv59TPV-H6bcDk1-VMttYP1KME,99
-bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/constraints/__init__.py,sha256=027WJxRLkybXZkmusfvR6iZayY2pDid7Tw6TTTeA6ko,64
-bayinx/constraints/lower.py,sha256=30y0l6PF-tbS9LR_tto9AvwmsvXq1ExU-v8DLrJD4g4,1446
-bayinx/core/__init__.py,sha256=Qmy0EjzqqKwI9F8rjmC9j6J8hiDw6A54yOck2WuQJkY,344
-bayinx/core/_constraint.py,sha256=Gx07ZT66VE2y-qZCmBDm3_y0wO4xQyslZW10Lec1_lM,761
-bayinx/core/_flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
-bayinx/core/_model.py,sha256=FJUyYVE9e2uTFamxtSMKY_VV2stiU2QF68Wl_7EAKEU,2895
-bayinx/core/_optimization.py,sha256=mmeVUqfFARz8F7q4LRl-uEwVWzekmzh-9o7PnuvsHZk,2651
-bayinx/core/_parameter.py,sha256=r20JedTW2lY0miNNh9y6LeIVAsGX1kP_rlGxphW_jZg,1080
-bayinx/core/_variational.py,sha256=ljhLi9vbw8wh-z0Eisf0C08IoKlgwpx0VONh_ES-HmI,6384
-bayinx/dists/__init__.py,sha256=BIrypqMnTLWK3a_zw8fYKMyuEMxP_qGsLfLeScias0o,118
-bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
-bayinx/dists/gamma2.py,sha256=HtB60LUQdPj4yDAHme2jsHNmLfrAKWsSZnDYkxAGaOI,1548
-bayinx/dists/normal.py,sha256=Yc2X8F7JoLYwprtK8bA2BPva1tAY7MEs3oSk5pMortI,3822
-bayinx/dists/posnormal.py,sha256=cOLCdd39DX3v8DD-seSIKNk4OfdNfaYaLzpCh_xBGyw,7150
-bayinx/dists/uniform.py,sha256=2ZQxEfAX5TFgSPuQ8joFDuFbd_NfmQ1GvmGGjusqvNQ,3461
-bayinx/dists/censored/__init__.py,sha256=UVihMbQgAzCoOk_Zt5wrumPv5-acuTzV3TYMB-U1gOc,49
-bayinx/dists/censored/gamma2/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
-bayinx/dists/censored/gamma2/r.py,sha256=dKAOYstufwgDwibQZHrJxA1d2gawj-7K3IkaCRCzNTg,2446
-bayinx/dists/censored/posnormal/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
-bayinx/dists/censored/posnormal/r.py,sha256=wMDt2Am1TD376ms8B-o6PFCJZXmUJd2-aBC-t9kidH4,3456
-bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-bayinx/mhx/opt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/mhx/vi/__init__.py,sha256=3T1dEpiiRge4tW-vpS0xBob_RbO1iVFnL3fVCRUawCM,205
-bayinx/mhx/vi/meanfield.py,sha256=iX4AeDG9jrLZd6d9NimuJ3O5zaoBXsD03JbgPgxVrfY,3917
-bayinx/mhx/vi/normalizing_flow.py,sha256=vzLu5H1G1-pBqhgHWmIZkUTyPE1DxC9vBwpiZeIyu1I,4712
-bayinx/mhx/vi/standard.py,sha256=LYgglaGQMGmXpzFR4eMJnXkl2PhBeggbXMvO5zJpf2c,1578
-bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
-bayinx/mhx/vi/flows/fullaffine.py,sha256=11y_A0oO3bkKDSz-EQ6Sf4Ec2M7vHZxw94EdvADwVYQ,1954
-bayinx/mhx/vi/flows/planar.py,sha256=2I2WzIskl8MRpJkK13FQE3vSF-077qo8gRed_EL1Pn8,1920
-bayinx/mhx/vi/flows/radial.py,sha256=e0GfuO-CL8SVr3YnEfsxStpyKcJlFpzMyjMp3sa38hg,2503
-bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.3.19.dist-info/METADATA,sha256=V1qTM781r9t4cDGSRkgvn4PQAoxtsD9cFrMj0YVhUHo,3087
-bayinx-0.3.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.19.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.19.dist-info/RECORD,,