Hey, I’m interested in creating and exporting a pullback model using ONNX. The pullback model wraps an existing model and, for a given input, computes the wrapped model’s output and its gradient w.r.t. to the input. (I’m using a virtual environment with Python 3.8.10 and PyTorch 1.10.1.)
The following code achieves what I want within PyTorch:
import torch
import torch.nn as nn


# dummy model we want to wrap
class BaseModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=3, padding=1)
    
    def forward(self, x):
        return self.conv1(x)


# the pullback model we want to export using ONNX
class PullbackModel(nn.Module):
    def __init__(self, forward_model):
        super().__init__()
        self.forward_model = forward_model
    
    def forward(self, x, dy):
        y = self.forward_model(x)
        y.backward(gradient=dy)
        dx = x.grad
        return y, dx

While within PyTorch, I can now compute exactly what I want (ignoring the possibly accumulating gradients for now), i.e.:
model = BaseModel()
pb_model = PullbackModel(model)

x = torch.randn(1, 2, 5, 5)  # example input
dy = torch.ones(1, 1, 5, 5)  # example output sensitivity 

# compute the wrapped model's output and the input sensitivity
# (i.e., the gradient w.r.t. the input and the given output sensitivity)
y, dx = pb_model(x, dy)

This works as expected but for my use-case I need to export the pullback model using ONNX for later inference. However, when I try
torch.onnx.export(pb_model, (x, dy), "pb-model.onnx", input_names=["x", "dy"], output_names=["y", "dx"])

I get the following error:
Traceback (most recent call last):
  File "pullback_onnx.py", line 117, in <module>
    torch.onnx.export(pb_model, (x, dy), "pb-model.onnx", input_names=["x", "dy"], output_names=["y", "dx"])
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/__init__.py", line 316, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 107, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 724, in _export
    _model_to_graph(model, args, verbose, input_names,
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 493, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 437, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 388, in _trace_and_get_graph_from_model
    torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
  File "~/my-venv/lib/python3.8/site-packages/torch/jit/_trace.py", line 1166, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "~/my-venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "~/my-venv/lib/python3.8/site-packages/torch/jit/_trace.py", line 127, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "~/my-venv/lib/python3.8/site-packages/torch/jit/_trace.py", line 118, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "~/my-venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "~/my-venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1090, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "pullback_onnx.py", line 94, in forward
    y.backward(gradient=dy)
  File "~/my-venv/lib/python3.8/site-packages/torch/_tensor.py", line 307, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "~/my-venv/lib/python3.8/site-packages/torch/autograd/__init__.py", line 154, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Found an unsupported argument type in the JIT tracer. File a bug report.

As there seemed to be a problem with the trace during torch.onnx.export, I tried converting the pullback model to a ScriptModule before exporting to ONNX.
For that, I directly traced the wrapped model during the pullback model’s __init__ using
    def __init__(self, forward_model):
        super().__init__()
        self.forward_model = torch.jit.trace(forward_model, torch.rand(1, 2, 5, 5))

and converted pb_model to a ScriptModule using
scripted_module = torch.jit.script(pb_model, (x, dy))

before trying to export the model again:
torch.onnx.export(scripted_module, (x, dy), "pb-model.onnx", input_names=["x", "dy"], output_names=["y", "dx"])

Now, I get a different error:
  File "pullback_onnx.py", line 118, in <module>
    torch.onnx.export(scripted_module, (x, dy), "pb-model.onnx", input_names=["x", "dy"], output_names=["y", "dx"])
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/__init__.py", line 316, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 107, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 724, in _export
    _model_to_graph(model, args, verbose, input_names,
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 497, in _model_to_graph
    graph = _optimize_graph(graph, operator_export_type,
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 216, in _optimize_graph
    graph = torch._C._jit_pass_onnx(graph, operator_export_type)
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/__init__.py", line 373, in _run_symbolic_function
    return utils._run_symbolic_function(*args, **kwargs)
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 1137, in _run_symbolic_function
    symbolic_fn = _find_symbolic_in_registry(domain, symbolic_name, opset_version,
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/utils.py", line 982, in _find_symbolic_in_registry
    return sym_registry.get_registered_op(op_name, domain, opset_version)
  File "~/my-venv/lib/python3.8/site-packages/torch/onnx/symbolic_registry.py", line 125, in get_registered_op
    raise RuntimeError(msg)
RuntimeError: Exporting the operator prim_grad to ONNX opset version 9 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.

I tried exporting using different opset versions which didn’t change the outcome. Looking up the “prim_grad” operator also didn’t yield any helpful results.
Is exporting such a model with ONNX even possible? If yes, what changes would be needed?
Thank you in advance!
PS: As this is my first topic, I’m not sure which category this question should belong to and if I provided all necessary information. Please let me know if there is something missing!