Functionalize in-place writes done through views in traced JIT graphs.
torch.jit.trace SSA-renames reads of a tensor mutated DIRECTLY by an
in-place op (later reads point at the op's output value), but a mutation
done THROUGH A VIEW (buf[:, :, i] = x, traced as slice/select +
aten::copy_) leaves later reads of the parent buffer pointing at the
pre-write value. Executing eagerly this is invisible (memory aliasing),
but any purely functional backend such as NNEF silently loses the write.
HF transformers hits this pattern in every recurrent/linear-attention
model (core_attn_out[:, :, i] = ..., cache copy_ updates, ...).
functionalize_view_inplace_copy rewrites each such aten::copy_ into
the functional aten::select_scatter / aten::slice_scatter chain and
re-points every later reader of the root buffer (and of the intermediate
view values) at the updated value.
Writes that cannot be expressed this way (strided slices, view chains
through transpose/reshape/... aliases) are left untouched; when the
pre-write value is observably read afterwards, a loud warning names the
offending source location (e.g. HF apply_interleaved_mrope, whose
lost write is numerically a no-op for text-only position ids).
Functions:
| Name | Description |
|---|---|
functionalize_view_inplace_copy |
Rewrite |
Rewrite copy_ through slice/select views into functional scatters.
Returns the number of aten::copy_ nodes rewritten. Chains that
cannot be functionalized (strided slice, non-slice/select alias in
the chain root) are left as-is; if their pre-write value is read
afterwards a warning names the write's source location, since the
exported graph then keeps the pre-write data.