Applying functions across Tensor dimensions similar to JAX’s vmap

[JR-1991]

Is there a way to apply an arbitrary function across a tensor dimension like Jax’s vmap? I haven’t found this feature in the documentation and might have looked for the wrong term.

I need to apply an ODE solver to the batch dimension of a tensor of initial conditions to generate a new tensor with integrated values. Currently, I’m integrating over the entire time domain for all initial conditions, but managing adaptive step sizes is challenging since they’re dictated by the stiffest trajectory.

If this functionality isn’t available, what are the best practices to achieve it? Thanks in advance!