Как умножить матрицу 2x3x3x3 на матрицу 2x3, чтобы получить матрицу 2x3 ⇐ Python

[Anonymous]

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

Код: Выделить всё

import torch

x_1 = torch.tensor([[1.,1.,1.]], requires_grad=True)
x_2 = torch.tensor([[2.,2.,2.]], requires_grad=True)
# Input to the network with dim 2x3 --> 2 Samples 3 Feature
x = torch.cat((x_1,x_2),dim=0)

def calculation(x):
c = torch.tensor([[1,2,3],[4,5,6],[7,8,9]]).float()
return (x@c)**2

c = torch.tensor([[1,2,3],[4,5,6],[7,8,9]]).float()

#output of the network with dimension 2x3 --> 3 outputs per Sample
y = calculation(x)

#Calculation of my jacobian with dimension 2x3x3 (one for each sample)
g = torch.autograd.functional.jacobian(calculation,(x),create_graph=True)
jacobian_summarized = torch.sum(g, dim=0).transpose(0,1)
#Calculation of my second order derivative 2x3x3x3 (On Tensor for each Sample)
gg = torch.autograd.functional.jacobian(lambda x: torch.sum(torch.autograd.functional.jacobian(calculation, (x), create_graph=True), dim=0).transpose(0,1),
(x),
create_graph=True)
second_order_derivative = torch.sum(gg, dim=0).transpose(1,2).transpose(0,1)

print('x:', x)
print('c:',c)
print('y:',y)
print('First Order Derivative:',jacobian_summarized)
print('Second Order Derivative:',second_order_derivative)

# Multiplication with for loop
result = torch.empty(0)
for i in range(y.shape[0]):
result_row = torch.empty(0)
for ii in range(y.shape[1]):
result_value = (x[i].unsqueeze(0))@second_order_derivative[i][ii]@(x[i].unsqueeze(0).T)
result_row = torch.cat((result_row, result_value), dim=1)
result = torch.cat((result, result_row))
print(result)

Я хотел бы знать, есть ли способ получить такой же результат умножения без необходимости использовать 2 для петли, а скорее простое умножение матриц

Подробнее здесь: https://stackoverflow.com/questions/793 ... 2x3-matrix