Вычисление матричной экспоненты: ValueError: expected A to be a square matrix

Хочу посчитать:

from jax.scipy.linalg import expm
import jax.numpy as jnp
from functools import lru_cache, reduce

num_qubits=2
theta = jnp.asarray(np.pi*np.random.random((15,2,2,2,2,2,2,2,2)))

def pauli_matrix(num_qubits):
    _pauli_matrices = jnp.array(
    [[[1, 0], [0, 1]], [[0, 1], [1, 0]], [[0, -1j], [1j, 0]], [[1, 0], [0, -1]]]
    )
    return reduce(jnp.kron, (_pauli_matrices for _ in range(num_qubits)))[1:]

def SpecialUnitary(num_qubits,theta):
    assert theta.shape[0] == 15
    A = jnp.tensordot(theta, pauli_matrix(num_qubits), axes=[[0], [0]])
    print(f'{A.shape= }{pauli_matrix(num_qubits).shape=}{theta.shape=}')
    return expm(1j*A/2)

SpecialUnitary(num_qubits,theta)

Размеры массивов:

A.shape= (2, 2, 2, 2, 2, 2, 2, 2, 4, 4), pauli_matrix(num_qubits).shape=(15, 4, 4), theta.shape=(15, 2, 2, 2, 2, 2, 2, 2, 2)

Traceback (most recent call last):
  File "/home/keckikon/quantum-neural-network/like/./main.py", line 125, in <module>
    loss_res, opt_state = update(opt_state, x, yy, step_size, pixels, encoding, optimizer = adam(1e-3, b1 = 0.9, b2 = 0.999))
  File "/home/keckikon/quantum-neural-network/like/model.py", line 57, in update
    loss_res, grads = dloss(params,x,y, pixels, encoding)
  File "/home/keckikon/quantum-neural-network/like/model.py", line 51, in dloss
    dloss = value_and_grad(loss, )(params,x,y, pixels, encoding)
  File "/home/keckikon/quantum-neural-network/like/model.py", line 42, in loss
    preds = predict(params, images, pixels, encoding)
  File "/home/keckikon/quantum-neural-network/like/model.py", line 114, in predict
    circ = Сircuit_(params[:4], image, measure=True, qsize=qsize)
  File "/home/keckikon/quantum-neural-network/like/quantum.py", line 77, in Сircuit_
    circ = Qft()[2,3,4,5,6]*Qft()[7,8,9,10,11]*Two_features(theta[::-1])*iQft()[2,3,4,5,6]*iQft()[7,8,9,10,11]*Qft()[2,3,4,5]*Qft()[7,8,9,10]*Two_features(theta2)[0,1,2,3,4,5,7,8,9,10]*iQft()[2,3,4,5]*iQft()[7,8,9,10]
  File "/home/keckikon/quantum-neural-network/like/quantum.py", line 163, in __init__
    self.u = SpecialUnitary(2, self.theta)#(15,2,2,2,..,2)
  File "/home/keckikon/quantum-neural-network/like/quantum.py", line 251, in SpecialUnitary
    return expm(1j*A/2)
  File "/home/keckikon/q_env/lib/python3.9/site-packages/jax/_src/scipy/linalg.py", line 436, in expm
    P, Q, n_squarings = _calc_P_Q(A)
  File "/home/keckikon/q_env/lib/python3.9/site-packages/jax/_src/scipy/linalg.py", line 455, in _calc_P_Q
    raise ValueError('expected A to be a square matrix')
ValueError: expected A to be a square matrix

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