batched_tensor_dot(tensor1, tensor2, contraction_modes, batched_modes)

def batched_tensor_dot_einsum(tensor1, tensor2, modes, batched_modes=()):
    """Tensor contraction between two tensors on specified modes
    
    Parameters
    ----------
    tensor1 : tl.tensor
    tensor2 : tl.tensor
    modes : int list or int
        modes on which to contract tensor1 and tensor2
    batched_modes : int or tuple[int]

    Returns
    -------
    contraction : tensor1 contracted with tensor2 on the specified modes
    """
    modes1, modes2 = _validate_contraction_modes(tensor1.shape, tensor2.shape, modes)
    batch_modes1, batch_modes2 = _validate_contraction_modes(tensor1.shape, tensor2.shape, batched_modes, batched_modes=True)

    start = ord('a')
    order_t1 = tl.ndim(tensor1)
    all_modes1 = [chr(start+i) for i in range(order_t1)]
    all_modes2 = [chr(start+i+order_t1) for i in range(tl.ndim(tensor2))]

    for m1, m2 in zip(modes1+batch_modes1, modes2+batch_modes2):
        all_modes2[m2] = all_modes1[m1]
    
    remaining_modes1 = [j for i, j in enumerate(all_modes1) if i not in modes1]
    remaining_modes2 = [j for i, j in enumerate(all_modes2) if i not in modes2+batch_modes2]
    remaining_modes = remaining_modes1 + remaining_modes2
    to_str = lambda x : ''.join(x)
    equation = f'{to_str(all_modes1)},{to_str(all_modes2)}->{to_str(remaining_modes)}'
    
    return tl.einsum(equation, tensor1, tensor2)


def batched_tensor_dot_matmul1(tensor1, tensor2, modes, batched_modes=()):
    modes1, modes2 = _validate_contraction_modes(tensor1.shape, tensor2.shape, modes)
    batch_modes1, batch_modes2 = _validate_contraction_modes(tensor1.shape, tensor2.shape, batched_modes, batched_modes=True)
    
    contraction_shape = [s for (i, s) in enumerate(tl.shape(tensor1)) if i in modes1]
    contraction_dim = prod(contraction_shape)
    batch_shape = [s for (i, s) in enumerate(tl.shape(tensor1)) if i in batch_modes1]
    
    # Prepare to reorganize the modes afterwards by moving bactch size back to their place
    # (while ommiting modes contracted over)
    final_modes = []
    n_batches = len(batch_modes1)
    batch_counter = 0
    free_counter = 0
    for i in range(tl.ndim(tensor1)):
        if i in modes1:
            continue
        elif i in batch_modes1:
            final_modes.append(batch_counter)
            batch_counter += 1
        else:
            final_modes.append(free_counter+n_batches)
            free_counter += 1

    # We will reorganize tensor1 to (batch_modes, new_modes1, contraction_modes)
    new_modes1 = [i for i in range(tensor1.ndim) if i not in batch_modes1+modes1]
    new_shape1 = [tl.shape(tensor1)[i] for i in new_modes1]
    tensor1 = tl.transpose(tensor1, batch_modes1 + new_modes1 + modes1)
    tensor1 = tl.reshape(tensor1, (*batch_shape, -1, contraction_dim))
    
    # Tensor2 will be (batch_modes, contraction_modes, new_modes2)
    new_modes2 = [i for i in range(tensor2.ndim) if i not in batch_modes2+modes2]
    new_shape2 = [tl.shape(tensor2)[i] for i in new_modes2]
    tensor2 = tl.transpose(tensor2, batch_modes2+modes2+new_modes2)
    tensor2 = tl.reshape(tensor2, (*batch_shape, contraction_dim, -1))
    
    res = tl.matmul(tensor1, tensor2)
    res = tl.reshape(res, (*batch_shape, *new_shape1, *new_shape2))

    final_modes += [i for i in range(res.ndim) if i not in final_modes]
    
    if final_modes:
        res = tl.transpose(res, final_modes)

    return res

def batched_tensor_dot_matmul2(tensor1, tensor2, modes, batched_modes=()):
    modes1, modes2 = _validate_contraction_modes(tensor1.shape, tensor2.shape, modes)
    batch_modes1, batch_modes2 = _validate_contraction_modes(tensor1.shape, tensor2.shape, batched_modes, batched_modes=True)

    contraction_shape = [s for (i, s) in enumerate(tl.shape(tensor1)) if i in modes1]
    contraction_dim = prod(contraction_shape)
    batch_shape = [s for (i, s) in enumerate(tl.shape(tensor1)) if i in batch_modes1]
    
    n_free = tl.ndim(tensor1) - len(modes1) - len(batch_modes1)
    # We will reorganize tensor1 by just moving the contraction modes to the end
    modes_begin1 = []
    modes_end1 = []
    shape_begin2 = []
    last_mode_is_batched = False
    last_mode = tl.ndim(tensor1) - 1
    for i, s in enumerate(tl.shape(tensor1)):
        if i in batch_modes1:
            modes_begin1.append(i)
            shape_begin2.append(s)
            last_mode_is_batched = True
        elif i in modes1:
            modes_end1.append(i)
        else:
            modes_begin1.append(i)
            if i != last_mode or n_free:
                shape_begin2.append(1)
            last_mode_is_batched = False
    tensor1 = tl.transpose(tensor1, modes_begin1+modes_end1)

    n_modes_1 = tl.ndim(tensor1) - len(modes1)
    shape = list(tl.shape(tensor1))[:n_modes_1]
    if last_mode_is_batched:
        shape += [1]
    elif shape_begin2:
        shape_begin2.pop(-1)

    tensor1 = tl.reshape(tensor1, (*shape, contraction_dim))
    
    # these are neither batch-size nor contraction modes: put them last
    new_modes2 = [i for i in range(tensor2.ndim) if i not in batch_modes2+modes2]
    new_shape2 = [tl.shape(tensor2)[i] for i in new_modes2]
    tensor2 = tl.transpose(tensor2, batch_modes2+modes2+new_modes2)

    if not new_modes2:
        squeeze_last = True
    else:
        squeeze_last = False
    tensor2 = tl.reshape(tensor2, (*shape_begin2, contraction_dim, -1))
    
    res = torch.matmul(tensor1, tensor2)
    
    out_shape = list(tl.shape(res))
    if squeeze_last:
        out_shape = out_shape[:-1]
    else:
        out_shape = out_shape[:-1] + new_shape2
    if last_mode_is_batched:
        out_shape.pop(n_modes_1)
    
    return tl.reshape(res, out_shape)

def batched_tensor_dot_sum(tensor1, tensor2, modes, batched_modes=()):
    modes1, modes2 = _validate_contraction_modes(tensor1.shape, tensor2.shape, modes)
    batch_modes1, batch_modes2 = _validate_contraction_modes(tensor1.shape, tensor2.shape, batched_modes, batched_modes=True)
    
    # Similarly, we reorganize tensor2 by moving the contraction modes before last, and others at the end
    shape1 = tensor1.shape + (1, )*(tensor2.ndim - len(modes2+batch_modes2))
    
    shape2 = []
    permute_modes2 = []
    batch_counter = 0
    contraction_counter = 0
    for i in range(tensor1.ndim):
        if i in batch_modes1:
            index = batch_modes2[batch_counter]
            permute_modes2.append(index)
            batch_counter += 1
            shape2.append(tensor2.shape[index])
        elif i in modes1:
            index = modes2[contraction_counter]
            permute_modes2.append(index)
            contraction_counter += 1
            shape2.append(tensor2.shape[index])
        else:
            shape2.append(1)
    permute_modes2 += [i for i in range(tl.ndim(tensor2)) if i not in permute_modes2]
    new_shape2 = [s for (i, s) in enumerate(tl.shape(tensor2)) if i not in batch_modes2+modes2]
    shape2 += new_shape2

    tensor1 = tl.reshape(tensor1, shape1)
    tensor2 = tl.transpose(tensor2, permute_modes2)
    tensor2 = tl.reshape(tensor2, shape2)
    res = tensor1 * tensor2

    for mode in sorted(modes1, reverse=True):
        res = tl.sum(res, axis=mode)

    return res