I am currently working on a module for a research paper where I need to perform a few calculations, like solving quadratic equations. Because the module is used a lot throughout the network and my memory is under pressure, I want to perform the calculations inplace. But the code is turning out to be quite ugly and borderline incomprehensible. Is there any way to improve the readability so that it’s not a massive chain of .mul_, .add_ and parentheses everywhere?
Also the autograd-errors are quite hard to debug if something is wrong in the inplace-calcuation.
Can something like the JIT-compiler help in this case? If yes, I can’t jit only part of the code, can I? I imagine the analysis for the compiler to be quite easy and well-understood.