Function replicate_to_local_data

axlearn/common/utils.py:658–667

Replicates and converts Tensors in `x` to local DeviceArrays. Args: x: The tensor to replicate. Returns: Replicated tensor.

(x: NestedTensor)

) -> NestedTensor:
    """Best-effort activation sharding hint: shards `x` by `partition_spec` if non-None.

    Manual axes (from an enclosing `shard_map`) in `partition_spec` are downgraded to
    `PartitionSpec.UNCONSTRAINED`; still-Auto axes apply as usual. Activation sharding is a
    hint XLA can override anyway, so silently dropping the Manual portion is safe. For
    parameters / state where a wrong axis is a real bug, use `with_sharding_constraint`.
    For `shard_map` `in_specs`/`out_specs` (which reject `UNCONSTRAINED`), use
    `manual_axes_to_none`.

    When all mesh axes are Manual (no Auto axes remain), `with_sharding_constraint` is
    skipped entirely — there's nothing left for the compiler to constrain, and the call
    would be a no-op wrapper at best (see `_manual_axes_to_unconstrained`).

    `partition_spec` may be a `PartitionSpec` or a flat sequence of axis entries (wrapped