flax.linen.logical_to_mesh#

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

Applies logical_to_mesh_axes to pytrees of logical PartitionSpecs.