Fix KeyError 'i23' with symbolic shapes and reshape bsyms

[shino16]

Partial fix for #2677. update_aliases.py manually inserts reshape bsyms for aliases with different shapes. In this PR, we additionally insert prims.shape to acquire the new symbolic shape in the trace.

The tests will fail due to this line.

lightning-thunder/thunder/core/update_aliases.py

Line 75 in 1a2c7d9

if arg.numel != arg_to_replace.numel:

It will disappear once #2760 is merged.

[Copilot]

Pull Request Overview

This PR fixes a KeyError 'i23' that occurs when using symbolic shapes with reshape bsyms in the alias update mechanism. The fix ensures that when reshaping aliased inputs with different shapes, the symbolic shape is properly acquired using prims.shape before performing the reshape operation.

Key Changes:

• Modified replace_args_with_alias_map to insert prims.shape bsym before reshape when using symbolic values
• Added test coverage for aliasing with viewed inputs of different shapes
• Updated existing tests to parameterize cache mode ("constant values" vs "symbolic values")

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
thunder/core/update_aliases.py	Updated reshape handling to include shape acquisition for symbolic values and changed from single bsym to tuple of bsyms
thunder/tests/test_update_aliases.py	Added cache parameterization to existing tests and new test for viewed inputs with different shapes

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[beverlylytle]

While this PR makes sense for what I imagine to be most cases, I think we are still not gracefully handling the following situation:

def f(x, y) : ...
jf = thunder.jit(f, cache='symbolic values')

x = torch.randn(4)

jf(x, x)

y = torch.randn(4)
z = y[0:2]

jf(y,z)

In the second application of jf, we won't go through the checks about numel and we will instead jump directly to a compiled function that will attempt to reshape z into y, which will produce an error. I'd really rather skip this reshaping and replacing hubbub, but I was running into trouble with numerical errors with nvfuser when I tried.

[shino16]

Hmm, that's a good point. I will wait until seeing how #2760 (comment) will end up.

[beverlylytle]

Looks great! Thanks @shino16!