def output_axes(self) -> dict[str, int]:
    """The axis names and sizes that will appear in the output only."""
    return {  # pytype: disable=bad-return-type
        name: size
        for name, size in self.weights.value.named_shape.items()
        if name in self.out_axis_names
    }
The axis names and sizes that will appear in the output only.
penzai/nn/linear_and_affine.py