xref: /aosp_15_r20/external/pytorch/torch/_export/db/gen_example.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1import os
2import sys
3
4import torch._export.db.examples as examples
5
6TEMPLATE = '''import torch
7
8def {case_name}(x):
9    """
10    """
11
12    return
13'''
14
15if __name__ == "__main__":
16    assert len(sys.argv) == 2
17    root_dir = examples.__name__.replace(".", "/")
18    assert os.path.exists(root_dir)
19    with open(os.path.join(root_dir, sys.argv[1] + ".py"), "w") as f:
20        print("Writing to", f.name, "...")
21        f.write(TEMPLATE.format(case_name=sys.argv[1]))
22