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

[ocramz]

PyTorch 2.13.0, released 2026-07-08, removed named tensors (PR #173895), deleting the Dimname overloads , including aten.prod.dim_Dimname.

torchax's amp.py still references that overload in a module-level CastPolicy dict at import time, so importing torchax with a newer PyTorch version crashes.

pytorch/pytorch#173895

[juicyzhou]

Hi, I would like to work on this compatibility fix. I confirmed that current main also references removed named-dimension overloads in both torchax/amp.py and torchax/decompositions.py. My plan is to make registration of those overloads conditional so torchax can import and initialize on PyTorch 2.13 while preserving behavior on older PyTorch versions, then add a regression test and run the targeted test suite against 2.13 and an older supported version. I will keep the change scoped to the named-tensor removal. Please let me know if you prefer a different compatibility approach.