flax.linen.logical_to_mesh_sharding

flax.linen.logical_to_mesh_sharding#

flax.linen.logical_to_mesh_sharding(tree, mesh, rules=None)[source]#

Convert pytrees of logical PartitionSpecs to shardings.