Function input_partition_spec

axlearn/common/utils.py:743–755

Returns partition spec for the input batch. We partition the inputs along all axes. For example, if the mesh has shape (64, 4) and axis names of ("data", "model"), the partition spec will be (("data", "model"), None...) so that the batch axis of every global tensor will be partitioned 2

()

    def _promote(orig, filtered):
        # `orig` is the user's dim; `filtered` is `manual_axes_to_none`'s result. A `None`
        # in `filtered` only becomes UNCONSTRAINED if it came from a Manual axis (or a
        # tuple of Manual axes) in `orig` — explicit `None` (replicate) must pass through.
        if filtered is not None:
            return filtered
        if isinstance(orig, str):
            return PartitionSpec.UNCONSTRAINED if orig in manual_axes else None
        if isinstance(orig, tuple):
            return PartitionSpec.UNCONSTRAINED if all(a in manual_axes for a in orig) else None
        return None

    filtered = manual_axes_to_none(spec)
    return PartitionSpec(*(_promote(o, f) for o, f in zip(spec, filtered)))


def get_current_abstract_or_physical_mesh() -> jax.sharding.AbstractMesh | jax.sharding.Mesh:
    """Returns the current abstract mesh if it's set, or the physical mesh otherwise."""