A nice simple walkthrough in this post, but would be nice if it was updated to show how to do this with sharding and the new jax.Array type introduced not too l

A nice simple walkthrough in this post, but would be nice if it was updated to show how to do this with sharding and the new jax.Array type introduced not too long ago https://github.com/google/jax/pull/11233/files https://github.com/google/jax/pull/11233/files