A deprecated alias for `host_to_global_array`. Please use `host_to_global_array` instead.
(*args, **kwargs)
| 927 | |
| 928 | def data_partition_type_to_spec( |
| 929 | partition: Union[DataPartitionType, Nested[PartitionSpec]], |
| 930 | ) -> Nested[PartitionSpec]: |
| 931 | """Returns a PartitionSpec for the given partition type.""" |
| 932 | if partition == DataPartitionType.FULL: |
| 933 | return input_partition_spec() |
| 934 | elif partition == DataPartitionType.REPLICATED: |
| 935 | return PartitionSpec(None) |
| 936 | elif partition == DataPartitionType.BATCH: |
| 937 | mesh = thread_resources.env.physical_mesh |
| 938 | batch_axis_names = tuple(name for name in mesh.axis_names if name in _BATCH_AXES) |
| 939 | seq_axis_names = tuple(name for name in mesh.axis_names if name in _SEQ_AXES) |
| 940 | return PartitionSpec(batch_axis_names, seq_axis_names) |