Custom functions

[guillaume-be]

Hello,

I am interested in implementing a custom function using this library, but I am not quite sure how to get started. I would like to have a custom implementation defining both forward and backward pass. I believe a related issue (~1 year old) was raised by a user attempting to do the same, and I was wondering if it is now possible: https://github.com/LaurentMazare/tch-rs/issues/70

Here is an example of what I am trying to accomplish, both in Libtorch and in Pytorch:

Libtorch:

https://github.com/pytorch/pytorch/issues/35736

Pytorch

https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html

Are the bindings to do so now available? I did not find reference to a Function trait in the current version of the library. Thank you!

[LaurentMazare]

Hello, I don't think we have any support for this at the moment. That being said, it's probably not very hard to support simple use cases such as the MulConstant case in the issue you referred to. Could you give a bit more color about the custom function that you would like to be able to implement to see how hard this would be to handle?

For unary functions, we could imagine passing two Rust defined closures to the C++ api for forward and backward and have a specific class defined in torch_api.cpp that inherits from Function and would be responsible for applying these closures. These closures will need access to the AutogradContext variable so we would also need to sketch an api for this.

[guillaume-be]

Hello,

Apologies for the delayed feedback. I am trying to implement a model using custom backward pass operations to reduce memory footprint during training (see https://ai.googleblog.com/2020/01/reformer-efficient-transformer.html and the reversible layers). The operations are more complex than MultConstant, see for example this entire customized module: https://github.com/huggingface/transformers/blob/24184e73c441397edd51e9068e0f49c0418d25ab/src/transformers/modeling_reformer.py#L1574

This is not a must have at this point, as I am more focused on inference, but would be a nice to have to offer the same training capabilities as in Pytorch and Libtorch. I believe these reversible layers are also used in several state of the art computer vision models.