import jax
import jax.numpy as jnp
from jax.experimental.maps import xmap

# dummy function
def func(cond, data):
   if cond: return data
   else: return data

# vectorisation
func_vmap = jax.vmap(func, in_axes=[None, 0], out_axes=0)
func_xmap = xmap(func, in_axes=[[...], ['axis']], out_axes=['axis'])

# jit compiling
func_vmap_jit = jax.jit(func_vmap, static_argnames=['cond'])
func_xmap_jit = jax.jit(func_xmap, static_argnames=['cond'])

# running
cond = True
data = jnp.ones(100)
out_vmap = func_vmap_jit(cond, data)
out_xmap = func_xmap_jit(cond, data)