Change the basis of an index of a tensor in numpy, by contracting with a matrix and reordering the resulting indices

tensor_change_basis.py

import numpy as np
from numpy.typing import NDArray

def tensor_change_basis(tensor: NDArray, index: int, basis_matrix: NDArray) -> NDArray:
    """
    Contract the arbitrary rank tensor `tensor` T_ijk...x...lmn
    and the rank 2 tensor `basis_matrix` M_wx
    to return T'_ijk...w...lmn, where x is the `index`-th index of `tensor`
    """
    rank = len(tensor.shape)
    permutation = tuple(range(1, index + 1)) + (0,) + tuple(range(index + 1, rank))
    # `np.tensordot` returns T'_wijk...lmn; we need to rearrange the indices with `permutation`
    return np.tensordot(basis_matrix, tensor, (1, index)).transpose(permutation)