Sharding Vector Fields in Jax Program

[cwoolfo1]

I have written a computer simulation that uses several 3D vector fields stored as arrays. I am trying to parallelize my code across multiple cores. For local operations in other codes, they use "ghost cells" and sharded fields to perform localized operations on multiple cores.

If I initialize my Jax arrays as sharded fields on multiple cores, will Jax do this automatically? Or is there more to it than that?