Diff Coverage

        dp_mode = self._reconcile_mode_with_context(mesh_context, dp_mode)
        self.pass_config.dp_mode = dp_mode

        if dp_mode == "ddp":
            dim = self._first_dim(names, ("fsdp_replicate", "dp_replicate", "dp"))
            if dim is None:
                raise ValueError(
                    "dp_mode='ddp' needs a replicate axis among "
                    "('fsdp_replicate', 'dp_replicate', 'dp') in the mesh "
                    f"dims {names}"
                )
            sub = fsdp_mesh[dim]
            self._register_mesh_group("dp_replicate", sub)
            self.pass_config.dp_replicate_degree = sub.size()
            return

        if dp_mode == "hsdp":
            rep_dim = self._first_dim(names, ("fsdp_replicate", "dp_replicate"))
            shard_dim = self._first_dim(names, ("fsdp_shard", "fsdp", "dp"))

        if dp_mode == "hsdp":
            rep_dim = self._first_dim(names, ("fsdp_replicate", "dp_replicate"))
            shard_dim = self._first_dim(names, ("fsdp_shard", "fsdp", "dp"))
            if rep_dim is None or shard_dim is None:
                raise ValueError(
                    "dp_mode='hsdp' needs both a replicate axis "
                    "('fsdp_replicate'/'dp_replicate') and a shard axis "
                    f"('fsdp_shard'/'fsdp'/'dp') in the mesh dims {names}"
                )

            raise ValueError(
                "dp_mode='ddp' requires a non-sharding mesh "
                f"(dp_shard_size=1), got dp_shard_size={shard}; use 'hsdp'"
            )
        return dp_mode

    @staticmethod
    def _first_dim(names: tuple, candidates: tuple) -> Optional[str]:
        """Return the first candidate axis present in ``names`` (or None)."""

        """Return the first candidate axis present in ``names`` (or None)."""
        for cand in candidates:
            if cand in names:
                return cand
        return None

    @staticmethod
    def _register_mesh_group(name: str, sub_mesh: Any) -> None:
        """Register ``sub_mesh``'s process group under ``name``."""

        """
        shard = self.pass_config.fsdp_degree
        replicate = self.pass_config.dp_replicate_degree
        if shard is None and replicate is None:
            raise ValueError(
                "dp_mode='hsdp' requires fsdp_degree and/or "
                "dp_replicate_degree to build the 2-D (dp_replicate, fsdp) mesh"
            )
        if shard is None:
            if world_size % replicate != 0:
                raise ValueError(
                    f"world_size ({world_size}) not divisible by "
                    f"dp_replicate_degree ({replicate})"
                )
            shard = world_size // replicate
        if replicate is None:
            if world_size % shard != 0:
                raise ValueError(
                    f"world_size ({world_size}) not divisible by fsdp_degree ({shard})"
                )
            replicate = world_size // shard
        if shard * replicate != world_size:
            raise ValueError(
                f"hsdp degrees fsdp_degree={shard} x "
                f"dp_replicate_degree={replicate} != world_size={world_size}"
            )
        return shard, replicate

        if self.pass_config.fsdp_enabled:
            kwargs["fsdp_group_name"] = "fsdp"
            # DDP / HSDP also reduce gradients on the replicate axis.
            if normalize_dp_mode(self.pass_config.dp_mode) in ("ddp", "hsdp"):
                kwargs["dp_replicate_group_name"] = "dp_replicate"

        return kwargs

    def _run_graph(self, **inputs):

    Raises:
        ValueError: When ``mode`` is not a string or is unknown.
    """
    if not isinstance(mode, str):
        raise ValueError(f"dp_mode must be a string, got {type(mode).__name__}")
    key = mode.strip().lower()
    return _DP_MODE_ALIASES.get(key, key)