jax

2d0df4d7 - Go via device_put in reshard's impl rule if the mesh in context and the mesh is not multi-process