Skip to content

Fix torch.stack and torch.cat recursion for CompositionScores - #1840

Merged
jlarson4 merged 2 commits into
TransformerLensOrg:devfrom
yuanwuyuan9:fix/composition-scores-torch-sequences
Sep 30, 2026
Merged

jlarson4 merged 2 commits into
TransformerLensOrg:devfrom
yuanwuyuan9:fix/composition-scores-torch-sequences

Conversation

@yuanwuyuan9

Copy link
Copy Markdown

Description

torch.stack and torch.cat recurse indefinitely when their input sequence contains CompositionScores: the current dispatch unwraps only top-level arguments and passes nested wrappers back to PyTorch.

Recursively unwrap wrappers in lists, tuples, and dictionary values before calling the torch function. Preserve ordinary arguments and the existing tensor-return behavior. No new dependencies.

Add 16 regression cases covering stack/cat, list/tuple inputs, positional/keyword arguments, and mixed wrapper/tensor inputs. All 16 fail with RecursionError before the fix and pass afterward.

Fixes #1839

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Validation

Check Result
New sequence-operation tests 16 passed
Existing CompositionScores protocol tests 8 passed
make unit-test 6784 passed, 61 skipped, 57 deselected, 6 xfailed, 306 warnings, 1 rerun
Gated-hooks benchmark test file, reruns disabled 4 passed
mypy No issues in 386 source files
Formatting and git diff --check Passed

The unit suite ran on CPU with OMP_NUM_THREADS=1 MKL_NUM_THREADS=1; uv synchronization was disabled to preserve the installed environment and lockfile. One test in test_benchmark_gated_hooks_fire.py required a rerun. All four tests in that file subsequently passed with --reruns 0; the initial failure cause remains undetermined. Full make test-pr and CI have not been verified locally.

Checklist

  • I have commented my code, particularly in hard-to-understand areas — updated the dispatch docstring to describe container unwrapping.
  • I have made corresponding changes to the documentation — no separate documentation update needed for this bug fix.
  • My changes generate no new warnings — targeted runs show the same SWIG warnings; a full-suite warning comparison was not performed.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes, with the configuration and rerun noted above.
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility.

@jlarson4 jlarson4 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @yuanwuyuan9, thanks for the contribution! Just one note about test location

Comment thread tests/unit/test_composition_scores_sequences.py Outdated
@jlarson4

Copy link
Copy Markdown
Collaborator

Looks good! Thanks for the contribution @yuanwuyuan9

@jlarson4
jlarson4 merged commit 886a5f5 into TransformerLensOrg:dev Sep 30, 2026
27 checks passed
@jlarson4 jlarson4 linked an issue Sep 30, 2026 that may be closed by this pull request
1 task done
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug Report] torch.stack and torch.cat recurse indefinitely with CompositionScores

2 participants