[Relax][PyTorch] Fix segfault in from_exported_program when model uses index_put_ with tuple output

[cchung100m]

Hi Committers,

This PR is trying to fix issues #18363. Any suggestions would be appreciated if you are available.

Root Cause

• When an ExportedProgram's FX graph output node returns a nested Python tuple (e.g., buffer mutation outputs + user-defined tuple returns), _translate_fx_graph() passes the raw nested structure directly to the Relax FFI Tuple constructor.
• The C++ Array initializer cannot handle heterogeneous/nested Python containers, causing a segmentation fault at expr.cc.
• Additionally, index_put_ (in-place write op) did not update self.env to alias the source tensor to the mutated output, causing subsequent FX nodes that read the same tensor to observe stale pre-mutation values.

Solution

• exported_program_translator.py
  • Added static method _flatten_output_args() that recursively walks any Python tuple/list, collects only relax.Expr leaves, and preserve explicit None outputs as Relax null objects.
  • Replaced the fragile assert isinstance(output_args, tuple | relax.Tuple) guard with a call to _flatten_output_args(), producing a clean flat tuple of relax.Expr before FFI construction.
• base_fx_graph_translator.py
  • In _index_put(), after emitting the relax.op.index_put(...) call, added an env alias update: self.env[source_node] = output when the target op name starts with index_put_, preserving correct in-place mutation semantics for downstream FX nodes.

[gemini-code-assist]

Code Review

This pull request introduces support for in-place index_put_ operations in the Torch FX translator by updating the environment with the operation's result. It also adds a _flatten_output_args utility to the ExportedProgram translator to handle nested output structures that could cause FFI issues. Feedback indicates that the flattening utility is currently defined but not invoked, and suggests a more robust approach for extracting target names from FX nodes to handle cases where the target is a string.

[cchung100m]

Hi @tlopex @mshr-h

This PR is trying to fix issues #18363. Any suggestions would be appreciated if you are available. 😄

[tlopex]

Thanks for the fix. One correctness case I’d like to clarify before approval: _index_put only updates self.env[node.args[0]]. If the mutated value is an alias of another FX node, later reads through the original node may still see the pre-mutation value.

Could you add a small regression test for mutation through an alias?