How to use shard_map in nnx modules

[acec2127]

Is it possible to use shard_map on the call method? I would like to have full control on model's weights and data shadings. Thanks

[cgarciae]

Hey! We haven't added nnx.shard_map (should be easy), but in the mean time you can use split, merge, and update to interact with any JAX API. You can use the pattern shown there inside a method for example.

[cgarciae]

I've create #4261 to track this.

[acec2127]

Many thanks for this quick answer. I haven't had the time to dig too deeply in the source code of NNX, but from what I understand a NNX module behaves similarly as a equinox module. Hence, if my reasoning is right, 'self' should reference to the actual module and I can simply decorate the call method with a shard_map and provide the right (in_specs, out_specs)?

Now, in equinox the module is itself is a pytree, so I can just take the module, partition it to obtain the parameters and then replace the leaves of the parameters pytree with the desired PartitionSpec to obtain the in_specs. Would that work similarly for nnx modules?