Does CoLA support torch.compile or torch.jit?

wilson-labs / cola

Compositional Linear Algebra

Apache License 2.0

325 stars 24 forks source link

Does CoLA support torch.compile or torch.jit? #83

Open adam-hartshorne opened 5 months ago

adam-hartshorne commented 5 months ago

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

mfinzi commented 5 months ago

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