bayinx 0.3.19__tar.gz → 0.4.1__tar.gz

bayinx-0.4.1/PKG-INFO

@@ -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/README.md

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

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

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

bayinx-0.4.1/src/bayinx/__init__.py

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

bayinx-0.4.1/src/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-0.3.19 → bayinx-0.4.1}/src/bayinx/core/__init__.py

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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