Does CoLA support torch.compile or torch.jit?

[adam-hartshorne]

CoLA supports JIT with JAX, just wondering if it supports torch.compile or torch.jit ?

[mfinzi]

That's a good question, I believe it works in some cases and not others. We plan to look into this over the next couple months so that we can support it better