Returns a FanAxes equivalent to this one where all fields are tuples.
(self)
| 45 | batch_axis: Union[tuple[int, ...], int] = () |
| 46 | |
| 47 | def canonicalize(self) -> "FanAxes": |
| 48 | """Returns a FanAxes equivalent to this one where all fields are tuples.""" |
| 49 | |
| 50 | def canonicalize(maybe_tuple: Union[None, int, Sequence[Union[int]]]) -> tuple[int, ...]: |
| 51 | if maybe_tuple is None: |
| 52 | return tuple() |
| 53 | if isinstance(maybe_tuple, int): |
| 54 | return (maybe_tuple,) |
| 55 | if isinstance(maybe_tuple, tuple): |
| 56 | return tuple(sorted(maybe_tuple)) |
| 57 | raise TypeError(f"Invalid type {type(maybe_tuple)} for data {maybe_tuple}.") |
| 58 | |
| 59 | axes = {} |
| 60 | for typ in self._fields: |
| 61 | axes[typ] = canonicalize(getattr(self, typ)) |
| 62 | return FanAxes(**axes) |
| 63 | |
| 64 | def __eq__(self, other): |
| 65 | if not isinstance(other, FanAxes): |