jax

2eb1888c - Make the vmap(jit) or vmap(wsc) with a concrete layout error more informative