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