spekk.transformations.for_all.python_vmap

Contents

spekk.transformations.for_all.python_vmap#

spekk.transformations.for_all.python_vmap(f, in_axes)[source]#

A simple Python implementation of JAX’s jax.vmap() based on for-loops.