I would like to know if it can be combined with the torch for training?

[WindCanDie]

I would like to know if it can be combined with the torch for training?

[2niuhe]

I have the same question. It seems that this project primarily focuses on translating the described computation graph into code, with all the examples provided illustrating the forward inference process. For the backward, one would need to describe the corresponding backward graph's computation process on their own. This is just my understanding, so it may not be correct.

[daneren]

I also want to know this question.

[Edenzzzz]

Even if you use Triton (not torch.compile), you will need to write your own backward. Torch.compile doesn't work for complex operators with transposes etc. such as Flash Attn