Cannot mutate inputs in aot_export workflow

What’s the rationale for disabling mutations to user inputs while tracing via torch.export?
Is it a soundness issue? Not implemented? Something else?

Context: I’m trying to run a modified version of Llama through torch.export where the Attention module needs to update the kv-cache, which is in turn given as an argument of forward.

The referenced restriction precludes the AOT flow.


class AdvancedIndexingModule(nn.Module):

    def __init__(self, x):
        self.x = x

    def forward(self, x, y):
        x[:10, 5 : 10 + 5] = y
        return self.x

module = AdvancedIndexingModule(torch.rand(20, 20))
inputs = (
    torch.rand(20, 20),
    torch.rand(10, 10),

ep = export.export(module, inputs)


Traceback (most recent call last):
  File "/home/anieto/Groq/Groq/Compiler/test/end_to_end/demos/llm_tools/", line 33, in <module>
    ep = export.export(module, inputs)
  File "/home/anieto/.local/lib/python3.10/site-packages/torch/export/", line 449, in export
    return export__RC__(
  File "/home/anieto/.local/lib/python3.10/site-packages/torch/_export/", line 258, in export__RC__
    return _export(
  File "/home/anieto/.local/lib/python3.10/site-packages/torch/_export/", line 567, in wrapper
    return fn(*args, **kwargs)
  File "/home/anieto/.local/lib/python3.10/site-packages/torch/_export/", line 713, in _export
    gm, graph_signature = aot_export_module(
  File "/home/anieto/.local/lib/python3.10/site-packages/torch/_functorch/", line 5157, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
  File "/home/anieto/.local/lib/python3.10/site-packages/torch/_functorch/", line 5339, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
  File "/home/anieto/.local/lib/python3.10/site-packages/torch/_dynamo/", line 244, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/anieto/.local/lib/python3.10/site-packages/torch/_functorch/", line 4569, in create_aot_dispatcher_function
    raise RuntimeError(f"""
Found following user inputs located at [0] are mutated. This is currently banned in the aot_export workflow.
If you need this functionality, please file a github issue.