Repository: pytorch
Issue: AOTAutograd not tracing AMP correctly when using PyTorch/XLA.
Labels: triaged,module: xla,module: amp (automated mixed precision),oncall: pt2,module: aotdispatch,module: pt2-dispatcher

### 🐛 Describe the bug

Running, for example, `torch.norm` with AMP results in a tensor with the incorrect data-type. In the program below, we expect `out` to be a `float32` tensor due to [its autocast rule][1]. However, turns out the result is actually of type `float16`.

```python
@torch.compile(backend="openxla")
def foo(x):
    with torch.cuda.amp.autocast(dtype=torch.float16):
        return torch.norm(x, p=2, dim=1)

inp = torch.rand((5, 5), dtype=torch.float16, device="xla")
out = foo(inp) 
# Expected: float32
# Got: float16
```

[1]: https://github.com/pytorch/pytorch/blob/afb3e5ed6a950fcb60b763b4eeddca9ad8d58283/aten/src/ATen/autocast_mode.h#L870-L879

Even though this looks like a PyTorch/XLA issue, inspecting the FX graph received from AOTAutograd, we can observe that AOT recorded `torch.ops.aten.norm.ScalarOpt_dim` instead of `torch.ops.aten.norm.ScalarOpt_dim_dtype`. Because of that, PyTorch/XLA has no access to the appended `at::kFloat` at the end of the call.

```python
class <lambda>(torch.nn.Module):
    def forward(self, arg0_1: "f16[5, 5]"):
         # File: examples/test.py:21 in foo, code: return torch.norm(x, p=2, dim=1)
         # Should be: torch.ops.aten.norm.ScalarOpt_dim_dtype(arg0_1, 2, [1], torch.float32)
        norm: "f16[5]" = torch.ops.aten.norm.ScalarOpt_dim(arg0_1, 2, [1]);  arg0_1 = None
        return (norm,)
```

### Expected Behavior

I would expect AOTAutograd to record the function that is called by the autocast layer.

### Observations

1. I can't replicate this issue with `aot_eager`. The FX graph I get from AOT is using `torch.ops.aten.linalg_vector_norm` instead of `torch.ops.aten.norm.ScalarOpt_dim`. Not sure where this happened, though.
2. A similar behavior also seems to happen with `torch.nn.Conv2d` (and, possibly, its other multi-dimensional variants), where the casting of the arguments won't get registered
3. Other operations that have autocast rules should also suffer from this, but I haven't tested them all

Because of (1), it feels like it's an XLA-specific issue inside PyTorch. I don't think it's something inside PyTorch/XLA, since this is a matter of which operation is recorded in the AOTAutograd tracing.

cc @bdhirsh @mcarilli @ptrblck @leslie-fang-intel @jgong5 @ezyang @chauhang @penguinwu @zou3519 @miladm @JackCaoG @ajam

### Versions

PyTorch version: 2.5.0a0+git7128504
Is debug build: True
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 11 (bullseye) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: Could not collect
CMake version: version 3.29.2
Libc version: glibc-2.31

Python version: 3.10.14 (main, Apr 24 2024, 12:45:30) [GCC 10.2.1 20210110] (64-bit runtime)
Python platform: Linux-6.10.4-200.fc40.x86_64-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY

PyTorch/XLA version: [ac7fd44f520d586c3d95fc929134c20946e60806][2]

[2]: https://github.com/pytorch/xla/commit/ac7fd44f520d586c3d95fc929134c20946e60806