Validate replicated axes orthogonality in dump comparator (#21026)
This commit is contained in:
@@ -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],
|
||||
|
||||
Reference in New Issue
Block a user