I’ve been following the development of the higher order ops in PyTorch nightlies for a little bit, and got a chance to try out while_loop. The best examples right now are in the tests, but as another, here’s a mandlebrot example:

https://www.science.org/content/blog-post/eureka

The error message you get actually tells you the fix, but I found it non-intuitive to what I was doing enough I was hesitant to actually just try the config:

I think I did know this at some point, but I was reminded today that unlike the (sometime) black box that is CUDA https://github.com/ROCm/ROCm is actually available on Github, which is operationally much nicer!

Whenever I am trying to do something fun (read: poorly considered) with torch.compile I find myself googling for the same handful of references, most of which are in the PyTorch Drive folder, so putting them here for my own reference.

Every time I have to set up a clean system I manage to mess up cuda somehow, so leaving this as an aide memoir. In general, there is a default nvidia-cuda-toolkit package that ships with Ubuntu-based systems, and you should ignore that. The right options are either:

A very good passing comment from Nikita at work. When folks train in eager mode they are generally able to leverage the full set of aten ops, with their attendant implementations in different backends.

https://github.com/pytorch-labs/torchft

The always prolific Lucidrains has been implementing AlphaFold3 in PyTorch:

https://lilianweng.github.io/posts/2024-02-05-human-data-quality/