The JIT optimization steps probably are among the most sophisticated bits in PyTorch (along with the dispatcher,…). For a deep dive on one of the fusers, I can also enthusiastically recommend Christian Sarofeen’s talk 6 (I think you need to register to see it).
Best regards
Thomas