Градиент через решение системы в PyTorch

Линейная система может стоять внутри обучаемого графа. Функция solve из модуля linalg находит вектор x из равенства A @ x = b. Если матрица или правая часть участвуют в графе, обратный проход доходит и до них.

Зададим симметричную матрицу и правую часть с записью градиента, решим систему, сложим компоненты ответа и выполним обратный проход:

import torch

A = torch.tensor([[2.0, 1.0], [1.0, 3.0]],
	requires_grad=True)
b = torch.tensor([1.0, 2.0], requires_grad=True)
x = torch.linalg.solve(A, b)
loss = x.sum()
loss.backward()
print(A.grad)
print(b.grad)

Для матрицы [[4.0, 0.0], [0.0, 2.0]] и правой части [2.0, 6.0] с записью градиента найдите решение, возьмите сумму компонент и выполните обратный проход. Выведите градиент правой части.

Решите систему с матрицей [[1.0, 0.0], [0.0, 1.0]] и b = [3.0, 5.0], где записывается градиент только у b. Суммируйте ответ и выведите градиент правой части после обратного прохода.

Для A = [[2.0, 0.0], [0.0, 2.0]] и b = [1.0, 1.0] с записью на обоих тензорах получите решение, умножьте первую компоненту на вторую и найдите обратный проход от этого произведения. Выведите градиент матрицы.