1import argparse 2import os 3from dataclasses import fields 4 5from torch.onnx._internal import diagnostics 6from torch.onnx._internal.diagnostics import infra 7 8 9def gen_docs(out_dir: str): 10 os.makedirs(out_dir, exist_ok=True) 11 for field in fields(diagnostics.rules): 12 rule = getattr(diagnostics.rules, field.name) 13 if not isinstance(rule, infra.Rule): 14 continue 15 if not rule.id.startswith("FXE"): 16 # Only generate docs for `dynamo_export` rules. Excluding rules for TorchScript 17 # ONNX exporter. 18 continue 19 title = f"{rule.id}:{rule.name}" 20 full_description_markdown = rule.full_description_markdown 21 assert ( 22 full_description_markdown is not None 23 ), f"Expected {title} to have a full description in markdown" 24 with open(f"{out_dir}/{title}.md", "w") as f: 25 f.write(f"# {title}\n") 26 f.write(full_description_markdown) 27 28 29def main() -> None: 30 parser = argparse.ArgumentParser( 31 description="Generate ONNX diagnostics rules doc in markdown." 32 ) 33 parser.add_argument( 34 "out_dir", metavar="OUT_DIR", help="path to output directory for docs" 35 ) 36 args = parser.parse_args() 37 gen_docs(args.out_dir) 38 39 40if __name__ == "__main__": 41 main() 42