Unify aten operator canonicalization

### 🚀 The feature, motivation and pitch

# Represent semantic ATen operator aliases centrally instead of only in JIT `NormalizeOps`

I am bit uncertain of this best way to describe this, but it's a bit surprising to me that we have generated alias that do not match and do not properly have ways to cannoncalize them for wrappers or prim decompositions in the PyTorch operator. Inspired by work on #186566 where `aten::pow(2.0...)` did not export properly to aten::exp2 because decomps were registered to aten::special_exp2.

## Summary

PyTorch has many pairs of ATen operators that are semantically equivalent but have distinct operator identities, for example:

```text
aten::absolute -> aten::abs
aten::clip -> aten::clamp
aten::special_erf -> aten::erf
aten::special_exp2 -> aten::exp2
aten::special_expm1 -> aten::expm1
```

These relationships are currently encoded in a handwritten table in:

```text
torch/csrc/jit/passes/normalize_ops.cpp
```

For example:

```cpp
static const std::unordered_map alias_map = {
{aten::absolute, aten::abs},
{aten::clip, aten::clamp},
// ...
{aten::special_erf, aten::erf},
{aten::special_exp2, aten::exp2},
{aten::special_expm1, aten::expm1},
// ...
};
```

The JIT `NormalizeOps` pass uses this table to rewrite alias operators into a canonical form:

```cpp
auto alias = getOperatorAliasMap().find(iter->kind());
if (alias != getOperatorAliasMap().end()) {
iter->replaceWithNewSymbol(alias->second);
iter.destroyCurrent();
return true;
}
```

The problem is that this alias information is only available to this JIT graph pass. Other PyTorch subsystems still treat the alias and canonical operator as unrelated identities.

## Concrete failure: `exp2`

This issue surfaced while fixing the derivative of `ldexp`.

The derivative multiplier is naturally written as:

```cpp
at::exp2(other)
```

instead of:

```cpp
at::pow(2.0, other)
```

However, using `at::exp2` exposed a missing reference or Meta implementation in some test paths.

At first glance, this is surprising because `aten::exp2` already has a PrimTorch reference:

```python
@_make_elementwise_unary_reference(
ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT
)
def exp2(a):
return prims.exp2(a)
```

The corresponding primitive is currently defined as:

```python
exp2 = _make_elementwise_unary_prim(
"exp2",
impl_aten=torch.special.exp2,
doc="",
type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT,
)
```

The important detail is that `torch.exp2` and `torch.special.exp2` are not the same callable or the same ATen operator:

```python
import torch

assert torch.special.exp2 is not torch.exp2
assert torch.ops.aten.special_exp2.default is not torch.ops.aten.exp2.default
```

Therefore the reference path changes operator identity:

```text
aten.exp2
-> torch._refs.exp2
-> prims.exp2
-> aten.special_exp2
```

The JIT alias map knows that:

```text
aten.special_exp2 -> aten.exp2
```

but PrimTorch, decomposition lookup, FakeTensor, Meta, export, and NumPy-reference machinery do not necessarily consult that map.

As a result, machinery with support for `aten.exp2` can unexpectedly encounter `aten.special_exp2` and report that its decomposition, Meta implementation, or NumPy reference is missing.

## Immediate fix for `exp2`

The primitive should use the canonical operator with the same name:

```diff
exp2 = _make_elementwise_unary_prim(
"exp2",
- impl_aten=torch.special.exp2,
+ impl_aten=torch.exp2,
doc="",
type_promotion=ELEMENTWISE_PRIM_TYPE_PROMOTION_KIND.DEFAULT,
)
```

This keeps the lowering path consistent:

```text
aten.exp2
-> torch._refs.exp2
-> prims.exp2
-> aten.exp2
```

This is the correct narrow fix for the concrete failure.

A regression test should exercise `aten.exp2` under the relevant reference, FakeTensor, or Meta mode, including integer inputs:

```python
import torch
from torch._prims.context import TorchRefsMode

x = torch.tensor([-2, -1, 0, 1], dtype=torch.int64)

with TorchRefsMode():
actual = torch.exp2(x)

expected = torch.tensor([0.25, 0.5, 1.0, 2.0])
torch.testing.assert_close(actual, expected)
```

## Broader structural issue

Changing `prims.exp2` fixes this particular bug, but the same class of failure can recur for every semantic alias in `normalize_ops.cpp`.

The current design has several separate notions of aliasing:

### Python API alias

Two Python names refer to the same callable:

```python
torch.special.exp2 is torch.exp2
```

This is false.

### Dispatcher identity

Two names resolve to the same dispatcher operator:

```python
torch.ops.aten.special_exp2.default
torch.ops.aten.exp2.default
```

These are also distinct.

### JIT normalization alias

A graph pass knows that one operator should be rewritten to another:

```text
aten.special_exp2 -> aten.exp2
```

This is the only alias relationship that currently exists for this pair, and it is local to the JIT normalization pass.

Consequently, each subsystem either needs to duplicate this knowledge or treats the two operators as unrelated.

Affected systems may include:

* PrimTorch
* decomposition lookup
* FakeTensor and Meta
* functionalization
* export
* NumPy-reference testing
* TorchDynamo and compiler IRs
* backend lowering
* operator analysis tools

## Proposed structural fix

Semantic operator aliases should be represented in a central, codegen-visible source of truth.

I would avoid calling this field simply `alias_of`, because “alias” already has a storage-aliasing meaning in ATen schemas. A name such as `canonical_op` or `semantic_alias_of` would be less ambiguous.

Conceptually:

```yaml
- func: special_exp2(Tensor self) -> Tensor
canonical_op: exp2
```

For overload-sensitive operators, the full overload could be specified:

```yaml
- func: special_exp2(Tensor self) -> Tensor
canonical_op: exp2.default
```

This metadata should not merge the two dispatcher identities. Keeping distinct schemas is important for backward compatibility and for systems such as `__torch_dispatch__`, custom backends, `torch.library`, tracing, and direct uses of:

```python
torch.ops.aten.special_exp2.default
```

Instead, it should declare that one operator has a canonical semantic equivalent.

## Codegen validation

Torchgen should validate every canonical relationship.

At minimum:

* the target operator exists;
* the relationship is acyclic;
* argument count and order match;
* argument types match;
* return types match;
* alias annotations match;
* mutability matches;
* overloads are compatible.

This mechanism should only describe direct semantic aliases.

Operators requiring transformations should remain separate normalization rules. For example, `rsub` requires argument reordering and is correctly handled by custom normalization logic rather than the simple alias map.

## Generated canonicalization API

Codegen should expose the relationship to both C++ and Python consumers.

For example:

```cpp
std::optional canonicalOperator(
const c10::OperatorName& op);
```

Conceptually:

```cpp
canonicalOperator("aten::special_exp2")
== "aten::exp2";
```

A Python equivalent could operate on exact overload identities:

```python
canonical_op(torch.ops.aten.special_exp2.default)
# torch.ops.aten.exp2.default
```

Using exact operator names and overloads is important because decomposition tables are keyed by `OpOverload`, not only by base symbol name.

## Generate the JIT map from this metadata

The existing JIT normalization behavior should remain, but its handwritten table should be generated from the central metadata.

Instead of maintaining:

```cpp
{aten::special_exp2, aten::exp2},
```

manually in `normalize_ops.cpp`, codegen would generate the alias-to-canonical map consumed by `NormalizeOps`.

This preserves current behavior while preventing JIT from being the only subsystem aware of these relationships.

## Integration with decompositions

For schema-compatible aliases, PyTorch could generate forwarding decompositions:

```python
@register_decomposition(aten.special_exp2.default)
def special_exp2_decomposition(x):
return aten.exp2.default(x)
```

This produces an explicit canonical graph:

```text
aten.special_exp2
-> aten.exp2
-> prims.exp2
```

Alternatively, decomposition registration could optionally propagate across semantic aliases:

```python
@register_decomposition(
aten.exp2.default,
include_semantic_aliases=True,
)
def exp2_reference(x):
return prims.exp2(x)
```

Generated forwarding decompositions are probably preferable because the resulting graph visibly uses the canonical operator rather than silently performing a canonical lookup while retaining the alias node.

## Compatibility

This proposal does not require deleting alias operators or making two schemas share a dispatcher identity.

The expected behavior would be:

```text
Dispatcher:
aten.special_exp2 and aten.exp2 remain distinct

Canonical metadata:
aten.special_exp2 -> aten.exp2

JIT:
rewrites special_exp2 to exp2

Decompositions:
forwards special_exp2 to exp2

PrimTorch:
uses canonical exp2

Export/compiler IR:
can normalize to canonical exp2
```

This preserves backward compatibility while giving compiler and reference subsystems a shared understanding of semantic equivalence.

## Suggested implementation sequence

1. Fix `prims.exp2` to use `impl_aten=torch.exp2`.
2. Add regression coverage for reference, FakeTensor, and Meta execution.
3. Introduce `canonical_op` or equivalent metadata in torchgen.
4. Validate schema compatibility and reject cycles.
5. Generate a C++ and Python canonicalization map.
6. Generate the JIT `NormalizeOps` alias map from that metadata.
7. Generate forwarding decompositions for directly compatible aliases.
8. Migrate the existing entries in `getOperatorAliasMap()` to the central metadata.

## Expected outcome

PyTorch should have one authoritative source of truth describing semantic operator aliases.

The dispatcher identities may remain distinct, but every relevant subsystem should be able to determine that:

```text
aten.special_exp2
```

has the canonical semantic form:

```text
aten.exp2
```

This would eliminate a class of missing decomposition, Meta, FakeTensor, export, and reference-support failures caused by alias relationships currently known only to the JIT.

### Alternatives

_No response_

### Additional context

_No response_

cc @ezyang @mruberry @manuelcandales @angelayi

Start with torch/csrc/jit/passes/normalize_ops.cpp, especially getOperatorAliasMap, and inspect the prims.exp2 implementation and the TorchRefsMode regression scenario described in the issue. Determine how canonical_op metadata could be validated and exposed to Python and C++, then generate the JIT map and forwarding decomposition behavior. Done means the exp2 regression is covered and semantic aliases have one validated source of truth without merging dispatcher identities.

Written by the indexing model from the issue text.