3 ms·
Yeah, one unspoken theme of this blog post is "look how nice torch.compile" is :) Fun fact, I had to put in extra work to get torch.compile working with my cod
by mrcslws 3y ago
Yeah, one unspoken theme of this blog post is "look how nice torch.compile" is :)
Fun fact, I had to put in extra work to get torch.compile working with my code, for understandable reasons. My library, Vexpr, literally runs an interpreter inside of Python, reading a big tree-like namedtuple-of-namedtuples "expression" data structure and evaluating it recursively. That data structure was way too fancy for torch.compile's guards, so I actually wrote code [1] that converts a Vexpr expression into a big Python code string and evals it, factoring the interpreter out of the code, then I pass that eval'd string into torch.compile.
One torch.compile capability I would be excited to see is compatibility with torch.vmap. One selling point of Vexpr is that you can use vmap with it, so I was sad when I found I couldn't use vmap and still support torch.compile. This made me convert a bunch of my GP kernels [2] to be batch-aware. (This missing capability is also understandable -- both vmap and compile are new.)
Anyway, I'm a fan of what y'all are doing!
[1] https://github.com/outergroup/vexpr/blob/e732e034768443386f9d6731f89cf16c206f2a5b/vexpr/to_python.py#L93-L98 https://github.com/outergroup/vexpr/blob/e732e034768443386f9...
[2] https://github.com/outergroup/outer-loop-cookbook/blob/5d94c431760fd1dfc10e49214859bd3536b15f45/mnist_project/src/gp/vexpr_handson_gp.py#L35-L37 https://github.com/outergroup/outer-loop-cookbook/blob/5d94c...
- voz_ 3y agoI spend a lot of sweat in the guards - I am very interested in how it failed! Can you say more? Did guard creation fail? or did guard check_fn perf overhead destroy it? > One torch.compile capability I would be excited to see is compatibility with torch.vmap We added support for torch.func.vmap, iirc - check out test_higher_order_ops.py, grep for vmap.