Своя операция с градиентом в PyTorch

Класс Function из модуля autograd задаёт пару шагов: прямой считает значение, обратный - как передать градиент дальше. Оба шага оформляют статическими методами forward и backward.

Пример умножает вход на два. В обратном шаге градиент тоже умножается на два. Вызов идёт через метод apply:

import torch

class DoubleMul(torch.autograd.Function):
	@staticmethod
	def forward(ctx, input):
		return input * 2

	@staticmethod
	def backward(ctx, grad_output):
		return grad_output * 2

x = torch.tensor(3.0, requires_grad=True)
y = DoubleMul.apply(x)
y.backward()
print(x.grad) # выведет tensor(2.)

Опишите операцию, которая в прямом шаге прибавляет единицу к входу, а в обратном передаёт градиент без изменения. Проверьте на 2.0 с записью. После обратного прохода выведите градиент аргумента.

Сделайте операцию утроения в прямом шаге и утроения градиента в обратном. Для 1.0 с записью получите результат и выведите градиент аргумента после обратного прохода.

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