`jax.vmap` Usage
==========================================


.. literalinclude:: ../../examples/vmap_example.py
   :language: python