Package: depyf
Have you ever felt overwhelmed by the complexities oftorch.compile? Diving into its workings can feel like black magic, with bytecode and Python internal details that many users fail to understand, hindering them from understanding and adapting totorch.compile.If you also face the problem, then you might be interested indepyf. As the logo suggests,depyfis a software tool to leverage advanced Python features (the Python snake symbol) to open up internal details (the internal gears symbol) of PyTorch's compilertorch.compile(the PyTorch logo), so that users can understand it, adapt to it, and tune their code (the debugger symbol) to get maximum performance benefit out of it.:warning: This project is developed under close collaborations with the PyTorch team. We frequently require new features from PyTorch to support better understanding oftorch.compile. Therefore,please use this project along with PyTorch nightly. Visit thePyTorch websitefor how to install nightly version of PyTorch.We recommend updating your PyTorch nightly installation every week or so.:warning: During development, we seek suggestions from the community quite a lot. You may find some early usage examples from some discussion forums or social media platforms.Please follow the latest documentation for how to use this tool.Whydepyf?If you want to understand bytecode generated bytorch.compile, thendepyfmight be the only choice for you. Below we tested several existing decompilers, they struggle to decompile simple Python bytecode across versions, and have poor support for PyTorch.DecompilerPython 3.8Python 3.9Python 3.10Python 3.11PyTorchdecompyle390% (72/80)××××uncompyle692.5% (74/80)××××pycdc75% (60/80)75% (60/80)75% (60/80)×19.3% (27/140)depyf100% (80/80)100% (80/80)100% (80/80)100% (80/80)100% (140/140)InstallationStable release:pip install depyfNightly version (recommended):pip install git+https://github.com/thuml/depyf.gitUsageThe main usage is quite simple: just wrap your code within a context manager:import torch
from torch import _dynamo as torchdynamo
from typing import List@torch.compiledef toy_example(a, b):x = a / (torch.abs(a) + 1)if b.sum() < 0:b = b * -1return x * b

def main():for _ in range(100):toy_example(torch.randn(10), torch.randn(10))

if __name__ == "__main__":-     main()+     import depyf+     with depyf.prepare_debug("./dump_src_dir"):+         main()Then you can see all the details oftorch.compileinside the directory./dump_src_dir. The details are organized into the following:full_code_for_xxx.pyfor each function usingtorch.compile__transformed_code_for_xxx.pyfor Python code associated with each graph.__compiled_fn_xxx.pyfor each computation graph and its optimization:Captured Graph: a plain forward computation graphJoint Graph: joint forward-backward graph fromAOTAutogradForward Graph: forward graph fromAOTAutogradBackward Graph: backward graph fromAOTAutogradkernel xxx: compiled CPU/GPU kernel wrapper from Inductor.If you want to use debugger to step through the above code, just add another context manager (and launch the script through debuggers):import torch
from torch import _dynamo as torchdynamo
from typing import List@torch.compiledef toy_example(a, b):x = a / (torch.abs(a) + 1)if b.sum() < 0:b = b * -1return x * b

def main():for _ in range(100):toy_example(torch.randn(10), torch.randn(10))

if __name__ == "__main__":import depyfwith depyf.prepare_debug("./dump_src_dir"):main()+     with depyf.debug():+         main()Callingdepyf.debug()will pause the program for you to set breakpoints, and then you can use debuggers to hit breakpoints in these files under the./dump_src_dirdirectory you specified above.ContactIf you have any question aboutdepyf, feel free to open issues to reach out! Any discussion/issue report/PR is welcome. Or contactyoukaichao@gmail.comif you have any other questions.