Hi Committers,
This PR is trying to fix issues
https://github.com/apache/tvm/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<Expr> 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.
---------
Co-authored-by: cchung100m <cchung100m@users.noreply.github.com>
Neo Chien
c12debb462