mixed base distribution

[SophieMLV]

Hi, I'm looking for an option to have a mixed base distribution for the normalizing flow.
So far I only see the option to have the one base distribution for all dimensions, ie

flow = masked_autoregressive_flow(
    subkey,
    base_dist=Normal(jnp.zeros(x.shape[1])),
)

With a mixed base distribution, I mean, when one has, for example, 2 dimensions and for the first one wants a normal distribution, while for the second one an exponential (or flat or whatever), is there a way to have something like

flow = masked_autoregressive_flow(
    subkey,
    base_dist=[Normal(jnp.zeros(1)),  Exponential(1)],
)

Thank you very much in advance!

[danielward27]

Hello, right now there isn't a direct way to to this in the package, so you'd have to implement a custom distribution. A simple implementation for stacking together scalar distributions could be something like this:

from flowjax.flows import AbstractDistribution
from typing import ClassVar
import jax.random as jr
import jax.numpy as jnp

class StackDistribution(AbstractDistribution):
    """Stack scalar distributions."""
    dists: tuple[AbstractDistribution]
    shape: tuple[int, ...]
    cond_shape: ClassVar[None] = None

    def __init__(self, *dists: AbstractDistribution):
        for dist in dists:
            if dist.shape != () or dist.cond_shape is not None:
                raise ValueError("Only scalar unconditional distributions supported.")

        self.dists = dists
        self.shape = (len(dists),)

    def _log_prob(self, x, condition=None):
        return sum(dist.log_prob(xi) for dist, xi in zip(self.dists, x))

    def _sample(self, key, condition=None):
        keys = jr.split(key, self.shape)
        return jnp.stack([dist.sample(k) for dist, k in zip(self.dists, keys)])

In certain cases it might be worth implementing a similar _sample_and_log_prob method, but this is probably not important for simple distributions. It probably makes sense to have this functionality in the package eventually, probably using slightly generalised version of the above, (with semantics similar to numpy.stack and numpy.concatenate).

One thing I'd note is that the support of the flow will not necessarily align with the support of the base distribution, if that is what you are expecting (it will depend on the transformer used).