Models with partial diff in forward pass

[ipcamit]

I want to export a model which can do partial diff against input variables during forward pass, using autograd. I am bit confused if it can be exported to an onnx object or not. The ONNX repository operators readme hints it might be possible (https://github.com/onnx/onnx/blob/master/docs/Operators.md#ai.onnx.preview.training.Gradient). But i get next to no search results on googling "ai.onnx.preview.training.Gradient".
Some discussions suggest it might not be possible (pytorch/pytorch#15644) with reason quoted as

ONNX export relies on TorchScript, and I believe there's currently no way to capture the backwards graph as TorchScript, so this is not feasible.

I do not understand that, TorchScript does support gradients, does it not (I can easily compile my modules with calls to autograd.grad using torch.jit.script())?

Discussion here is bit more ambiguous for me microsoft/onnxruntime#5892

Can anyone tell if it is possible to export models in onnx with grad calls in forward function?

[malteal]

@ipcamit Did you find a solution? I have the same problem.

[ipcamit]

No. There were several more issues. Like certain functions were not supported for double precision. So eventually we abandoned onnx, and went full TorchScript.

[ipcamit]

Just out of curiosity why do you need onnx ? Why not torch script directly?

[malteal]

No.. :( Sadly the framework where I have to implement the model requires it to be an ONNX model...

[ipcamit]

I don't think it will be possible in near future as it will mean onnx having a full fledged autodiff library. One workaround that we did briefly consider was having backend in TensorFlow. As for simple enough models onnx->TF does work (https://github.com/onnx/onnx-tensorflow). And move grad operation out of model. But even that was severely limiting and certain functions were incompatible (i think ceil was an issue). So even if there is a workaround, I will suspect that it will be bit dicey and not general enough.
Lastly if you have will, and enough manpower, one of the idea we were considering was writing our own onnx backend. It is not trivial, but not too hard either. ONNX object provides all the functions that model is calling, and there are finite number of them. So you can use do one time port of common ones like Linear layers, non-linear functions etc using C++. And grad can be provided by autodiff library Enzyme (https://enzyme.mit.edu/). But it was way too ambitious and further away from our goal, so it never went past the crazy idea stage. Plus now there is issue of Enzyme limitations (it is impossibly slow on vector<vector> kind of C++ objects, and could not differentiate Eigen3 C++ library's eigenvalue finder).
If it is an opensource or academic project, and you want to explore the enzyme side, I am willing to contribute as well.