a = np.random.randn((2, 2, 3, 6))
b = np.random.randn((2, 2, 3))
np.allclose(np.einsum("...ijk,...k->...ij", a, tensorsolve(a, b)), b) # tensorsolve(a, b).shape == (2, 6)

import numpy as np

def btensorsolve(a, b, axes=None):
    """
    Solve the tensor equation ``a x = b`` for x.

    It is assumed that all indices of `x` are summed over in the product,
    together with the rightmost indices of `a`, as is done in, for example,
    ``tensordot(a, x, axes=x.ndim)``.

    Parameters
    ----------
    a : array_like
        Coefficient tensor, of shape ``b.shape + Q``. `Q`, a tuple, equals
        the shape of that sub-tensor of `a` consisting of the appropriate
        number of its rightmost indices, and must be such that
        **there exists i that**
        ``prod(Q) == prod(b.shape[i:])``
    b : array_like
        Right-hand tensor, which can be of any shape.
    axes : tuple of ints, optional
        Axes in `a` to reorder to the right, before inversion.
        If None (default), no reordering is done.

    Returns
    -------
    x : ndarray, shape Q

    Raises
    ------
    LinAlgError
        If `a` is singular or not 'square' (in the above sense).

    See Also
    --------
    numpy.tensordot, tensorinv, numpy.einsum

    Examples
    --------
    >>> import numpy as np
    >>> rng = np.random.default_rng()
    >>> a = rng.normal(size=(2, 2*3, 4, 2, 3, 4))
    >>> b = rng.normal(size=(2, 2*3, 4))
    >>> x = np.linalg.tensorsolve(a, b)
    >>> x.shape
    (2, 2, 3, 4)
    >>> np.allclose(np.einsum('...ijklm,...klm->...ij', a, x), b)
    True

    """
    # https://github.com/numpy/numpy/blob/
    # e7a123b2d3eca9897843791dd698c1803d9a39c2/numpy/linalg/_linalg.py#L291
    an = a.ndim
    if axes is not None:
        allaxes = list(range(0, an))
        for k in axes:
            allaxes.remove(k)
            allaxes.insert(an, k)
        a = a.transpose(allaxes)

    # find right dimensions
    # a = [2 (dim1) 2 2 3 (dim2) 2 6]
    # b = [2 (dim1) 2 2 3 (dim2)]
    axis_sol_last = b.ndim
    if a.shape[:axis_sol_last] != b.shape:
        raise ValueError(
            f"Shapes of a and b are incompatible: " f"{a.shape} and {b.shape}"
        )

    # the dimention of the linear system
    sol_dim = np.prod(a.shape[axis_sol_last:])
    sol_dim_ = 1
    for axis_batch_last in range(axis_sol_last - 1, -1, -1):
        sol_dim_ *= a.shape[axis_batch_last]
        if sol_dim_ == sol_dim:
            break
    else:
        raise ValueError("Unable to divide batch dimensions and solution dimensions")

    a_ = a.reshape(a.shape[:axis_batch_last] + (sol_dim, sol_dim))
    b_ = b.reshape(b.shape[:axis_batch_last] + (sol_dim, 1))
    x = np.linalg.solve(a_, b_)
    return x.reshape(a.shape[:axis_batch_last] + a.shape[axis_sol_last:])