Skip to main content
Ctrl+K

flax.linen.logical_to_mesh_sharding#

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

Convert pytrees of logical PartitionSpecs to shardings.

Parameters:
  • tree (Any) –

  • mesh (Mesh) –

  • rules (Optional[Sequence[Tuple[str, Union[str, Tuple[str], None]]]]) –

Return type:

Any