Apply a function to 1D array slices along an axis.
This is implemented internally with applications of jax.vmap(), rather than iterating the function over a loop.
jax.vmap()
Apply a function to 1D array slices along an axis.
This is implemented internally with applications of
jax.vmap(), rather than iterating the function over a loop.