# Export sub-graphs at the Aten IR level

**URL:** <https://dev-discuss.pytorch.org/t/export-sub-graphs-at-the-aten-ir-level/1936>\
**Category:** compiler\
**Created:** [March 13, 2024, 3:02pm UTC](https://dev-discuss.pytorch.org/t/export-sub-graphs-at-the-aten-ir-level/1936 "2024-03-13T15:02:53Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![rawwds](https://avatars.discourse-cdn.com/v4/letter/r/f08c70/32.png) [@rawwds](https://dev-discuss.pytorch.org/u/rawwds)\
**Post date:** [March 13, 2024, 3:02pm UTC](https://dev-discuss.pytorch.org/t/export-sub-graphs-at-the-aten-ir-level/1936/1 "2024-03-13T15:02:53Z")

</div>

Regarding the `torch.export` documentation, it is designed to trace the full graph at the Aten IR level, but potential failures in some corner cases. In such instances, **is it possible to export multiple subgraphs, at the Aten IR level as well?**

The motivation behind this is that the full graph might be too restrictive. Even if we only obtain the subgraph, there is still space for optimization for each of them.

cc @voz@gmagogsfm @desertfire @SherlockNoMad (possible owner, kindly ignore it if not 🙂

---

<div class="post-metadata">

**Author:** ![voz](https://avatars.discourse-cdn.com/v4/letter/v/439d5e/32.png) [@voz](https://dev-discuss.pytorch.org/u/voz)\
**Post date:** [March 13, 2024, 3:44pm UTC](https://dev-discuss.pytorch.org/t/export-sub-graphs-at-the-aten-ir-level/1936/2 "2024-03-13T15:44:52Z")

</div>

Dynamo backends can be written as functions! One relatively straightforward approach here would be to pass a function to compile or torch.\_dynamo.optimize that then invokes aot\_autograd (or any other torch ir-\> aten ir step of your choosing) and records the functions in a list or structure of your choosing.

For an example of an existing function that “records” subgraphs in a custom small backend, check out `explain`.  
[https://github.com/pytorch/pytorch/blob/92a2b214f8e9217c5d94c7583817b22f822fa925/torch/\_dynamo/eval\_frame.py#L701](https://github.com/pytorch/pytorch/blob/92a2b214f8e9217c5d94c7583817b22f822fa925/torch/_dynamo/eval_frame.py#L701)

Note, that this isn’t really export anymore if you do it this way, since the very things that make export what it is seem to be different from what you want (whole graph) - but the step above would give you a single call way to extract all the graphs out. But it should give what you asked for.

Note, also, that in doing so, you silently are skipping whatever code couldn’t be captured and got delegated to eager. So, user beware!

Answered by mobile, apologies for being terse and any formatting issues.

---

<div class="post-metadata">

**Author:** ![rawwds](https://avatars.discourse-cdn.com/v4/letter/r/f08c70/32.png) [@rawwds](https://dev-discuss.pytorch.org/u/rawwds)\
**Post date:** [March 14, 2024, 3:29am UTC](https://dev-discuss.pytorch.org/t/export-sub-graphs-at-the-aten-ir-level/1936/3 "2024-03-14T03:29:09Z")

</div>

Hi @voz, thanks a lot for getting back to me so quickly. I’ve had a quick look, and it is exactly what I was searching for. I’ve got another question regarding the transformation from `torch IR` to `aten IR`. Could you provide an example of how to perform these transformations or direct me to the main entry point in the code?