Any examples of model parallelism using flax + jax > 0.4.1's sharding APIs?

[tensorpro]

I've seen the siren model parallelism example, but it uses an older version of jax.

I was wondering if there's any more recent examples that use the newer APIs?
And also, I was wondering whether we still need to re-define Flax modules with named axes like the siren colab does?
It kind of seems like we can use existing flax modules with parameter sharding in the initializers.

[tensorpro]

I guess newer examples aren't really necessary, since all that really changes is where we import Mesh / PartitionSpec from?

And it does seem like as long as the params are sharded in the variable dictionary, we can still re-use flax modules?

[cgarciae]

After a quick search, none of our current examples use pjit, Mesh, or PartitionSpec so I guess there is nothing to update 😅. The Scale up Flax Modules on multiple devices with pjit is our best approach on the topic, @IvyZX might release an example using the new jax.Array sharding APIs in the future.

[tensorpro]

Gotcha, thanks for the pointer Cristian!

On a tangential note, the docs there talk about using flax + pjit while some of Jax's pjit docs now advocate using jax.jit. Could the flax docs be updated to also include how things should be done using the jax.jit api?