jax.experimental.maps module# API# xmap(fun, in_axes, out_axes, *[, ...]) Assign a positional signature to a program that uses named array axes.