[APS] Implement __torch_function__ for KeyedTensor (pytorch#161683)

Summary:
X-link: meta-pytorch/torchrec#3329

Pull Request resolved: pytorch#161683

1.  There are a bunch of `torch.ops.aten` operations that can't handle `KeyedTensor`: The error was occurring because these ops expects a regular `Tensor` but was receiving a `KeyedTensor` object.

2.  Implement `__torch_function__` for `KeyedTensor`, so when these incompatible operations are called with a `KeyedTensor`, the `__torch_function__` method automatically delegates the op to the underlying values tensor from the `KeyedTensor` and returns a new `KeyedTensor` with updated values.

Test Plan:
```
buck2 run mode/opt fbcode//aps_models/ads/gmp:launcher_with_publish mode=mtml_mobile_cvr_model/managed/Y2025Q2/local_mode_mtml_mobile_cvr_model_733415799_v0_fork +training.ir_serializer=manifold
```

MAST job: https://fburl.com/mlhub/pp937uxf

Rollback Plan:

Reviewed By: malaybag

Differential Revision: D81047278

Author: jd7-tr

torch/_subclasses/fake_impls.py

@@ -817,10 +817,16 @@ def assert_tensor_metadata(
             f"Tensor layout mismatch! Expected: {layout}, Got: {t.layout()}"
         )
     if device is not None:
-        assert t.device == device, (
-            f"Tensor device mismatch! Expected: {device}, Got: {t.device}"
+        assert t.device.type == device.type, (
+            f"Tensor device type mismatch! Expected: {device.type}, Got: {t.device.type}"
+        )
+        assert (
+            (not device.index)
+            or (not t.device.index)
+            or (t.device.index == device.index)
+        ), (
+            f"Tensor device index mismatch! Expected: {device.type}, Got: {t.device.type}"
         )
-
 
 # NB: this must be ordered after local_scalar_dense
 @register_op_impl(lambda func: torch.Tag.data_dependent_output in func.tags)