Threading in random key to `AbstractBijection.{transform_and_log_det}` and relatives

[amanjitsk]

Hi,

How would you recommend threading in random keys for bijection methods such as transform_and_log_det and others? Or is that not currently possible? I figure this would facilitate approximate computation of log determinants using stochastic estimators.

[danielward27]

Maybe something like this?

from typing import ClassVar

import equinox as eqx
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import Array, ArrayLike, PRNGKeyArray

from flowjax.bijections import AbstractBijection


class LocWithRandomLogDet(AbstractBijection):
    """Location transform with random log det (for illustration)."""

    shape: tuple[int, ...]
    cond_shape: ClassVar[None] = None
    loc: Array
    key: PRNGKeyArray

    def __init__(
        self,
        key: PRNGKeyArray,
        loc: ArrayLike = 0,
    ):
        self.loc = jnp.asarray(loc)
        self.shape = self.loc.shape
        self.key = key

    def transform_and_log_det(self, x, condition=None):
        y = x + self.loc
        stochastic_log_det = jr.normal(self.key)
        return y, stochastic_log_det

    def inverse_and_log_det(self, y, condition=None):
        x = y - self.loc
        stochastic_log_det = jr.normal(self.key)
        return x, stochastic_log_det

    def set_key(self, key: PRNGKeyArray):
        return eqx.tree_at(lambda b: b.key, self, key)


key = jr.key(0)
bijection = LocWithRandomLogDet(key, loc=2)

for _ in range(3):
    key, subkey = jr.split(key)
    bijection = bijection.set_key(subkey)
    y, log_det = bijection.transform_and_log_det(1)
    print(y, log_det)


# If you have these in a complex model and don't want to manually update the seeds,
# we can map over the model.

import jax

from flowjax.bijections import Chain
from flowjax.distributions import Normal, Transformed


def update_keys(model, key):
    # Generate key for each LocWithRandomLogDet with same tree structure
    treedef = jax.tree_util.tree_structure(
        model,
        is_leaf=lambda leaf: isinstance(leaf, LocWithRandomLogDet),
    )
    keys = jax.random.split(key, treedef.num_leaves)
    keys = jax.tree_util.tree_unflatten(treedef, keys)

    # Set the keys
    return jax.tree_util.tree_map(
        lambda l, key: l.set_key(key) if isinstance(l, LocWithRandomLogDet) else l,
        model,
        keys,
        is_leaf=lambda leaf: isinstance(leaf, LocWithRandomLogDet),
    )


dist = Transformed(
    base_dist=Normal(),
    bijection=Chain([LocWithRandomLogDet(key, loc=2), LocWithRandomLogDet(key, loc=2)]),
)

for _ in range(3):
    key, subkey = jr.split(key)
    dist = update_keys(dist, subkey)
    print(dist.log_prob(1))

To use the training loops, you'd have to modify the loss function used to make sure you update the key before each call to transform_and_log_det/inverse_and_log_det.

[amanjitsk]

Oh wow these are some great ideas! I was also leaning toward the tree approach but did not want to change the code to manually update the keys. Instead I was thinking of just using the key as the condition, or perhaps even re-defining bijections with an optional third parameter for the key (this way I can have a consistent call signature and pass in the key).

[danielward27]

Looking back, you'd have to be careful with my suggestion e.g. log_prob over a batch would share a key, which I imagine wouldn't be what you want (so you'd e.g. have to vmap over keys and data to avoid that). At least in the main package, I don't think I'd add a key argument to every bijection, just because of the additional clutter and compexity - you'd have to add key splitting logic in a lot of places to support it package wide. Passing as a condition could be used (again be careful of unintended key reuse), but I think the challenge again is you still need to manually add the key splitting whatever method is used, which may be easy or hard depending on the use case.

[amanjitsk]

Yep I figured! I don't think passing the key in as condition is that great of a solution either, especially in the presence of an actual condition that is used.

[amanjitsk]

Coming back to this again, it was hard (and often tedious) to modify the AbstractBijection class to allow for stochastic transformations (in either the map itself, or the log determinants). This was eventually going to be a footgun anyway, as equinox does not have the mutability that NNX has now.

But I do think we could modify condition to make it a PyTree (with array leaves) instead of an array, and this would allow for passing in arbitray number of arrays, including keys. I think this would also make the bijection object more greatly more flexible, by allowing for arbitrary condition objects that may be used. For instance,

from jax import tree_util as jtu
import equinox as eqx
from jaxtyping import Array, PRNGKeyArray

class AbstractCondition(eqx.Module):
    key:  eqx.AbstractVar[PRNGKeyArray]

    def eval_shape(self):
        """Helper for evaluating shapes to pass in to AbstractBijection."""
        return jtu.tree_map(self, lambda x: x.shape, is_leaf=eqx.is_array) # or something to the effect of this

and subsequent condition objects can inherit from AbstractCondition.

Then, all the shape checking for vectorization (and for compatibility of bijections, and distributions etc.) can be replaced with their jtu.tree_map counterparts. What do you think about this modification?

I am happy to work on a pull request if you don't see any major hurdles with this change, but I feel like this could be a worthwhile change. Of course, one downside is that nesting condition arrays into PyTrees might slow down things a bit, but the trade-off is that we can pass all sorts of array (-like) objects into the condition without having to change code in multiple places.

Do you anticipate any major downsides if a change like this were to be made? One thing that I thought of is that it makes passing in the cond_shape a bit more involved (and something like the eval_shape method could be used), but once we do have that PyTree of shapes, then everything else downstream would only require a careful application of jtu.tree_map.

[danielward27]

There are definitely times where a PyTree for the condition would be nice. My initial thoughts are

• Would you expect special handling of keys, or just treat them as another array? e.g. should dist.log_prob, if given a single key in condition, but a batch of x, internally split the key, or just broadcast the same key? Unless I'm missing something, I think you'd need special treatment at least in the loss functions to avoid key reuse?
• I think the important thing to know is the canonical (i.e. not batched) shape(s) for the condition, so you'd replace cond_shape with some pytree of tuples, denoting the canonical shape (which would still just be a tuple[int, ...] for simple array conditions). I do wonder though if the broadcasting for PyTree inputs could become a point of confusion.

I'd definitely take a look at a pull request for it, but without spending a bit more time thinking about the implications, it's hard to say whether I'd merge it.

[amanjitsk]

Thanks for the quick response. Here is what I think:

• As it concerns special handling of keys, no. It would be identical to the way it is now: Simply calling the vectorized version as in

  flowjax/flowjax/distributions.py

  Line 107 in cabcb42

  return self._vectorize(self._log_prob)(x, condition)

  This would just mean the caller should pass in a vectorized condition (along all leaves by default). Of course we could always provide more options like a PyTree of bools which specifies whether or not to vectorize the array in the corresponding leaf, but this is probably overkill for now (something like the in_axes specification for eqx.filter_vmap).

On the other hand, I do see what you are saying, but my initial thought was why use jnp.vectorize instead of jax.vmap (which only adds one dimension to the output). Of course, there might be scenarios where this would be useful (especially if we'd like conditional generation for multiple inputs across multiple conditions), but I would think that it'd be best to just align the condition and x arrays using vmap instead (with an additional call to reshape prior to and after calling vmap if needed).

• It probably will appear confusing at first as it concerns broadcasting of a PyTree, but I think that could be partially resolved if we replace jnp.vectorize with jnp.vmap in the dist.log_prob. Of course, this also means that the VectorizedBijection and other vectorized operations might not necessarily align with this.

But I think a simple (but tedious) resolution to all the above would be either to:

• to replace jnp.vectorize with jnp.vmap, or better yet eqx.filter_vmap. This way we could allow for more control over which leaves to vmap over using in_axes for the condition, while retaining all of its functionality.
• or to instead replace the input x with a PyTree for which all arrays share the same leading dimensions when querying dist.log_prob (although this would require rethinking the entire bijection business, with in_shape and out_shape attributes, and also seems overkill).

This is just what I thought of immediately, but I'd have to think about this more. Also, with the possibility of passing in PyTree inputs or condition, we will have to think more carefully about whether vectorize is perhaps too unpredictable, and instead vmap with an in_axes is preferred. But I understand obviously this adds more complexity to the code. It might also make sense to look at how NNX implements its rng splitting features (obviously ignoring the fact that they allow for mutable PyTrees, which can be addressed with eqx.tree_at) for some ideas. Something as simple as adding randomness should be supported without the need for adding some much more code on a case-by-case basis (for the end user).

What do you think?

[danielward27]

So, the advantage of jnp.vectorize is that it automatically handles the numpy-style broadcasting rules, which keeps things pretty simple in the implementation, e.g. automatically includes things like using many x values with a single condition, and vice versa, arbitrary batch dimensions, plus reasonable error messages. Another option, if condition is a pytree of arrays, you could probably flatten the condition pytree and use a little wrapper function taking the arrays as args to still make use of jnp.vectorize, I don't think that would be too tricky and could avoid breaking changes, but again the automatic broadcasting across a PyTree feels potentially problematic.

I would be against facilitating accepting anything like in_axes in distribution methods, because we can just wrap the method then use eqx.filter_vmap however we like, e.g. following from the above example, to vmap over batches of keys stored on the pytree:

batch_size = 10
random_transforms = []  # Each transform will have 10 keys
key, subkey = jr.split(jr.key(0))

for key in jr.split(subkey, 3):
    transform = LocWithRandomLogDet(key, 2)
    transform = eqx.tree_at(
        lambda x: x.key,
        transform,
        replace_fn=lambda key: jr.split(key, batch_size),
    )
    random_transforms.append(transform)

dist = Transformed(
    base_dist=Normal(),
    bijection=Chain(random_transforms),
)


def get_in_axes(dist):  # Map only over keys
    def map_fn(leaf):
        if isinstance(leaf, LocWithRandomLogDet):
            in_axes = jax.tree.map(lambda _: None, leaf)
            return eqx.tree_at(lambda x: x.key, in_axes, 0, is_leaf=lambda x: x is None)
        return None

    return jax.tree.map(
        map_fn,
        is_leaf=lambda x: isinstance(x, LocWithRandomLogDet),
        tree=dist,
    )


in_axes = get_in_axes(dist)


def log_prob(dist, x):
    return dist.log_prob(x)


# Vmap over keys, for one input
eqx.filter_vmap(log_prob, in_axes=(in_axes, None))(dist, 1)

# Vmap over keys, for same number of inputs
eqx.filter_vmap(log_prob, in_axes=(in_axes, 0))(dist, jnp.ones(10))

Sure, I get this isn't the most convenient thing, but I think there will always be some friction to the task of vectorization over a subset of parameters.

I do think there is merit to a stricter approach which avoids vectorize, and errors if batches are passed, which would e.g. avoid confusion if using PyTree conditions, or avoid accidental key reuse from broadcasting for a key stored on a distribution PyTree. In that instance, I wonder if the best bet would be to add broadcast-free methods which error on batches (e.g. log_prob_single and sample_single, and relying on the user to apply eqx.filter_vmap whenever they want.

I understand this isn't really providing any proper answers, so apologies for that! I'm not certain what the best approach is.

[danielward27]

A slightly tidier summary of the first approach is here:

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import PRNGKeyArray

from flowjax.bijections import AbstractBijection
from flowjax.distributions import Normal, Transformed


class Key(eqx.Module):
    """Simple wrapper, for facilitating filtering."""

    key: PRNGKeyArray


class RandomTransform(AbstractBijection):
    """Random location transform."""

    key: Key
    shape = ()
    cond_shape = None

    def __init__(self, key: PRNGKeyArray):
        self.key = Key(key)

    def transform_and_log_det(self, x, condition=None):
        return x + jr.normal(self.key.key, shape=x.shape), jnp.zeros(())

    def inverse_and_log_det(self, y, condition=None):
        return y - jr.normal(self.key.key, shape=y.shape), jnp.zeros(())


def update_keys(model, key: PRNGKeyArray):
    treedef = jax.tree_util.tree_structure(
        model,
        is_leaf=lambda leaf: isinstance(leaf, Key),
    )
    keys = jax.random.split(key, treedef.num_leaves)
    keys = jax.tree_util.tree_unflatten(treedef, keys)
    return jax.tree_util.tree_map(
        lambda x, key: Key(key) if isinstance(x, Key) else x,
        model,
        keys,
        is_leaf=lambda leaf: isinstance(leaf, Key),
    )


key, subkey = jr.split(jr.key(0))
dist = Transformed(Normal(), RandomTransform(key))


# Be careful, the keys will be shared in batched methods!
dist.sample(jr.key(1), (1000,)).mean()  # -2.48 as transform key is shared!


# But we can define wrappers to avoid this, e.g.
def sample(dist, key, n: int = 1):
    @eqx.filter_vmap(in_axes=(None, eqx.if_array(0)))
    def _sample(dist, key):
        sample_key, transform_key = jr.split(key)
        dist = update_keys(dist, transform_key)
        return dist.sample(sample_key)

    return _sample(dist, jr.split(key, n))



sample(dist, jr.key(0), n=1000).mean()  # -0.03

# We should be able to still make use of the training loop if we want,
# using a custom loss function
import paramax
from jaxtyping import Array, Float

from flowjax.train.loops import fit_to_data


@eqx.filter_jit
def maximum_likelihood_loss(
    params,
    static,
    x: Array,
    condition: Array | None = None,
    key: PRNGKeyArray | None = None,
) -> Float[Array, ""]:
    dist = paramax.unwrap(eqx.combine(params, static))
    in_axes = [0, None, 0, None if condition is None else 0]

    @eqx.filter_vmap(in_axes=in_axes)
    def _log_prob(key, dist, x, condition):
        dist = update_keys(dist, key)
        return dist.log_prob(x, condition)

    return -_log_prob(jr.split(key, x.shape[0]), dist, x, condition)


key, subkey = jr.split(key)
data = jr.normal(subkey, (5000,))


key, subkey = jr.split(key)
fit_to_data(key, dist, data)

Not saying this is the end result, but I think it's quite reasonable.

[amanjitsk]

Thank you! I appreciate the quick responses!

Ya, no worries this is all to brainstorm ideas anyway. I don't think there is a perfect solution, but your last piece of code seems reasonable. Its just that I thought it will scale much better if there was a way to thread keys in directly to the calls to Bijection and log_prob methods, but your last comment (I think) already suggests a fairly model-agnostic way of using keys.