Enamex:

I was hoping there’d be a way to work around this specifically  The JIT already works with a self argument of an nn.Module subtype. How can we make this available for free functions?  This restriction greatly limits the composability of free functions if I want to use script .


The short answer is not really. The JIT supports classes insofar as they are static - if you JIT-compile Modules, you don’t give it the class source but rather an instance. It will then go through the data members and see and process their types, assuming they will be fixed (which isn’t the case in Python in general).
Maybe you might find more success in taking the “JIT” part more literally, e.g. JIT-compiling a local function from a dynamic function.


Honestly not sure at that point if there’d be a worthwhile speed either. This’s unrelated, but I’ve had some trouble understanding what the JIT does with what it sees, and intuiting where my code might be giving it more trouble than it can handle.


I’m affraid the JIT isn’t a magic bullet for optimization but “only” does specific things that are highly desirable.