spmd# flax.nnx.get_partition_spec(tree)[source]# Extracts a PartitionSpec tree from a PyTree containing Variable values. flax.nnx.get_named_sharding(tree, mesh)[source]# flax.nnx.with_partitioning(initializer, sharding, mesh=None, **metadata)[source]# A wrapper over any initializer to add sharding annotation data to a Variable.