SplitInfinity:

Yeah, modules cannot be passed around because there is no Module type in the JIT type system. There is no way to annotate a function correctly to make this work.


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. 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.