Говорят, что невозможно возненавидеть кого-то, если сначала не полюбил его. Не знаю, справедливо ли это в целом, но это определённо описывает моё отношение к NumPy.

NumPy — это ПО для выполнения вычислений с массивами на Python. Оно невероятно популярно и очень сильно повлияло на все популярные библиотеки машинного обучения, например, на PyTorch. Эти библиотеки во многом имеют те же самые проблемы, но для конкретики я рассмотрю NumPy.

NumPy упрощает выполнение простых задач. Пусть A — это матрица 5×5, x — вектор длиной 5, и вам нужно найти такой вектор y, что Ay=x. В NumPy это будет выглядеть так:

y = np.linalg.solve(A, x)

Всё так изящно и чётко!

Но, допустим, ситуация стала немножечко сложнее. Пусть A — это стек из 100 матриц 5×5 , заданный как массив 100×5×5. И пусть x — это стек 100 векторов длиной 5, заданный как массив 100×5. Допустим, нам нужно вычислить Aᵢyᵢ=xᵢ для 1≤i≤100.

Если бы мы могли использовать циклы, то это было бы просто:

y = np.empty_like(x)
for i in range(100):
    y[i,:] = np.linalg.solve(A[i,:,:], x[i,:])

Но мы не можем использовать циклы. В какой-то степени это ограничение связано с тем, что Python циклы медленны. Но сегодня всё вычисляется на GPU и у нас есть большие массивы, поэтому, вероятно, вы не захотите пользоваться массивами ни на одном из языков. Чтобы все транзисторы трудились в поте лица, нам нужно вызывать особые функции GPU, которые как бы разобьют массивы на кучу маленьких частей и обработают их параллельно.

Хорошо здесь то, что NumPy знает об этих специальных процедурах (по крайней мере, если вы используете JAX или CuPy), и если корректно вызвать np.linalg.solve, то она их применит.

Плохо здесь то, что никто не знает, как это сделать.

Не верите мне? Тогда скажите, что из этого подходит для решения?

y = linalg.solve(A,x)
y = linalg.solve(A,x,axis=0)
y = linalg.solve(A,x,axes=[[1,2],1])
y = linalg.solve(A.T, x.T)
y = linalg.solve(A.T, x).T
y = linalg.solve(A, x[None,:,:])
y = linalg.solve(A,x[:,:,None])
y = linalg.solve(A,x[:,:,None])[:,:,0]
y = linalg.solve(A[:,:,:,None],x[:,None,None,:])
y = linalg.solve(A.transpose([1,2,0]),x[:,:,None]).T

Никто не знает. И позвольте мне показать вам кое-что ещё. Вот документация:

Прочитайте и помедитируйте над прочитанным. А теперь осознайте: вы всё ещё не знаете, как вычислить Aᵢyᵢ=xᵢ для всех i за раз. Возможно ли это вообще? Может, я соврал, когда утверждал это?

Насколько я понимаю, люди пробуют произвольные вариации, пока не подберут то, что, на их взгляд, работает.

Чем плоха NumPy

Задача NumPy — применять операции к массивам. Если у массива две или меньше размерностей, то всё в порядке. Но если вы делаете что-то хотя бы чуть более сложное, то вам неизбежно понадобится операция, которую нужно применить к некоторым размерностям массива A, к каким-то другим размерностям B и ещё одним размерностям массива C. А у NumPy нет теоретического обоснования того, как это выразить.

Позвольте продемонстрировать, что я имею в виду. Предположим, что:

• A — это массив K×L×M

• B — это массив L×N

• C — это массив K×M

Допустим, для каждого k и n нам нужно вычислить среднее по размерностям L и M. То есть нам нужно следующее:

   Dkn = 1/(LM) × ∑lm Aklm Bln Ckm.

Чтобы сделать это, есть два варианта. Первый — использовать гротескный трюк выравнивания размерностей:

D = np.mean(
        np.mean(
            A[:,:,:,None] *
            B[None,:,None,:] *
            C[:,None,:,None],
        axis=1),
    axis=1)

«Какого чёрта?», — спросите вы? Почему повсюду указаны None? При индексировании массива в NumPy можно написать None для вставки новой размерности. A имеет форму K×L×M, но A[:,:,:,None] имеет форму K×L×M×1. Аналогично, B[None,:,None,:] —это 1×L×1×N, а C[:,None,:,None] — это K×1×M×1. При их перемножении NumPy «транслирует» все размерности размера 1 для получения массива K×L×M×N. Затем вызовы np.mean выполняют усреднение среди размерностей L и M.

Я считаю, что это плохо. Я годами пользовался NumPy, но всё равно было невозможно писать подобный код без постоянного совершения ошибок.

К тому же, это почти невозможно читать. Чтобы доказать это, я подбросил монетку и добавил баг, если выпала решка. Есть ли там баг? А вы уверены? Никто точно не знает.

Второй вариант — отчаянно пытаться быть умным. Жизнь коротка и драгоценна, но если вы потратите приличную её часть на чтение документации NumPy, то рано или поздно сможете найти функцию np.tensordot; возможно, у вас получится заставить её выполнять основную часть работы:

D = (1/L) * np.mean(
                np.tensordot(A, B, axes=[1,0]) *
                C[:,:,None],
            axis=1)

Этот код корректен. (Гарантирую.) Но почему он работает? Что именно делает np.tensordot? У вас появилось бы хоть какое-то представление о происходящем в коде, если бы вы встретили его в другом контексте?

Вот, как это бы сделал я, если бы можно было использовать циклы:

D = np.zeros((K,N))  
for k in range(K):  
    for n in range(N):  
        a = A[k,:,:]  
        b = B[:,n]  
        c = C[k,:]  
        assert a.shape == (L,M)  
        assert b.shape == (L,)  
        assert c.shape == (M,)  
        D[k,n] = np.mean(a * b[:,None] * c[None,:])

Людям, много писавшим на NumPy, этот код может показаться неуклюжим. Подозреваю, что это частично вызвано стокгольмским синдромом. Но мы с уверенностью можем сказать, что он понятен.

На практике, всё часто бывает ещё хуже. Допустим, A имеет форму M×K×L , а не K×L×M. Если работать с циклами, то особой проблемы не возникает. Но NumPy вынуждает писать чудовищные конструкции типа A.transpose([1,2,0]). Или это должно быть A.transpose([2,0,1])? Какие формы они создают? Этого не знает никто.

Циклы были бы лучше.

Ладно, признаю, я соврал

Есть ещё и третий вариант:

D = 1/(L*M) * np.einsum('klm,ln,km->kn', A, B, C)

Если вы никогда не видели соглашения Эйнштейна, то это может показаться пугающим. Но помните, что наша цель — найти

   D_kn = 1/(LM) × ∑_lm A_klm B_ln C_km.

Строка в приведённом выше коде, по сути, создаёт метки для индексов каждых из входящих данных (klm,ln,km) и целевой индекс для вывода (->kn). Затем np.einsum перемножает соответствующие элементы входящих данных и выполняет суммирование по всем индексам, не находящимся в выводе.

Лично я считаю, что np.einsum — это один из немногих аспектов NumPy, которые действительно хороши. Возня со строками немного утомляет, но оно того стоит, потому что в целом функцию становится проще понять, она полностью явная, достаточно обобщённая и мощная.

Остаётся только вопрос: как np.einsum реализует всё это? Она использует индексы. Или, если точнее, она добавляет крошечный язык предметной области, основанный на индексах. Она не страдает от изъянов архитектуры NumPy, потому что отказывается играть по обычным правилам NumPy.

Но возможности np.einsum ограничены. (Einops способна на чуть большее.) Что, если нам нужно применить какую-то другую функцию для разных размерностей каких-то массивов?Такой функции, как np.linalg.einsolve, не существует. А если мы создадим собственную функцию, то для неё точно не будет «эйнштейновской» версии.

Мне кажется, удобство np.einsum демонстрирует, что NumPy зашла куда-то не туда.

Интерлюдия

Вот картина, кажущаяся аналогом обсуждаемой нами темы.

Где NumPy зашла не туда?

Вот, чего мне хочется от языка для работы с массивами. Я говорю не конкретно о синтаксисе, но было бы здорово, если бы:

1. Когда я хочу сделать что-то, то существовал бы «очевидный» способ реализации.

2. Когда я читаю какой-то код, чтобы было «очевидно», что он делает.

Разве не было бы это здорово? Думаю, NumPy не решает этих задач из-за своего первородного греха: она забрала у нас индексы и заменила их на трансляцию. А трансляция не умещается в возможности индексов.

Я не люблю трансляцию NumPy

Базовый инструмент NumPy — трансляция. Возьмём такой пример кода:

A = np.array([[1,2],[3,4],[5,6]])
B = np.array([10,20])
C = A * B
print(C)

Его результат:

[[ 10  40]
 [ 30  80]
 [ 50 120]]

Здесь A — массив 3×2, B  массив длиной 2. При их перемножении B «транслируется» в форму A, то есть первый столбец A умножается на B[0]=10, а второй умножается на B[1]=20.

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

Другая причина — неявность и нечитаемость. Иногда A*B выполняет умножение поэлементно, а иногда делает что-то более сложное. Поэтому каждый раз, когда мы видим A*B, нам приходится разбираться, какой конкретно применяется случай трансляции.

Но главная проблема трансляции заключается в том, что она инфицирует всё остальное. Ниже я объясню, что имеется в виду.

Мне не нравится индексирование NumPy

Вот вам головоломка. Посмотрите на код:

A = np.ones((10,20,30,40))
i = np.array([1,2,3])
j = np.array([[0],[1]])
B = A[:,i,j,:]

Какую форму имеет B?

Оказывается, правильный ответ — это 10×2×3×40. Индексы i и j транслируются в форму 2×3, а потом... хм... что-то происходит. Попробуйте убедить себя, что это имеет какой-то смысл.

Справились? Так, а теперь попробуйте эту задачку:

C = A[:,:,i,j]
D = A[:,i,:,j]
E = A[:,1:4,j,:]

Какую форму они имеют?

• C имеет форму 10×20×2×3. Это кажется логичным, учитывая произошедшее с B в примере выше.

• А как насчёт D? Он имеет форму 2×3×10×30. Откуда взялись 2 и 3 в начале?

• А что вы думаете про E? «Срезы» в Python исключают конечную точку, так что 1:4 эквивалентно [1,2,3], что эквивалентно i, поэтому E эквивалентен B. Ха-ха, шучу! E имеет форму 10×3×2×1×40.

Да, именно так и происходит. Проверьте сами, если не верите мне! Я понимаю, почему NumPy ведёт себя так, ведь я переварил документ из пяти тысяч слов, в котором объясняются принципы индексирования NumPy. Но мне бы хотелось вернуть это время себе.

(Цепочка рассуждений DeepSeek 76 раз использовала слово «подождите». Она решила всё правильно с первого раза, но когда я попробовал снова, она почему-то ошиблась в B, C, и D, а E решила правильно.)

Это безумие. Для работы с базовыми функциями мы не должны решать дурацкие логические задачки.

Вы можете подумать: «Так, можно ограничиться только простым индексированием». Звучит неплохо, только иногда нам необходимо расширенное индексирование. И даже вы делаете что-то простое, всё равно нужно быть внимательным, чтобы избегать дурацких случаев.

Из-за этого всё снова становится нечитаемым. Даже если вы читаете код, в котором используется простое индексирование, то как узнать, что оно простое? Если вы видите A[B,C], то это может делать почти что угодно. Чтобы понять это, вам нужно помнить формы A, B и C, а также проработать все случаи. И, разумеется, A, B и C часто создаются другим кодом, о котором вам тоже нужно думать…

Я не люблю функции NumPy

Почему в NumPy появилась такая запутанная функция np.linalg.solve(A,B)? Наверно, разработчики сначала сделали так, чтобы она работала, когда A — это 2D-массив, а b — 1D- или 2D-массив, как математическая запись A⁻¹b или A⁻¹B.

Пока всё нормально. Но потом кто-то, вероятно решил добавить 3D-массив. Если бы допускалось использование циклов, то можно было просто использовать старую функцию с циклами. Но циклы использовать нельзя. Значит, остаётся три варианта:

1. Можно было добавить дополнительные аргументы axes, чтобы пользователь мог указать, с какими размерностями работать. Допустим, мы могли бы записать solve(A,B,axes=[[1,2],1]).

2. Можно было добавить разные функции с разными названиями для разных ситуаций. Допустим, solve_matrix_vector решает одну задачу, а solve_tensor_matrix — другую.

3. Можно было добавить соглашение: некий произвольный вариант, по которому solve внутри выравнивает размерности. Тогда уже проблема пользователя, как он будет разбираться и соответствовать этому соглашению.

Все эти варианты плохи, ведь ни один из них не учитывает тот факт, что различных случаев существует комбинаторное количество. Разработчики NumPy выбрали... все три варианта. У некоторых функций есть аргументы axes. У некоторых есть разные версии с разными названиями. У некоторых есть соглашения. У некоторых есть соглашения И аргументы axes. А у некоторых вообще нет векторизованной версии.

Но самый серьёзный изъян NumPy таков: допустим, вы создаёте функцию, решающую задачу с массивами какой-то заданной формы. Как теперь применить эти конкретные размерности к массивам бóльших размеров? Ответ таков: переписать функцию с нуля гораздо более сложным образом. Базовый принцип программирования — это абстракция, решение простых задач и использование этих решений в качестве строительных блоков более сложных задач. NumPy не позволяет нам этого делать.

Прошу вашего внимания

Приведу последний пример, чтобы показать, что я имею в виду. Каждый раз, когда я начинаю ныть по поводу NumPy, люди всегда хотят увидеть пример с самовниманием — основным трюком, лежащим в основе современных языковых моделей. Ну ладно. Вот реализация, которую я скромно считаю лучшей по сравнению со всеми 227 версиями, которые я нашёл, загуглив «self-attention numpy»:

# реализация самовнимания вашего покорного слуги

input_dim = 4  
seq_len = 4  
d_k = 5  
d_v = input_dim  
  
X = np.random.randn(seq_len, input_dim)  
W_q = np.random.randn(input_dim, d_k)  
W_k = np.random.randn(input_dim, d_k)  
W_v = np.random.randn(input_dim, d_v)  
  
def softmax(x, axis):  
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))  
    return e_x / np.sum(e_x, axis=axis, keepdims=True)  
  
def attention(X, W_q, W_k, W_v):  
    d_k = W_k.shape[1]  
    Q = X @ W_q  
    K = X @ W_k  
    V = X @ W_v  
    scores = Q @ K.T / np.sqrt(d_k)  
    attention_weights = softmax(scores, axis=-1)  
    return attention_weights @ V  

Неплохо. Кое-что с axis кажется не очень понятным, ну да ладно.

Но языковым моделям очень важно многоголовое внимание, при котором параллельно многократно реализуется внимание, а затем результаты объединяются. Как нам его реализовать?

Для начала давайте представим, что живём в нормальном мире, где разрешается использовать абстракции. Тогда можно просто вызвать представленную выше функцию в цикле:

# реализация многоголового самовнимания вашего покорного слуги,
# если бы мы могли использовать циклы

n_head = 2
  
X = np.random.randn(seq_len, input_dim)  
W_q = np.random.randn(n_head, input_dim, d_k)  
W_k = np.random.randn(n_head, input_dim, d_k)  
W_v = np.random.randn(n_head, input_dim, d_v)  
W_o = np.random.randn(n_head, d_v, input_dim // n_head)

def multi_head_attention(X, W_q, W_k, W_v, W_o):  
    projected = []  
    for n in range(n_head):  
        output = attention(X,
                           W_q[n,:,:],
                           W_k[n,:,:],
                           W_v[n,:,:])  
        my_proj = output @ W_o[n,:,:]  
        projected.append(my_proj)  
    projected = np.array(projected)  
  
    output = []  
    for i in range(seq_len):  
        my_output = np.ravel(projected[:,i,:])  
        output.append(my_output)  
    return np.array(output)  

Выглядит глупо, правда? Да, и большое спасибо! Умничать плохо.

Мы живём в ненормальном мире. Поэтому приходится делать так:

# реализация самовнимания вашего покорного слуги
# с векторизацией и полной запутанностью

def multi_head_attention(X, W_q, W_k, W_v, W_o):  
    d_k = W_k.shape[-1]  
    Q = np.einsum('si,hij->hsj', X, W_q)  
    K = np.einsum('si,hik->hsk', X, W_k)  
    V = np.einsum('si,hiv->hsv', X, W_v)  
    scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k)  
    weights = softmax(scores, axis=-1)
    output = weights @ V  
    projected = np.einsum('hsv,hvd->hsd', output, W_o)  
    return projected.transpose(1, 0, 2).reshape(
        seq_len, input_dim)  

Ха-ха-ха-ха!

Что же делать?

Я просто высказался, что NumPy — это худший язык для работы с массивами, кроме всех других языков работы с массивами. Но какой смысл жаловаться, если я не могу предложить ничего лучшего?

Ну, на самом деле, у меня есть, что предложить. Я создал прототип «улучшенной» NumPy, сохранившей всю её мощь, но избавившейся от острых углов. Мне думалось, что статья будет коротким мотивирующим введением, но когда я начал писать, меня обуяла злоба и очнулся я, только написав три тысячи слов.

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