Validate replicated axes orthogonality in dump comparator (#21026)

This commit is contained in:
fzyzcjy
2026-03-20 22:02:40 +08:00
committed by GitHub
parent 2f01950a0e
commit cc22601d28
3 changed files with 261 additions and 0 deletions

View File

@@ -138,6 +138,11 @@ def _validate_explicit_replicated(
f"Axes {{{conflict_names}}} declared as both sharded and replicated"
)
_validate_replicated_axes_orthogonal(
explicit_replicated_axes=explicit_replicated_axes,
parallel_infos=parallel_infos,
)
candidate_axes: set[ParallelAxis] = (
all_axes - sharded_axes - explicit_replicated_axes
)
@@ -178,6 +183,33 @@ def _validate_explicit_replicated(
)
def _validate_replicated_axes_orthogonal(
*,
explicit_replicated_axes: frozenset[ParallelAxis],
parallel_infos: list[dict[ParallelAxis, AxisInfo]],
) -> None:
"""Every pair of explicitly replicated axes must be fully orthogonal (no dependency)."""
axes: list[ParallelAxis] = sorted(explicit_replicated_axes, key=lambda a: a.value)
if len(axes) < 2:
return
violations: list[str] = []
for i, axis_a in enumerate(axes):
for axis_b in axes[i + 1 :]:
for parent, child in [(axis_a, axis_b), (axis_b, axis_a)]:
if _is_dependent_axis(parallel_infos, parent=parent, child=child):
violations.append(
f"'{parent.value}' determines '{child.value}'"
f"remove '{child.value}:replicated'"
)
if violations:
details = "; ".join(violations)
raise ValueError(
f"Explicitly-replicated axes overlap (not orthogonal): {details}"
)
def _validate(
*,
axes_to_validate: set[ParallelAxis],