# What is the mental model for torch.compile to deal with stateful objects in general?

**URL:** <https://dev-discuss.pytorch.org/t/what-is-the-mental-model-for-torch-compile-to-deal-with-stateful-objects-in-general/2080>\
**Category:** compiler\
**Created:** [May 13, 2024, 7:10am UTC](https://dev-discuss.pytorch.org/t/what-is-the-mental-model-for-torch-compile-to-deal-with-stateful-objects-in-general/2080 "2024-05-13T07:10:13Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![youkaichao](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/youkaichao/32/1278_2.png) [@youkaichao](https://dev-discuss.pytorch.org/u/youkaichao)\
**Post date:** [May 13, 2024, 7:10am UTC](https://dev-discuss.pytorch.org/t/what-is-the-mental-model-for-torch-compile-to-deal-with-stateful-objects-in-general/2080/1 "2024-05-13T07:10:13Z")

</div>

I’m trying to get a clear mental model for `torch.compile` to deal with stateful objects.

Traditionally, JIT compilers such as [numba](https://numba.pydata.org/) only compile for arrays and several Python basic data structures. It is easy to figure out when they re-compile: just check the metadata of input, e.g. `dtype, shape, type`.

In the case of `torch.compile`, the situation is quite complicated. The promise of `torch.compile` is essentially it can work for any function `f(*args, **kwargs)` with arbitary input, because it can understand Python bytecode. This makes the re-compile condition extremely difficult to figure out. Currently, it is kind of a black box.

The re-compile condition becomes even more difficult to figure out, when some arguments are stateful, e.g. `dict`, `list`, and notably, `nn.Module` . Recently, I encountered [such a problem](https://github.com/pytorch/pytorch/issues/125652), and had to use a global flag `export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1` to achieve what I want.

That’s why I’m asking. Does `torch.compile` have a mental model for the re-compile condition (or essentially the guard system)? e.g. what does it guard on? object type? object id? object content (for stateful objects)? In terms of stateful objects, how can we guard on the content though? We know Python is very flexible and dynamic, the content of a stateful object can change quite dramatically.

cc @jansel @anijain2305

---

<div class="post-metadata">

**Author:** ![jansel](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/jansel/32/74_2.png) [@jansel](https://dev-discuss.pytorch.org/u/jansel)\
**Post date:** [May 13, 2024, 6:45pm UTC](https://dev-discuss.pytorch.org/t/what-is-the-mental-model-for-torch-compile-to-deal-with-stateful-objects-in-general/2080/2 "2024-05-13T18:45:12Z")

</div>

The handling here is specific to `torch.nn.Module()` instances. Depending on the config, `torch.compile()` might specialize on the `id()` of a nn.Module. This allows additional optimizations (freezing, copy-free-cudagraphs, etc) since the module parameters are known at a compile time.

---

<div class="post-metadata">

**Author:** ![youkaichao](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/youkaichao/32/1278_2.png) [@youkaichao](https://dev-discuss.pytorch.org/u/youkaichao)\
**Post date:** [May 13, 2024, 7:06pm UTC](https://dev-discuss.pytorch.org/t/what-is-the-mental-model-for-torch-compile-to-deal-with-stateful-objects-in-general/2080/3 "2024-05-13T19:06:59Z")

</div>

While I know this specific issue is related with `nn.Module` instances, I’m trying to find a general mental model for how Dynamo places guards for stateful objects.

---

<div class="post-metadata">

**Author:** ![jansel](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/jansel/32/74_2.png) [@jansel](https://dev-discuss.pytorch.org/u/jansel)\
**Post date:** [May 13, 2024, 8:01pm UTC](https://dev-discuss.pytorch.org/t/what-is-the-mental-model-for-torch-compile-to-deal-with-stateful-objects-in-general/2080/4 "2024-05-13T20:01:06Z")

</div>

Dynamo guards on nn.Module `id()s` unless you mark the module dynamic (which happens automatically if you mutate `self`) or disable that behavior by config.
