def linear_transform(inp: Tensor, transform_matrix: Tensor, mean_vector: Tensor, dim: int = 0) -> Tensor:
    r"""Given a transformation matrix and a mean vector, this function will flatten the input tensor along the given
    dimension and subtract the mean vector from it. Then the dot product with the transformation matrix will be
    computed and then the resulting tensor is reshaped to the original input shape.

    .. math::

        \mathbf{X}_{T} = (\mathbf{X - \mu})(T)

    Args:
        inp: Input data :math:`X`.
        transform_matrix: Transform matrix :math:`T`.
        mean_vector: mean vector :math:`\mu`.
        dim: Batch dimension.

    Shapes:
        - inp: :math:`(D_0,...,D_{\text{dim}},...,D_N)` is a batch of N-D tensors.
        - transform_matrix: :math:`(\Pi_{d=0,d\neq \text{dim}}^N D_d, \Pi_{d=0,d\neq \text{dim}}^N D_d)`
        - mean_vector: :math:`(1, \Pi_{d=0,d\neq \text{dim}}^N D_d)`

    Returns:
        Transformed data.

    Example:
        >>> # Example where dim = 3
        >>> inp = torch.ones((10,3,4,5))
        >>> transform_mat = torch.ones((10*3*4,10*3*4))
        >>> mean = 2*torch.ones((1,10*3*4))
        >>> out = linear_transform(inp, transform_mat, mean, 3)
        >>> print(out.shape, out.unique())  # Should a be (10,3,4,5) tensor of -120s
        torch.Size([10, 3, 4, 5]) tensor([-120.])

        >>> # Example where dim = 0
        >>> inp = torch.ones((10,2))
        >>> transform_mat = torch.ones((2,2))
        >>> mean = torch.zeros((1,2))
        >>> out = linear_transform(inp, transform_mat, mean)
        >>> print(out.shape, out.unique()) # Should a be (10,2) tensor of 2s
        torch.Size([10, 2]) tensor([2.])

    """  # noqa: D205
    inp_size = inp.size()

    if dim >= len(inp_size) or dim < -len(inp_size):
        raise IndexError(
            f"Dimension out of range (expected to be in range of [{-len(inp_size)},{len(inp_size) - 1}], but got {dim}"
        )

    if dim < 0:
        dim = len(inp_size) + dim

    feat_dims = concatenate([torch.arange(0, dim), torch.arange(dim + 1, len(inp_size))])

    perm = concatenate([tensor([dim]), feat_dims])
    perm_inv = torch.argsort(perm)

    new_order: List[int] = perm.tolist()
    inv_order: List[int] = perm_inv.tolist()

    feature_sizes = tensor(inp_size[0:dim] + inp_size[dim + 1 : :])
    num_features: int = int(torch.prod(feature_sizes).item())

    inp_permute = inp.permute(new_order)
    inp_flat = inp_permute.reshape((-1, num_features))

    inp_center = inp_flat - mean_vector
    inp_transformed = inp_center.mm(transform_matrix)

    inp_transformed = inp_transformed.reshape(inp_permute.size())

    inp_transformed = inp_transformed.permute(inv_order)

    return inp_transformed
Given a transformation matrix and a mean vector, this function will flatten the input tensor along the given
    dimension and subtract the mean vector from it. Then the dot product with the transformation matrix will be
    computed and then the resulting tensor is reshaped to the original input shape.

    .. math::

        \mathbf{X}_{T} = (\mathbf{X - \mu})(T)

    Args:
        inp: Input data :math:`X`.
        transform_matrix: Transform matrix :math:`T`.
        mean_vector: mean vector :math:`\mu`.
        dim: Batch dimension.

    Shapes:
        - inp: :math:`(D_0,...,D_{\text{dim}},...,D_N)` is a batch of N-D tensors.
        - transform_matrix: :math:`(\Pi_{d=0,d\neq \text{dim}}^N D_d, \Pi_{d=0,d\neq \text{dim}}^N D_d)`
        - mean_vector: :math:`(1, \Pi_{d=0,d\neq \text{dim}}^N D_d)`

    Returns:
        Transformed data.

    Example:
        >>> # Example where dim = 3
        >>> inp = torch.ones((10,3,4,5))
        >>> transform_mat = torch.ones((10*3*4,10*3*4))
        >>> mean = 2*torch.ones((1,10*3*4))
        >>> out = linear_transform(inp, transform_mat, mean, 3)
        >>> print(out.shape, out.unique())  # Should a be (10,3,4,5) tensor of -120s
        torch.Size([10, 3, 4, 5]) tensor([-120.])

        >>> # Example where dim = 0
        >>> inp = torch.ones((10,2))
        >>> transform_mat = torch.ones((2,2))
        >>> mean = torch.zeros((1,2))
        >>> out = linear_transform(inp, transform_mat, mean)
        >>> print(out.shape, out.unique()) # Should a be (10,2) tensor of 2s
        torch.Size([10, 2]) tensor([2.])

    
kornia/enhance/zca.py