pytorch
c2182632 - [LTC] Mark Step Indicator (#76840)

Commit
2 years ago
[LTC] Mark Step Indicator (#76840) Proposed solution for #76826 Basically adds a context which is only "active" when called in mark step. Any backend can then use this to check if within a mark step context. I've also added an example warning in the TS backend so that we now see the following: ```python >>> import torch >>> import torch._lazy >>> import torch._lazy.ts_backend >>> torch._lazy.ts_backend.init() >>> a = torch.tensor([1, 2, 3, 4], device="lazy") >>> b = torch.tensor([5, 6, 7, 8], device="lazy") >>> c = a + b >>> c [W ts_backend_impl.cpp:187] Compile outside of mark step tensor([ 6, 8, 10, 12], device='lazy:0') >>> d = a * b >>> torch._lazy.mark_step() >>> d tensor([ 5, 12, 21, 32], device='lazy:0') ``` Though it was mainly for example and will be happy to remove if this warning is not desired. Fixes #76826 CC: @wconstab @desertfire @henrytwo @ke1337 Pull Request resolved: https://github.com/pytorch/pytorch/pull/76840 Approved by: https://github.com/desertfire
Author
Committer
Parents
Loading