keras/src/ops/nn.py/BatchNorm/_check_shape
class BatchNorm:    def _check_shape(self, name, shape, expected_shape):
        if shape != expected_shape:
            raise ValueError(
                f"Arguments `{name}` must be a vector of length "
                f"`x.shape[axis]`. Expected: `{expected_shape}`. "
                f"Received: `{shape}."
            )