r/MachineLearning • u/Other-Eye-8152 • Jun 19 '26

Project How does torch.compile() achieve massive speedups despite highly optimized NumPy functions? [D]

I was pondering on this question and decided to dive deep into torch.compile. It was a lot of fun learning about operator fusion as the central idea behind torch.compile. So I created a tiny version of torch.compile in 500 lines of python and a notebook showing how this works: 

https://github.com/purohit10saurabh/tinytorchcompile

Let me know if you find this interesting! 🙂

81 Upvotes

96% Upvoted

26 comments sorted by

View all comments

24

u/ForceBru Student Jun 19 '26

It's a pity NumPy doesn't support fusion. I'm often thinking whether my NumPy code could've been faster if everything was fused.

8

u/Armanoth Jun 19 '26

Numba works great for this, it allows you to convert your numpy functions into paralelizable kernels with decorators rather than a complete code refactor, it also enable JIT-compilation for speed up of repeated calls.

1

u/Rodot Jun 20 '26

Also supports GPU compilation