openxla / shardy

MLIR-based partitioning system
Apache License 2.0
45 stars 7 forks source link

Use AxisRefsWithTail instead of ArrayRef of AxisRefs, so the tailing AxisRefAttr can be modified on factor-axis pairs. #205

Closed copybara-service[bot] closed 2 weeks ago

copybara-service[bot] commented 2 weeks ago

Use AxisRefsWithTail instead of ArrayRef of AxisRefs, so the tailing AxisRefAttr can be modified on factor-axis pairs.

It is to prepare the cases that we consider modified axes to assign to a factor, for example, a trimmed version of the existing axes. For example, if factor i is assigned to {"y":(4)2 }, then some other factor with a sharding {"x", "y", "z"} can now be assigned to {"x", "y":(1)4} and since the tailing axis "y":(1)4 is not in the original sharding sharding {"x", "y", "z"}, it can not be referred using ArrayRef, hence the tailing axis is AxisRefAttr.

This change is no-op, as the tailing axis is always empty.