I would love to see a cleaned-up implementation of that (and I do have notes on the subject but no time to try and implement them). Note that to be useful you

I would love to see a cleaned-up implementation of that (and I do have notes on the subject but no time to try and implement them). Note that to be useful you need backward differentiation (a function with one output but several outputs) and compatibility with tests and loops. See the Zygote.jl library (Julia) for a working example in a different language (their whitepaper is a great starting point if you want to implement something like that).