Pytorch的“`autograd`”如何处理非数学函数？

import torch

class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # 定义非数学函数的正向传播逻辑
        output = input * 2
        ctx.save_for_backward(input)  # 保存输入张量，用于反向传播计算梯度
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # 定义反向传播逻辑以计算梯度
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()  # 计算输入张量的梯度
        grad_input *= 2  # 非数学函数的导数为2
        return grad_input

# 创建输入张量
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# 使用自定义函数进行计算
output = MyFunction.apply(x)

# 计算梯度
output.backward(torch.ones_like(x))

# 打印输入张量的梯度
print(x.grad)