[BE] Get rid of unnecessary inner_torch_dispatch method (#132769)
Signed-off-by: Edward Z. Yang <ezyang@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/132769
Approved by: https://github.com/Skylion007, https://github.com/bdhirsh
ghstack dependencies: #132674, #132675, #132421, #132062, #132767
diff --git a/torch/fx/experimental/proxy_tensor.py b/torch/fx/experimental/proxy_tensor.py
index ae666d3..23c61f8 100644
--- a/torch/fx/experimental/proxy_tensor.py
+++ b/torch/fx/experimental/proxy_tensor.py
@@ -1302,7 +1302,12 @@
kwargs: Optional[Dict[str, object]] = None,
) -> object:
with set_original_aten_op(func):
- return self.inner_torch_dispatch(func, types, args, kwargs)
+ kwargs = kwargs or {}
+
+ if func in (prim.device.default,):
+ return func(*args, **kwargs)
+
+ return proxy_call(self, func, self.pre_dispatch, args, kwargs)
def __enter__(self) -> Self:
# Stash and store the previous proxy mode (there may or may not be one)
@@ -1325,20 +1330,6 @@
return b
- def inner_torch_dispatch(
- self,
- func: OpOverload,
- types: Tuple[torch._C._TensorMeta, ...],
- args: Tuple[object, ...] = (),
- kwargs: Optional[Dict[str, object]] = None,
- ) -> object:
- kwargs = kwargs or {}
-
- if func in (prim.device.default,):
- return func(*args, **kwargs)
-
- return proxy_call(self, func, self.pre_dispatch, args, kwargs)
-
@classmethod
def is_infra_mode(cls) -> bool:
return True
