xref: /aosp_15_r20/external/pytorch/torch/_lazy/debug.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# mypy: allow-untyped-defs
2import torch._C._lazy
3
4
5def render_ir_graph(tensors):
6    """Return a text dump of the LTC IR graph in dot format for the tensors.
7    The text can be processed by tools like dot to be rendered in pdf,png etc."""
8    return torch._C._lazy._get_tensors_dot(tensors)
9
10
11def dump_ir(tensors, ir_format):
12    """Return a dump of the tensors in the specified format.
13    Valid format are
14    - text: for LTC IR
15    - backend: for the activate backend IR
16    """
17    if ir_format == "text":
18        return torch._C._lazy._get_tensors_text(tensors)
19    elif ir_format == "backend":
20        return torch._C._lazy._get_tensors_backend(tensors)
21    else:
22        raise RuntimeError(f"Unrecognized IR format: {ir_format}")
23