More generally, I’m looking for a way to treat a function body as a “template” that I can provide specific types for, and ask the JIT explicitly to give me a new version for it. So it doesn’t just see its full name and go “I’ve already seen this; here’s the cached version”.

I don’t think it is a publicly advertised feature, but you could try using @torch.jit._overload. You can find examples of how to use this decorator in test/test_jit.py, like test_function_overloads. The caveat here is that the one body you supply will be used with all overloaded signatures because there is no concept of TypeVar in the JIT typing system. So you might have do type refinement to write one body that works for all types (see my_conv in test_function_overloading_isinstance, also in test_jit.py).

By your confirmation though, this’s probably very difficult. I couldn’t get the JIT to accept a function that takes in a  nn.Module  instance as an argument either (besides  self.  accesses) so there’re multiple issues here.

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.