Add Leafwise bijector#83
Conversation
Leafwise applies a tree of bijectors to a matching tree of inputs, analogous to TensorFlow's JointMap and the pytree-level opposite of Chain (which composes bijectors sequentially on a single input).
| return isinstance(node, AbstractBijector) | ||
|
|
||
|
|
||
| class Leafwise(AbstractFwdLogDetJacBijector, AbstractInvLogDetJacBijector): |
There was a problem hiding this comment.
Can you clarify the specific use case for this style? I'm just thinking in general, pytrees can be handled by the user and ideally we can make things work for pytrees naturally, without requiring additional classes. I guess this is just doing that as a utility like jax.tree.map(lambda bi, x: bi(x), bijectors, inputs?
There was a problem hiding this comment.
By this style do you mean these types of classes? Essentially it is to prevent a ton of boilerplate and if conditions for pytree "power users".
Leafwise specifically is very similar to a Joint distribution. A user may have a tree of distributions that encapsulate a bijective transformation over another tree. Without this class:
- Maintaining semantic clarity becomes difficult. Instead of the user writing
bijector.forward(tree)they have to manually map over the leaves, even though the equivalent operation is semantically represented by a single bijector - Abstraction layers become muddled. For example, lets say a user actually has some bijector that caters for the trees shape. Now they have to use an if condition to check "if they have a pytree of bijectors" or a "single bijector" which complicates code very quickly
It would be similar to having to manually sum log probs for a Joint distribution: technically the user could do it themselves (maybe even implementing the class manually), but it probably should just be in distreqx itself.
This is why I think TensorFlow provides classes like JointMap - it is not that the user cannot accomplish the same task themselves (as you've demonstrated): it is to keep the abstraction layer clean. Otherwise classes like Chain should technically also not exist (as the user could just write a for loop while carrying through log dets).
If you would like a concrete example, have a look at the second codeblock in Parax's documentation.
One more such utility class (not yet submitted for PR but one I've found I will need) is a Combined distribution, which is essentially eqx.combine over a tree of distributions with some none-leaves. Again, technically doable by an end-user, but provides much better pytree enablement if part of the base library.
There was a problem hiding this comment.
I see, makes sense, lets make sure this is clear in class docstring
| return isinstance(node, AbstractBijector) | ||
|
|
||
|
|
||
| class Leafwise(AbstractFwdLogDetJacBijector, AbstractInvLogDetJacBijector): |
There was a problem hiding this comment.
I see, makes sense, lets make sure this is clear in class docstring
|
|
||
| def __init__(self, bijectors: PyTree[AbstractBijector]): | ||
| """Initializes a TreeMap bijector.""" | ||
| leaves = jax.tree_util.tree_leaves(bijectors, is_leaf=_is_bijector) |
There was a problem hiding this comment.
should we check to make sure log det shapes match? e.g. Leafwise({"a": Shift(jnp.ones(4)), "b": Shift(jnp.ones(2))}) raises TypeError: add got incompatible shapes for broadcasting: (4,), (2,) on forward_log_det_jacobian.
Wrapping them in Block would work tho.
|
|
||
| def __init__(self, bijectors: PyTree[AbstractBijector]): | ||
| """Initializes a TreeMap bijector.""" | ||
| leaves = jax.tree_util.tree_leaves(bijectors, is_leaf=_is_bijector) |
2 participants
Adds a
Leafwisebijector, analogous to tensorflow's JointMap and opposite toChain.This applies a tree of bijectors to a tree of inputs.