Оптимальное умножение двух 3D-массивов переменной размерности

Я не думаю, что еще что-то нужно сделать.
В зависимости от того, как распределены N_i, например, если половина N_i равна 5, а другая половина 10, может быть полезно иметь два X (один M×5 X для 5 первых столбцов и один (M/2)× 5 X для остальных). То есть, вообще говоря, работать по столбцам, а не по "X"
Или, в некоторых других случаях, заполнить X нулями (например, если N_i равно 10000 для 99% i и 9999 для остальных%).

Но, грубо говоря, numpy предназначен для обработки однородного массива измерений.

Но есть один момент: никогда не используйте hstackstackvstack, append или подобные вещи в цикле for. Это очень медленно. Это подразумевает перераспределение нового массива и копирование всего содержимого предыдущего массива в новый.

Либо используйте простые списки Python для накопления результатов. И делайте укладку только в конце

res=[]
for i in range(M):
    R_i = R[i]
    X_i = X[i]
    res.append(np.dot(R_i, X_i))

np.hstack(res)

Или заранее выделите массив результатов и сохраните результат по мере итерации.

result=np.empty((3,sum(x.shape[1] for x in X)))
k=0
for i in range(M):
    R_i = R[i]
    X_i = X[i]
    result[:,k:k+X_i.shape[1]]=np.dot(R_i,X_i)
    k+=X_i.shape[1]

Я знаю, что я не @mozway и не @hpaulj (это относится к комментарию @chrslg), но действительно кажется, что есть реальное решение с einsum:

np.einsum("ijk,ki->ji",
          np.repeat(R, [x.shape[1] for x in X], axis=0),
          np.hstack(X))

Вот полный код, с которым я тестировал:

import numpy as np
from timeit import Timer

np.random.seed(0)

L, M, timeit_n_times = 3, 5, 10_000
R = [np.random.rand(L, L) for _ in range(M)]
X = []
for i in range(M):
    N_i = np.random.randint(1, 6)
    X_i = np.random.rand(L, N_i)
    X.append(X_i)

def original(R, X):    
    result = np.zeros((L, 0))
    for i in range(M):
        R_i = R[i]
        X_i = X[i]
        result = np.hstack((result, np.dot(R_i, X_i)))
    return result

def list_stack(R, X):
    result = []
    for i in range(M):
        R_i = R[i]
        X_i = X[i]
        result.append(np.dot(R_i, X_i))
    return np.hstack(result)

def preallocate(R, X):
    result=np.empty((L, sum(x.shape[1] for x in X)))
    k=0
    for i in range(M):
        R_i = R[i]
        X_i = X[i]
        result[:, k:k+X_i.shape[1]] = np.dot(R_i, X_i)
        k += X_i.shape[1]
    return result

def with_einsum(R, X):
    return np.einsum("ijk,ki->ji",
                     np.repeat(R, [x.shape[1] for x in X], axis=0),
                     np.hstack(X))
    
assert np.allclose(original(R, X), list_stack(R, X))
assert np.allclose(original(R, X), preallocate(R, X))
assert np.allclose(original(R, X), with_einsum(R, X))

print("original", Timer(lambda: original(R, X)).timeit(timeit_n_times))
print("list_stack", Timer(lambda: list_stack(R, X)).timeit(timeit_n_times))
print("preallocate", Timer(lambda: preallocate(R, X)).timeit(timeit_n_times))
print("with_einsum", Timer(lambda: with_einsum(R, X)).timeit(timeit_n_times))

Некоторые наблюдения:

• Время расчета list_stack() и preallocate() незначительно отличается, но всегда быстрее, чем original(), что также предлагается и подразумевается в ответе @chrslg.
• Будет ли with_einsum() быстрее или медленнее остальных, во многом зависит от формы задачи:
  • Для маленьких, но многих L×L матриц (L==3 в вопросе) with_einsum() выигрывает.
    Вот результаты для L, M, timeit_n_times = 3, 500, 10_000:

    original 19.07941191700229
    list_stack 6.437358160997974
    preallocate 7.638774587001535
    with_einsum 3.6944152869982645
    

  • Для больших, но немногих L×L матриц with_einsum() проигрывает.
    Вот результаты для L, M, timeit_n_times = 100, 5, 100_000:

    original 6.661783802999707
    list_stack 4.67112236200046
    preallocate 4.94899292899936
    with_einsum 28.233751856001618
    

Я не проверял влияние размера и вариации N_i.

Обновлять

Основываясь на обновлении вопроса и мыслях @chrslg, я попробовал, будет ли заполнение нулями жизнеспособным способом. Итог: вероятно, нет (по крайней мере, не в данном примере).

Вот новый код тестирования:

from timeit import Timer
import numpy as np

M, N, timeit_n_times = 30, 100, 10_000

np.random.seed(0)
p = np.random.randint(M, size=N)
R = np.random.rand(M, 3, 3)
X = np.random.rand(3, N)

einsum = lambda: np.einsum('ijk,ki->ji', R[p], X)

def direct():
    res = np.zeros((3, N))
    for i in range(N):
        res[:, i] = np.dot(R[p[i]], X[:, i])
    return res

# Rearrange data for padding
max_count = np.max(np.unique(p, return_counts=True)[1])
counts = np.zeros(M, dtype=int)

X_padded = np.zeros((M, max_count, 3))
for i, idx in enumerate(p):
    X_padded[idx, counts[idx]] = X[:, i]
    counts[idx] += 1

padded = lambda: np.einsum('ijk,ilk->ilj', R, X_padded)
    
result_einsum = einsum()
result_direct = direct()
result_padded_raw = padded()

# Extract padding result (reverse steps of getting from X to X_padded)
result_padded = np.zeros((3, N))
counts = np.zeros(M, dtype=int)
for i, idx in enumerate(p):
    result_padded[:, i] = result_padded_raw[idx, counts[idx]]
    counts[idx] += 1

assert np.allclose(result_einsum, result_direct)
assert np.allclose(result_padded, result_direct)

print("einsum", Timer(einsum).timeit(timeit_n_times))
print("direct", Timer(direct).timeit(timeit_n_times))
print("padded", Timer(padded).timeit(timeit_n_times))

Здесь мы сначала переставляем данные в X_padded, который имеет M×max_count×3-образную форму, где max_count — максимальное количество ссылок на один из индексов в R из p. Мы итеративно заполняем X_padded для каждого индекса, оставляя неиспользуемое пространство заполненным нулями. В конце концов, мы снова можем использовать версию einsum для вычисления результата (padded = lambda: np.einsum('ijk,ilk->ilj', R, X_padded)). Если мы хотим сравнить результат с результатами других методов, нам нужно переставить его обратно, по сути инвертируя шаги перехода от X к X_padded.

Наблюдения:

• С точки зрения памяти мы имеем:

  • Сохранено на стороне R, так как нам больше не нужно копировать отдельные матрицы 3×3.
  • Оплачивается на стороне X, так как нам нужно заполнение нулями.

• По скорости мы, кажется, немного проиграли;
  вот результаты для M, N, timeit_n_times = 30, 100, 100_000:

  einsum 0.7587459350002064
  direct 13.305958173999898
  padded 1.502786388000004
  

  Обратите внимание, что сюда даже не входит время на перестановку данных. Так что, вероятно, отступы здесь не помогут – по крайней мере, не так, как я пытался.