4 ms·
Pretty cool to see people using compile in the wild :)
by voz_ 3y ago
Pretty cool to see people using compile in the wild :)
- mrcslws 3y agoYeah, 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.