ragged arrays in jax

raax.py

import jax
import chex

@chex.dataclass
class OptionalPyTree:
  array: jax.Array
  exists: bool

def make_optional_pytree(array, exists):
    return OptionalPyTree(
        array=array,
        exists=exists
    )
    
x = jax.numpy.array([[1., 2., 3.], [4., 5., 6.]])
y = jax.vmap(make_optional_pytree, in_axes=(0, 0))(jax.numpy.zeros((2, 3)), jax.numpy.array([False, True]))

def optional_vmap(func, zero_value, in_axes=0, out_axes=0):
    def optional_apply(*args):
        args = tuple(
            make_optional_pytree(arg, True) if not isinstance(arg, OptionalPyTree) else arg
            for arg in args
        )

        exists = jax.numpy.all(
            jax.numpy.array(
                jax.tree.leaves(
                    jax.tree.map(
                        lambda x: x.exists,
                        args,
                        is_leaf=lambda tree: isinstance(tree, OptionalPyTree)
                    )
                )
            )
        )
        return jax.lax.cond(
            exists,
            lambda: func(
                *(jax.tree.map(
                    lambda x: x.array,
                    args,
                    is_leaf=lambda tree: isinstance(tree, OptionalPyTree)
                ))
            ),
            lambda: zero_value
        )
    return jax.vmap(optional_apply, in_axes=in_axes, out_axes=out_axes)

print(optional_vmap(lambda x, y: x + y, jax.numpy.array([0., 0., 0.]))(x, y))