bayinx 0.1.0__tar.gz → 0.2.3__tar.gz

bayinx-0.2.3/.github/workflows/publish.yml

@@ -0,0 +1,29 @@
+name: "Publish"
+
+on:
+  release:
+    types: ["published"]
+
+jobs:
+  run:
+    name: "Build and publish release"
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write # Required for trusted publishing
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          enable-cache: true
+
+      - name: Set up Python
+        run: uv python install 3.12
+
+      - name: Build
+        run: uv build
+
+      - name: Publish
+        run: uv publish

{bayinx-0.1.0 → bayinx-0.2.3}/.gitignore

@@ -5,6 +5,11 @@ build/
 dist/
 wheels/
 *.egg-info
+.pytest_cache
+.benchmarks
 
 # Virtual environments
 .venv
+
+# Other
+.ruff_cache

bayinx-0.2.3/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: bayinx
+Version: 0.2.3
+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>
+
+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 equivalent 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{M})$, or set up contrasts and evaluate $P(\mu_1 - \mu_2 > 0 | \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.2.3/README.md

@@ -0,0 +1,27 @@
+# <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 equivalent 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{M})$, or set up contrasts and evaluate $P(\mu_1 - \mu_2 > 0 | \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.2.3/pyproject.toml

@@ -0,0 +1,21 @@
+[project]
+name = "bayinx"
+version = "0.2.3"
+description = "Bayesian Inference with JAX"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "equinox>=0.11.12",
+    "jax>=0.4.38",
+    "jaxtyping>=0.2.36",
+    "optax>=0.2.4",
+    "pytest>=8.3.5",
+    "pytest-benchmark>=5.1.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.pytest.ini_options]
+addopts = "-q --benchmark-min-rounds=30 --benchmark-columns=rounds,mean,median,stddev --benchmark-group-by=func"

bayinx-0.2.3/src/bayinx/__init__.py

@@ -0,0 +1 @@
+from bayinx.core.model import Model as Model

bayinx-0.2.3/src/bayinx/core/__init__.py

@@ -0,0 +1,3 @@
+from bayinx.core.flow import Flow as Flow
+from bayinx.core.model import Model as Model
+from bayinx.core.variational import Variational as Variational

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

@@ -0,0 +1,68 @@
+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.3/src/bayinx/core/model.py

@@ -0,0 +1,55 @@
+from abc import abstractmethod
+from typing import Any, Callable, Dict
+
+import equinox as eqx
+from jaxtyping import Array, Scalar
+
+from bayinx.core.utils import __MyMeta
+
+
+class Model(eqx.Module, metaclass=__MyMeta):
+    """
+    A superclass used to define probabilistic models.
+
+    # Attributes
+    - `params`: A dictionary of JAX Arrays representing parameters of the model.
+    - `constraints`: A dictionary of functions that constrain their corresponding parameter.
+    """
+
+    params: Dict[str, Array]
+    constraints: Dict[str, Callable[[Array], Array]]
+
+    @abstractmethod
+    def eval(self, data: Any) -> Scalar:
+        pass
+
+    def __init_subclass__(cls):
+        # Add constrain method
+        def constrain_pars(self: Model) -> Dict[str, Array]:
+            """
+            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: Model) -> 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.3/src/bayinx/core/utils.py

@@ -0,0 +1,54 @@
+from typing import Callable, Dict
+
+import equinox as eqx
+from jaxtyping import Array
+
+
+class __MyMeta(type(eqx.Module)):
+    """
+    Metaclass to ensure attribute types are respected.
+    """
+
+    def __call__(cls, *args, **kwargs):
+        obj = super().__call__(*args, **kwargs)
+
+        # Check parameters are a Dict of JAX Arrays
+        if not isinstance(obj.params, Dict):
+            raise ValueError(
+                f"Model {cls.__name__} must initialize 'params' as a dictionary."
+            )
+
+        for key, value in obj.params.items():
+            if not isinstance(value, Array):
+                raise TypeError(f"Parameter '{key}' must be a JAX Array.")
+
+        # Check constraints are a Dict of functions
+        if not isinstance(obj.constraints, Dict):
+            raise ValueError(
+                f"Model {cls.__name__} must initialize 'constraints' as a dictionary."
+            )
+
+        for key, value in obj.constraints.items():
+            if not isinstance(value, Callable):
+                raise TypeError(f"Constraint for parameter '{key}' must be a function.")
+
+        # Check that the constrain method returns a dict equivalent to `params`
+        t_params: Dict[str, Array] = obj.constrain_pars()
+
+        if not isinstance(t_params, Dict):
+            raise ValueError(
+                f"The 'constrain' method of {cls.__name__} must return a Dict of JAX Arrays."
+            )
+
+        for key, value in t_params.items():
+            if not isinstance(value, Array):
+                raise TypeError(f"Constrained parameter '{key}' must be a JAX Array.")
+
+            if not value.shape == obj.params[key].shape:
+                raise ValueError(
+                    f"Constrained parameter '{key}' must have same shape as unconstrained counterpart."
+                )
+
+        # Check transform_pars
+
+        return obj

bayinx-0.2.3/src/bayinx/core/variational.py

@@ -0,0 +1,159 @@
+from abc import abstractmethod
+from typing import Any, Callable, Self, Tuple
+
+import equinox as eqx
+import jax
+import jax.lax as lax
+import jax.numpy as jnp
+import jax.random as jr
+import optax as opx
+from jaxtyping import Array, Float, Key, PyTree, Scalar
+from optax import GradientTransformation, OptState, Schedule
+
+from bayinx.core import Model
+
+
+class Variational(eqx.Module):
+    """
+    A superclass used to define variational methods.
+
+    # Attributes
+    - `_unflatten`: A static function to transform draws from the variational distribution back to a `Model`.
+    - `_constraints`: A static partitioned `Model` with the constraints of the `Model` used to initialize the `Variational` object.
+    """
+
+    _unflatten: Callable[[Float[Array, "..."]], Model]
+    _constraints: Model
+
+    @abstractmethod
+    def sample(self, n: int, key: Key) -> Array:
+        """
+        Sample from the variational distribution.
+        """
+        pass
+
+    @abstractmethod
+    def eval(self, draws: Array) -> Array:
+        """
+        Evaluate the variational distribution at `draws`.
+        """
+        pass
+
+    @abstractmethod
+    def elbo(self, n: int, key: Key, data: Any = None) -> Array:
+        """
+        Evaluate the ELBO.
+        """
+        pass
+
+    @abstractmethod
+    def elbo_grad(self, n: int, key: Key, data: Any = None) -> PyTree:
+        """
+        Evaluate the gradient of the ELBO.
+        """
+        pass
+
+    @abstractmethod
+    def filter_spec(self):
+        """
+        Filter specification for dynamic and static components of the `Variational`.
+        """
+        pass
+
+    def __init_subclass__(cls):
+        """
+        Construct methods that are shared across all VI methods.
+        """
+
+        def eval_model(self, draws: Array, data: Any = None) -> Array:
+            """
+            Reconstruct models from variational draws and evaluate their posterior density.
+
+            # Parameters
+            - `draws`: A set of variational draws.
+            - `data`: Data used to evaluate the posterior(if needed).
+            """
+            # Unflatten variational draw
+            model: Model = self._unflatten(draws)
+
+            # Combine with constraints
+            model: Model = eqx.combine(model, self._constraints)
+
+            # Evaluate posterior density
+            return model.eval(data)
+
+        cls.eval_model = jax.vmap(eqx.filter_jit(eval_model), (None, 0, None))
+
+        def fit(
+            self,
+            max_iters: int,
+            data: Any = None,
+            learning_rate: float = 1,
+            tolerance: float = 1e-4,
+            var_draws: int = 1,
+            key: Key = jr.PRNGKey(0),
+        ) -> Self:
+            """
+            Optimize the variational distribution.
+
+            # Parameters
+            - `max_iters`: Maximum number of iterations for the optimization loop.
+            - `data`: Data to evaluate the posterior density with(if available).
+            - `learning_rate`: Initial learning rate for optimization.
+            - `tolerance`: Relative tolerance of ELBO decrease for stopping early.
+            - `var_draws`: Number of variational draws to draw each iteration.
+            - `key`: A PRNG key.
+            """
+            # Construct scheduler
+            schedule: Schedule = opx.cosine_decay_schedule(
+                init_value=learning_rate, decay_steps=max_iters
+            )
+
+            # Initialize optimizer
+            optim: GradientTransformation = opx.chain(
+                opx.scale(-1.0), opx.nadam(schedule)
+            )
+            opt_state: OptState = optim.init(eqx.filter(self, self.filter_spec()))
+
+            # Optimization loop helper functions
+            @eqx.filter_jit
+            def condition(state: Tuple[Self, OptState, Scalar, Key]):
+                # Unpack iteration state
+                self, 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
+
+                # Update iteration
+                i = i + 1
+
+                # Update PRNG key
+                key, _ = jr.split(key)
+
+                # Compute gradient of the ELBO
+                updates: PyTree = self.elbo_grad(var_draws, key, data)
+
+                # Compute updates
+                updates, opt_state = optim.update(
+                    updates, opt_state, eqx.filter(self, self.filter_spec())
+                )
+
+                # Update variational distribution
+                self: Self = eqx.apply_updates(self, updates)
+
+                return self, opt_state, i, key
+
+            # Run optimization loop
+            self = lax.while_loop(
+                cond_fun=condition,
+                body_fun=body,
+                init_val=(self, opt_state, jnp.array(0, jnp.uint32), key),
+            )[0]
+
+            return self
+
+        cls.fit = eqx.filter_jit(fit)

bayinx-0.2.3/src/bayinx/dists/bernoulli.py

@@ -0,0 +1,33 @@
+import jax.lax as lax
+from jaxtyping import Array, ArrayLike, Real, UInt
+
+
+# MARK: Functions ----
+def prob(x: UInt[ArrayLike, "..."], p: Real[ArrayLike, "..."]) -> Real[Array, "..."]:
+    """
+    The probability mass function (PMF) for a Bernoulli distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `p`:     The probability parameter(s).
+
+    # Returns
+    The PMF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
+    """
+
+    return lax.pow(p, x) * lax.pow(1 - p, 1 - x)
+
+
+def logprob(x: UInt[ArrayLike, "..."], p: Real[ArrayLike, "..."]) -> Real[Array, "..."]:
+    """
+    The log probability mass function (log PMF) for a Bernoulli distribution.
+
+    # Parameters
+    - `x`:  Value(s) at which to evaluate the log PMF.
+    - `p`:  The probability parameter(s).
+
+    # Returns
+    The log PMF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
+    """
+
+    return lax.log(p) * x + (1 - x) * lax.log(1 - p)

bayinx-0.2.3/src/bayinx/dists/normal.py

@@ -0,0 +1,83 @@
+# MARK: Imports ----
+import jax.lax as _lax
+
+## Typing
+from jaxtyping import Array, Real
+
+# MARK: Constants
+_PI = 3.141592653589793
+
+
+# MARK: Functions ----
+def prob(
+    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
+) -> Real[Array, "..."]:
+    """
+    The probability density function (PDF) for a Normal distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `mu`:     The mean/location parameter(s).
+    - `sigma`:  The non-negative standard deviation parameter(s).
+
+    # Returns
+    The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+
+    return _lax.exp(-0.5 * _lax.square((x - mu) / sigma)) / (
+        sigma * _lax.sqrt(2.0 * _PI)
+    )
+
+
+def logprob(
+    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
+) -> Real[Array, "..."]:
+    """
+    The log of the probability density function (log PDF) for a Normal distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the log PDF.
+    - `mu`:     The mean/location parameter(s).
+    - `sigma`:  The non-negative standard deviation parameter(s).
+
+    # Returns
+    The log of the PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+
+    return -_lax.log(sigma * _lax.sqrt(2.0 * _PI)) - 0.5 * _lax.square((x - mu) / sigma)
+
+
+def uprob(
+    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
+) -> Real[Array, "..."]:
+    """
+    The unnormalized probability density function (uPDF) for a Normal distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the uPDF.
+    - `mu`:     The mean/location parameter(s).
+    - `sigma`:  The non-negative standard deviation parameter(s).
+
+    # Returns
+    The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+
+    return _lax.exp(-0.5 * _lax.square((x - mu) / sigma)) / sigma
+
+
+def ulogprob(
+    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
+) -> Real[Array, "..."]:
+    """
+    The log of the unnormalized probability density function (log uPDF) for a Normal distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the log uPDF.
+    - `mu`:     The mean/location parameter(s).
+    - `sigma`:  The non-negative standard deviation parameter(s).
+
+    # Returns
+    The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+
+    return -_lax.log(sigma) - 0.5 * _lax.square((x - mu) / sigma)

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

@@ -0,0 +1,5 @@
+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.3/src/bayinx/machinery/variational/flows/__init__.py

@@ -0,0 +1,3 @@
+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.3/src/bayinx/machinery/variational/flows/affine.py

@@ -0,0 +1,68 @@
+from functools import partial
+from typing import Callable, Dict, Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, Float, Scalar
+
+from bayinx.core import Flow
+
+
+class Affine(Flow):
+    """
+    An affine flow.
+
+    # Attributes
+    - `params`: A dictionary containing the JAX Arrays representing the scale and shift 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):
+        """
+        Initializes an affine flow.
+
+        # Parameters
+        - `dim`: The dimension of the parameter space.
+        """
+        self.params = {
+            "shift": jnp.zeros(dim),
+            "scale": jnp.zeros((dim, dim)),
+        }
+
+        self.constraints = {"scale": lambda m: jnp.tril(jnp.exp(m))}
+
+    @eqx.filter_jit
+    def forward(self, draws: Array) -> Array:
+        params = self.constrain_pars()
+
+        # Extract parameters
+        shift: Array = params["shift"]
+        scale: Array = params["scale"]
+
+        # Compute forward transformation
+        draws = draws @ scale + shift
+
+        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
+        shift: Array = params["shift"]
+        scale: Array = params["scale"]
+
+        # Compute forward transformation
+        draws = draws @ scale + shift
+
+        # Compute ladj
+        ladj: Scalar = jnp.log(jnp.diag(scale)).sum()
+
+        return ladj, draws