Назначение массива с помощью Autograd (не снова :())

У меня есть довольно нетривиальная функция, которую я хочу отличить от autograd, но мне недостаточно мастера numpy, чтобы понять, как это сделать без присвоения массива.

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

import autograd.numpy as anp
from autograd import jacobian

def alpha(x):
    return anp.exp(-(x - 10) ** 2) / (x + 1)


def f(x):
    # Matrix getting constructed
    k = anp.zeros((x.shape[0], x.shape[0]))

    # loop over some random 3 dimensional vectors
    for element in anp.random.randint(0, x.shape[0], (x.shape[0], 3)):

        # select 3 values from x
        x_ijk = anp.array([[x[i] for i in element]])

        norm = anp.linalg.norm(
            x_ijk @ anp.vstack((element, element)).transpose()
        )

        # make some matrix from the element
        m = element.reshape(3, 1) @ element.reshape(1, 3)

        # alpha is an arbitrary differentiable function R -> R
        alpha_value = alpha(norm)

        # combine m matricies into k scaling by alpha_value
        n = m.shape[0]
        for i in range(n):
            for j in range(n):
                k[element[i], element[j]] += m[i, j] * alpha_value

    return k @ x


print(jacobian(f)(anp.random.rand(10)))
# And course we get an error
# k[element[i], element[j]] += m[i, j] * alpha_value
# ValueError: setting an array element with a sequence.

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

После написания вышесказанного я тривиально переключился на PyTorch, и код работает нормально. Но я бы все же предпочел использовать autograd

#pytorch version
import torch
from torch.autograd.gradcheck import zero_gradients


def alpha(x):
    return torch.exp(x)


def f(x):
    # Matrix getting constructed
    k = torch.zeros((x.shape[0], x.shape[0]))

    # loop over some random 3 dimensional vectors
    for element in torch.randint(0, x.shape[0], (x.shape[0], 3)):

        # select 3 values from x
        x_ijk = torch.tensor([[1. if n == e else 0 for n in range(len(x))] for e in element]) @ x
        norm = torch.norm(
            x_ijk @ torch.stack((torch.tanh(element.float() + 4), element.float() - 4)).t()
        )

        m = torch.rand(3, 3)

        # alpha is an arbitrary differentiable function R -> R
        alpha_value = alpha(norm)

        n = m.shape[0]
        for i in range(n):
            for j in range(n):
                k[element[i], element[j]] += m[i, j] * alpha_value

    print(k)
    return k @ x


x = torch.rand(4, requires_grad=True)
print(x, '\n')
y = f(x)
print(y, '\n')
grads = []
for val in y:
    val.backward(retain_graph=True)
    grads.append(x.grad.clone())
    zero_gradients(x)
if __name__ == '__main__':
    print(torch.stack(grads))