Class FanAxes

axlearn/common/param_init.py:24–135

FanAxes describes axis indices corresponding to input, output, and batch axes. Note: axes not listed in {in,out,batch}_axis are assumed to be the “receptive field” (convolution kernel spatial axes). Used for https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.initializers.vari

class FanAxes(NamedTuple):
    """FanAxes describes axis indices corresponding to input, output, and batch axes.

    Note: axes not listed in {in,out,batch}_axis are assumed to be the “receptive field”
    (convolution kernel spatial axes).

    Used for
    https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.initializers.variance_scaling.html.
    """

    class AxisType(Enum):
        IN_AXIS = "in_axis"
        OUT_AXIS = "out_axis"
        BATCH_AXIS = "batch_axis"
        NONE = None

    # Input axis or sequence of axes of the fan "input" dimension.
    in_axis: Union[tuple[int, ...], int]
    # Output axis or sequence of axes of the fan "output" dimension.
    out_axis: Union[tuple[int, ...], int]
    # Batch axis or sequence of axes that should be ignored when computing fan.
    batch_axis: Union[tuple[int, ...], int] = ()

    def canonicalize(self) -> "FanAxes":
        """Returns a FanAxes equivalent to this one where all fields are tuples."""

        def canonicalize(maybe_tuple: Union[None, int, Sequence[Union[int]]]) -> tuple[int, ...]:
            if maybe_tuple is None:
                return tuple()
            if isinstance(maybe_tuple, int):
                return (maybe_tuple,)
            if isinstance(maybe_tuple, tuple):
                return tuple(sorted(maybe_tuple))
            raise TypeError(f"Invalid type {type(maybe_tuple)} for data {maybe_tuple}.")

        axes = {}
        for typ in self._fields:
            axes[typ] = canonicalize(getattr(self, typ))
        return FanAxes(**axes)

    def __eq__(self, other):
        if not isinstance(other, FanAxes):
            return False
        return self.canonicalize()._asdict() == other.canonicalize()._asdict()

    # We can't insert into other locations besides the start and end
    # because FanAxes doesn't know the number of dimensions of the parameter shape.
    def _insert_axis(self, axis: int, *, axis_type: "FanAxes.AxisType") -> "FanAxes":
        """Returns a copy of this where a new axis of the given type has been inserted
        at the given index, moving the other axes accordingly.

        Only 0 and -1 are supported for locations.

        Args:
            axis: The index of the new axis.
            axis_type: The type of axis to insert.

        Returns: