Revert "Add nvprims.var_mean (pytorch#83508)"

This reverts commit 7e7694b.

Reverted pytorch#83508 on behalf of https://github.com/facebook-github-bot due to Diff reverted internally

Author: pytorchmergebot

test/test_decomp.py

@@ -188,7 +188,6 @@ def op_assert_equal(test_case, op, test_dtype, orig, decomp, args, kwargs):
             1e-3,
             1e-3,
         ),
-        (torch.float64, torch.ops.aten.native_layer_norm.default): (1e-6, 1e-6),
     }
     if (test_dtype, op) in tol_table:
         rtol, atol = tol_table[(decomp.dtype, op)]

test/test_prims.py

@@ -353,31 +353,6 @@ def _wrapper(a):
             self.assertTrue(result.is_contiguous)
             self.assertEqual(_wrapper(a), result)
 
-    @onlyCUDA
-    @skipCUDAIfRocm
-    @dtypes(torch.float16, torch.float32)
-    @parametrize("correction", [0, 1])
-    @parametrize("keepdim", [True, False])
-    def test_var_mean(self, device, dtype, correction, keepdim):
-        from torch.fx.experimental.proxy_tensor import make_fx
-        from torch._prims.context import TorchRefsNvfuserCapabilityMode
-
-
-        def _wrapper(a):
-            return torch.var_mean(a, [0, 1], correction=correction, keepdim=keepdim)
-
-        make_arg = partial(make_tensor, device=device, dtype=dtype)
-
-        with TorchRefsNvfuserCapabilityMode():
-            gm = make_fx(_wrapper)(make_arg((5, 5)))
-
-        call_function_nodes = list(filter(lambda n: n.op == "call_function", gm.graph.nodes))
-        includes_nvprims_var_mean = any(
-            torch.ops.nvprims.var_mean.main == node.target
-            for node in call_function_nodes
-        )
-        self.assertTrue(includes_nvprims_var_mean)
-
     @onlyCUDA
     @skipCUDAIfRocm
     @dtypes(torch.float32)

torch/_prims/__init__.py

@@ -20,7 +20,6 @@
     DimsType,
     Number,
     NumberType,
-    RETURN_TYPE,
     ShapeType,
     StrideType,
     TensorLike,
@@ -281,6 +280,20 @@ def TensorMeta(
 #
 # Common datastructures and helpers
 #
+
+# Describes the return type of the primitive:
+#
+#   - NEW, a new tensor is created
+#   - VIEW, a view of an input tensor is returned
+#   - INPLACE, one or more input tensors is modified
+#
+# these descriptors are mututally exclusive and exhaustive.
+class RETURN_TYPE(Enum):
+    NEW = (0,)
+    VIEW = (1,)
+    INPLACE = (2,)
+
+
 def _wrap_tensor_meta(f):
     def wrap(t):
         if (

torch/_prims/context.py

@@ -12,7 +12,6 @@
 import torch._refs.nn.functional
 import torch._refs.special
 import torch.overrides
-from torch._prims.nvfuser_executor import NvfuserPrimOperatorSupport
 
 from torch._prims_common import torch_function_passthrough
 from torch.fx.experimental.proxy_tensor import get_isolated_graphmodule
@@ -205,42 +204,15 @@ def _is_func_unsupported_nvfuser(torch_function_mode, func, args, kwargs):
     ):
         gm = get_isolated_graphmodule(func, args, kwargs)
 
-    supported_ops = NvfuserPrimOperatorSupport()
     call_function_nodes = filter(lambda n: n.op == "call_function", gm.graph.nodes)
     any_unsupported = any(
-        not supported_ops.is_node_supported(None, node) for node in call_function_nodes
+        not _is_node_supported_nvfuser(node) for node in call_function_nodes
     )
     return any_unsupported
 
 
-class TorchRefsNvfuserCapabilityMode(TorchRefsMode):
-    def __init__(self):
-        super().__init__(
-            strict=False,
-            should_fallback_fn=_is_func_unsupported_nvfuser,
-            prims_mode_cls=NvfuserPrimsMode,
-        )
-
-    def _is_var_mean(self, func):
-        return "torch.var_mean" == torch.overrides.resolve_name(func) or (
-            (
-                isinstance(func, torch._ops.OpOverload)
-                or isinstance(func, torch._ops.OpOverloadPacket)
-            )
-            and "aten.var_mean" in str(func)
-        )
-
-    def __torch_function__(
-        self,
-        orig_func: Callable,
-        types: Sequence,
-        args: Sequence[Any] = (),
-        kwargs: Dict = None,
-    ):
-        if kwargs is None:
-            kwargs = {}
-        # First we intercept calls for nvfuser-specific prims bypassing generic torch._refs
-        if self._is_var_mean(orig_func):
-            return torch.ops.nvprims.var_mean(*args, **kwargs)
-        # Then we use TorchRefsMode to interpret the rest
-        return super().__torch_function__(orig_func, types, args, kwargs)
+TorchRefsNvfuserCapabilityMode = functools.partial(
+    TorchRefsMode,
+    should_fallback_fn=_is_func_unsupported_nvfuser,
+    prims_mode_cls=NvfuserPrimsMode,
+)

torch/_prims/nvfuser_executor.py

@@ -57,8 +57,6 @@ def make_nvfuser_fusion(gm: GraphModule, *nv_args_templates):
     # PROTOTYPE nvfuser executor
     # Everything in the graph must support nvfuser
     for node in gm.graph.nodes:
-        if node.op == "call_function" and "getitem" in node.name:
-            continue
         if (
             node.op == "call_function"
             and getattr(node.target, "impl_nvfuser", None) is None
@@ -79,10 +77,6 @@ def _to_nvfuser_constant(arg):
 
         class FusionInterpreter(torch.fx.Interpreter):
             def call_function(self, target, args, kwargs):
-                # This handles tuple unpacking
-                if "getitem" in str(target):
-                    assert isinstance(args[0], tuple)
-                    return target(*args, **kwargs)
                 args = tuple(map(_to_nvfuser_constant, args))
                 target = target.impl_nvfuser
                 args = (fd,) + args
@@ -138,7 +132,6 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         return (
             node.op == "call_function"
             and getattr(node.target, "impl_nvfuser", None) is not None
-            or "getitem" in node.name  # getitem is a special case
         )
 
 

torch/_prims/nvfuser_prims.py

@@ -5,31 +5,24 @@
 # can be added in the future for the corresponding higher-level torch/aten
 # functions.
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 
 import torch
 
 from torch._prims_common import (
     DimsSequenceType,
-    ELEMENTWISE_TYPE_PROMOTION_KIND,
     getnvFuserDtype,
     ShapeType,
     TensorLikeType,
 )
 
-from torch._prims_common.wrappers import (
-    backwards_not_supported,
-    elementwise_type_promotion_wrapper,
-)
+from torch._prims_common.wrappers import backwards_not_supported
 
 nvprim_namespace = "nvprims"
 nvprim = torch.library.Library(nvprim_namespace, "DEF")
 nvprim_impl = torch.library.Library(
     nvprim_namespace, "IMPL", "CompositeExplicitAutograd"
 )
-nvprim_implicit_impl = torch.library.Library(
-    nvprim_namespace, "IMPL", "CompositeImplicitAutograd"
-)
 nvprim_autograd_impl = torch.library.Library(nvprim_namespace, "IMPL", "Autograd")
 nvprim_meta_impl = torch.library.Library(nvprim_namespace, "IMPL", "Meta")
 
@@ -241,23 +234,6 @@ def _var_nvfuser(
     return fd.ops.var(a, dims, correction, keep_dims)
 
 
-def _var_mean_nvfuser(
-    fd: Any,
-    a: TensorLikeType,
-    dims: DimsSequenceType,
-    unbiased: Optional[bool] = None,
-    keepdim: bool = False,
-    *,
-    correction: int,
-):
-    # Unbiased arg shouldn't be set when this function is called
-    assert unbiased is None
-    # Ignore keepdim arg, because currently it's automatically converted into nvfuser's symbolic scalar
-    # keepdim is handled by the reference implementation
-    keepdim = False
-    return fd.ops.var_mean(a, dims, correction, keepdim)
-
-
 def _amax_nvfuser(
     fd: Any,
     a: TensorLikeType,
@@ -280,112 +256,12 @@ def _amin_nvfuser(
 _nvfuser_impls["convert_element_type"] = _convert_element_type_nvfuser
 _nvfuser_impls["sum"] = _sum_nvfuser
 _nvfuser_impls["var"] = _var_nvfuser
-_nvfuser_impls["var_mean"] = _var_mean_nvfuser
 _nvfuser_impls["amax"] = _amax_nvfuser
 _nvfuser_impls["amin"] = _amin_nvfuser
 
 
-def register_var_mean():
-    """This function is used to register the var_mean function in torch.ops.nvprims module."""
-    name = "var_mean.main"
-
-    # This overload must be default for correct dispatching of var_mean(Tensor, bool)
-    nvprim.define("var_mean(Tensor inp, bool unbiased) -> (Tensor, Tensor)")
-
-    # This signature tries to combine several overloads of the torch.var_mean function into one overload.
-    nvprim.define(
-        f"{name}(Tensor inp, int[1]? dim=None, bool? unbiased=None, bool keepdim=False, *, int? correction=None)"
-        + " -> (Tensor, Tensor)"
-    )
-
-    # This function is used for device="meta" Tensors.
-    def _meta_var_mean(inp, dim=None, unbiased=None, keepdim=False, *, correction=None):
-        if torch._prims_common.is_complex_dtype(inp.dtype):
-            output_dtype = torch._prims_common.corresponding_real_dtype(inp.dtype)
-        else:
-            output_dtype = inp.dtype
-        var = torch._prims._reduction_meta(inp, dim, output_dtype=output_dtype)
-        mean = torch._prims._reduction_meta(inp, dim, output_dtype=inp.dtype)
-        if keepdim:
-            output_shape = [
-                inp.shape[i] if i not in dim else 1 for i in range(inp.ndim)
-            ]
-            broadcast_dims = [i for i in range(inp.ndim) if i not in dim]
-            var = torch.ops.nvprims.broadcast_in_dim(var, output_shape, broadcast_dims)
-            mean = torch.ops.nvprims.broadcast_in_dim(
-                mean, output_shape, broadcast_dims
-            )
-        return (var, mean)
-
-    # This function is used under _AutoDispatchBelowAutograd context
-    def _prim_impl(inp, dim=None, unbiased=None, keepdim=False, *, correction=None):
-        correction = torch._prims_common.set_correction(unbiased, correction)
-        return torch.var_mean(inp, dim, correction=correction, keepdim=keepdim)
-
-    nvprim_impl.impl(name, _prim_impl)
-    nvprim_meta_impl.impl(name, _meta_var_mean)
-
-    prim_packet = torch.ops.nvprims.var_mean
-    prim = prim_packet.main
-
-    def _unbiased_overload_impl(inp, unbiased):
-        return prim(inp, dim=None, unbiased=unbiased)
-
-    nvprim_implicit_impl.impl("var_mean", _unbiased_overload_impl)
-
-    @elementwise_type_promotion_wrapper(
-        type_promoting_args=("a",),
-        type_promotion_kind=ELEMENTWISE_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT,
-    )
-    def _var_mean_ref(a, dim=None, unbiased=None, keepdim=False, *, correction=None):
-        correction = torch._prims_common.set_correction(unbiased, correction)
-        # reduces over all dimensions if dim=() is passed
-        if dim == () or dim == []:
-            dim = None
-        dim = torch._prims_common.reduction_dims(a.shape, dim)
-
-        # For complex tensors eager computes the variance as the sum of variances of
-        # the real and imaginary parts
-        # TODO: Creating a complex tensor from real and imaginary parts is not supported
-        if torch._prims_common.is_complex_dtype(a.dtype):
-            raise NotImplementedError("Complex tensors are not supported")
-
-        var_mean = prim(a, dim, correction=correction)
-
-        if keepdim:
-            output_shape = [a.shape[i] if i not in dim else 1 for i in range(a.ndim)]
-            broadcast_dims = [i for i in range(a.ndim) if i not in dim]
-            var, mean = var_mean
-            var = torch.ops.nvprims.broadcast_in_dim(var, output_shape, broadcast_dims)
-            mean = torch.ops.nvprims.broadcast_in_dim(
-                mean, output_shape, broadcast_dims
-            )
-            var_mean = (var, mean)
-        return var_mean
-
-    def _var_mean_autograd(
-        a, dim=None, unbiased=None, keepdim=False, *, correction=None
-    ):
-        # This wrapper is needed to convert prims calls inside
-        # elementwise_type_promotion_wrapper to nvprims calls
-        from torch._prims.context import NvfuserPrimsMode
-
-        with NvfuserPrimsMode():
-            return backwards_not_supported(_var_mean_ref)(
-                a, dim, unbiased, keepdim, correction=correction
-            )
-
-    nvprim_autograd_impl.impl(name, _var_mean_autograd)
-
-    for p in (prim_packet, prim):
-        p.__doc__ = "Computes the variance and mean of x over the list of dimensions specified in the dim argument"
-        p.impl_nvfuser = _nvfuser_impls["var_mean"]
-        p.return_type = torch._prims_common.RETURN_TYPE.NEW  # type: ignore[attr-defined]
-
-
 def register_nvprims():
     """Registers all nvFuser primitives in the torch.ops.nvprims module."""
-    register_var_mean()
     for name in nvprim_names:
         main_prim = getattr(torch.ops.prims, name)
 

torch/_prims_common/__init__.py

@@ -1033,19 +1033,6 @@ class REDUCTION_OUTPUT_TYPE_KIND(Enum):
     ALWAYS_BOOL = (3,)
 
 
-# Describes the return type of the primitive:
-#
-#   - NEW, a new tensor is created
-#   - VIEW, a view of an input tensor is returned
-#   - INPLACE, one or more input tensors is modified
-#
-# these descriptors are mututally exclusive and exhaustive.
-class RETURN_TYPE(Enum):
-    NEW = (0,)
-    VIEW = (1,)
-    INPLACE = (2,)
-
-
 # TODO: document type promotion kinds
 def elementwise_dtypes(
     *_args,
@@ -1361,23 +1348,6 @@ def reduction_dims(shape: ShapeType, dims: Optional[Sequence]) -> Tuple[int, ...
     return dims
 
 
-def set_correction(
-    unbiased: Optional[bool] = None,
-    correction: Optional[int] = None,
-):
-    if correction is not None and unbiased is not None:
-        raise RuntimeError("cannot specify both correction and unbiased arguments")
-    elif correction is None and unbiased is None:
-        correction = 1
-    elif correction is None and unbiased is not None:
-        correction = 0 if unbiased is False else 1
-    if not isinstance(correction, int):
-        raise ValueError("correction argument should be integer")
-    if correction < 0:
-        raise ValueError("correction argument should be non-negative")
-    return correction
-
-
 def check_in_bounds_for_storage(
     a: torch.TypedStorage, shape: ShapeType, strides: StrideType, storage_offset: int
 ):

torch/_prims_common/wrappers.py

@@ -118,11 +118,9 @@ def _fn(*args, **kwargs):
 
             result = fn(**bound.arguments)
 
-            if isinstance(result, TensorLike):
-                return _maybe_convert_to_dtype(result, result_dtype)
-            if isinstance(result, Sequence):
-                return tuple(_maybe_convert_to_dtype(x, result_dtype) for x in result)
-            raise AssertionError(f"Unhandled result type: {type(result)}")
+            # FIXME?: assumes result is a single tensor
+            assert isinstance(result, TensorLike)
+            return _maybe_convert_to_dtype(result, result_dtype)
 
         _fn.__signature__ = sig  # type: ignore[attr-defined]
         return _fn