maxtext
maxtext copied to clipboard
Support `allow_split_physical_axes` for multislice hybrid meshes such as FSDP on ICI and DCN.
allow_split_physical_axes is only supported for device meshes atm but we also should support this for hybrid meshes. This is useful when we want to use FSDP across DCN and ICI + split the physical axis.