add test

thuml · Oct 23, 2023 · 0c9f0a7 · 0c9f0a7
1 parent ef0b64a
commit 0c9f0a7
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/tests/test_pytorch/test_dump_function.py b/tests/test_pytorch/test_dump_function.py
@@ -0,0 +1,15 @@
+import torch
+
+@torch.compile(backend="eager")
+def toy_example(a, b):
+    x = a / (torch.abs(a) + 1)
+    if b.sum() < 0:
+        b = b * -1
+    return x * b
+
+toy_example(torch.randn(10), torch.randn(10))
+
+from depyf.explain import dump_src
+src = dump_src(toy_example)
+with open("./dump_src_dump_function.py", "w") as f:
+    f.write(src)