Skip to content

Fix PyTorch 2.13 named overload compatibility#102

Open
juicyzhou wants to merge 2 commits into
google:mainfrom
juicyzhou:agent/fix-pytorch-2.13-named-overloads
Open

Fix PyTorch 2.13 named overload compatibility#102
juicyzhou wants to merge 2 commits into
google:mainfrom
juicyzhou:agent/fix-pytorch-2.13-named-overloads

Conversation

@juicyzhou

@juicyzhou juicyzhou commented Jul 14, 2026

Copy link
Copy Markdown

Fixes #101.

Summary

  • avoid dereferencing named-tensor ATen overloads that were removed in PyTorch 2.13
  • remove named-tensor entries from the decomposition list
  • keep the two autocast registrations only when an older PyTorch version still provides those overloads
  • add regression coverage for importing torchax and constructing the default environment
  • add a focused PyTorch 2.13 compatibility CI job while keeping the existing PyTorch 2.8 full test matrix unchanged

Root cause

PyTorch 2.13 removed named tensors and their ATen overloads. torchax.amp accessed two removed overloads at module import time, and torchax.decompositions referenced another 20 removed overloads while constructing the default environment. Accessing any missing overload raises AttributeError before torchax can be used.

The removed decomposition entries are now deleted directly. The two autocast entries are registered conditionally so older PyTorch versions retain their existing behavior.

Validation

  • PyTorch 2.8.0, Python 3.12: pytest -q test/test_torch_compat.py test/test_amp.py — 6 passed
  • PyTorch 2.13.0, Python 3.12: pytest -q test/test_torch_compat.py test/test_amp.py — 6 passed
  • PyTorch 2.13.0: import torchax and torchax.default_env() succeed
  • ruff check and ruff format --check pass for all changed Python files

@google-cla

google-cla Bot commented Jul 14, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Comment thread torchax/decompositions.py Outdated
@juicyzhou
juicyzhou marked this pull request as ready for review July 19, 2026 16:18
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.

torch named tensors API (Dimname etc.) incompatible with torch >= 2.13

2 participants