augax._dispatch

augax._dispatch = <jax._src.custom_batching.custom_vmap object>[source]