Revert "Add sym_log2 (pytorch#137980)"

This reverts commit 5d450d7.

Reverted pytorch#137980 on behalf of https://github.com/jeanschmidt due to lint broke from this onwards on main ([comment](pytorch#137980 (comment)))

Author: pytorchmergebot

test/test_dynamic_shapes.py

@@ -499,24 +499,14 @@ def test_sym_int(self):
             str(shape_env.guards[2][0]), """Eq(TruncToInt(2.0*ToFloat(s2)), 6)"""
         )
 
-    def test_sym_log2(self):
-        shape_env = ShapeEnv()
-        a0 = create_symint(shape_env, 4)
-        r = torch._sym_log2(a0)
-        self.assertEqual(r, 2.0)
-        self.assertIsInstance(r, torch.SymFloat, msg=type(r))
-        self.assertExpectedInline(
-            str(shape_env.guards[0][0]), """Eq(OpaqueUnaryFn_log2(ToFloat(s0)), 2.0)"""
-        )
-
     def test_sym_sqrt(self):
         shape_env = ShapeEnv()
         a0 = create_symint(shape_env, 4)
         r = torch._sym_sqrt(a0)
         self.assertEqual(r, 2)
         self.assertIsInstance(r, torch.SymFloat, msg=type(r))
         self.assertExpectedInline(
-            str(shape_env.guards[0][0]), """Eq(OpaqueUnaryFn_sqrt(ToFloat(s0)), 2.0)"""
+            str(shape_env.guards[0][0]), """Eq(OpaqueUnaryFn_sqrt(s0), 2.0)"""
         )
 
     def test_sym_floor(self):
@@ -550,8 +540,7 @@ def test_sym_trunc(self):
         self.assertEqual(r, 2)
         self.assertIsInstance(r, torch.SymInt, msg=type(r))
         self.assertExpectedInline(
-            str(shape_env.guards[1][0]),
-            """Eq(TruncToInt(OpaqueUnaryFn_sqrt(ToFloat(s0))), 2)""",
+            str(shape_env.guards[1][0]), """Eq(TruncToInt(OpaqueUnaryFn_sqrt(s0)), 2)"""
         )
 
     def test_sym_ceil(self):

torch/__init__.py

@@ -876,8 +876,6 @@ def _get_sym_math_fn(name):
     def fn(a):
         if overrides.has_torch_function_unary(a):
             return overrides.handle_torch_function(fn, (a,), a)
-        if isinstance(a, SymInt):
-            a = torch.sym_float(a)
         if hasattr(a, f"__sym_{name}__"):
             return getattr(a, f"__sym_{name}__")()
         return getattr(math, name)(a)
@@ -897,7 +895,6 @@ def fn(a):
     "asin",
     "acos",
     "atan",
-    "log2",
 ):
     __sym_name = f"_sym_{__name}"
     __fn = _get_sym_math_fn(__name)

torch/fx/experimental/sym_node.py

@@ -1,8 +1,4 @@
 # mypy: allow-untyped-defs
-
-from __future__ import annotations
-
-
 """
 This file does three things:
 - Contains the definition of SymNode
@@ -149,12 +145,12 @@ def compute_hint():
         )
         self.fx_node = tx_validation_en and fx_node
 
-    def with_shape_env(self, shape_env: ShapeEnv) -> SymNode:
+    def with_shape_env(self, shape_env: "ShapeEnv") -> "SymNode":
         return SymNode(
             self._expr, shape_env, self.pytype, self._hint, self.constant, self.fx_node
         )
 
-    def _value_eq(self, other: SymNode) -> bool:
+    def _value_eq(self, other: "SymNode") -> bool:
         # Purposely don't include the shape_env in the eq.
         return (
             self._expr == other._expr
@@ -285,121 +281,121 @@ def _graph_repr(self) -> builtins.str:
 
     # These methods call the metaprogrammed methods, they're hand written
     # here so we get good stack traces
-    def abs(self) -> SymNode:
+    def abs(self) -> "SymNode":
         return self._abs()  # type: ignore[attr-defined]
 
-    def pos(self) -> SymNode:
+    def pos(self) -> "SymNode":
         return self._pos()  # type: ignore[attr-defined]
 
-    def round(self, ndigits=None) -> SymNode:
+    def round(self, ndigits=None) -> "SymNode":
         return self._round(ndigits)  # type: ignore[attr-defined]
 
-    def trunc(self) -> SymNode:
+    def trunc(self) -> "SymNode":
         return self._trunc()  # type: ignore[attr-defined]
 
-    def add(self, other) -> SymNode:
+    def add(self, other) -> "SymNode":
         return self._add(other)  # type: ignore[attr-defined]
 
-    def sub(self, other) -> SymNode:
+    def sub(self, other) -> "SymNode":
         return self._sub(other)  # type: ignore[attr-defined]
 
-    def mul(self, other) -> SymNode:
+    def mul(self, other) -> "SymNode":
         return self._mul(other)  # type: ignore[attr-defined]
 
-    def mod(self, other) -> SymNode:
+    def mod(self, other) -> "SymNode":
         return self._mod(other)  # type: ignore[attr-defined]
 
-    def float_pow(self, other) -> SymNode:
+    def float_pow(self, other) -> "SymNode":
         return self._float_pow(other)  # type: ignore[attr-defined]
 
-    def pow_by_natural(self, other) -> SymNode:
+    def pow_by_natural(self, other) -> "SymNode":
         return self._pow_by_natural(other)  # type: ignore[attr-defined]
 
-    def and_(self, other) -> SymNode:
+    def and_(self, other) -> "SymNode":
         return self._and_(other)  # type: ignore[attr-defined]
 
-    def or_(self, other) -> SymNode:
+    def or_(self, other) -> "SymNode":
         return self._or_(other)  # type: ignore[attr-defined]
 
-    def float_truediv(self, other) -> SymNode:
+    def float_truediv(self, other) -> "SymNode":
         return self._float_truediv(other)  # type: ignore[attr-defined]
 
-    def int_truediv(self, other) -> SymNode:
+    def int_truediv(self, other) -> "SymNode":
         return self._int_truediv(other)  # type: ignore[attr-defined]
 
-    def int_floordiv(self, other) -> SymNode:
+    def int_floordiv(self, other) -> "SymNode":
         return self._int_floordiv(other)  # type: ignore[attr-defined]
 
-    def lshift(self, other) -> SymNode:
+    def lshift(self, other) -> "SymNode":
         return self._lshift(other)  # type: ignore[attr-defined]
 
-    def rshift(self, other) -> SymNode:
+    def rshift(self, other) -> "SymNode":
         return self._rshift(other)  # type: ignore[attr-defined]
 
-    def sym_not(self) -> SymNode:  # noqa: F811
+    def sym_not(self) -> "SymNode":  # noqa: F811
         return self._sym_not()  # type: ignore[attr-defined]
 
-    def eq(self, other) -> SymNode:
+    def eq(self, other) -> "SymNode":
         return self._eq(other)  # type: ignore[attr-defined]
 
-    def ne(self, other) -> SymNode:
+    def ne(self, other) -> "SymNode":
         return self._ne(other)  # type: ignore[attr-defined]
 
-    def gt(self, other) -> SymNode:
+    def gt(self, other) -> "SymNode":
         return self._gt(other)  # type: ignore[attr-defined]
 
-    def lt(self, other) -> SymNode:
+    def lt(self, other) -> "SymNode":
         return self._lt(other)  # type: ignore[attr-defined]
 
-    def le(self, other) -> SymNode:
+    def le(self, other) -> "SymNode":
         return self._le(other)  # type: ignore[attr-defined]
 
-    def ge(self, other) -> SymNode:
+    def ge(self, other) -> "SymNode":
         return self._ge(other)  # type: ignore[attr-defined]
 
-    def floor(self) -> SymNode:
+    def floor(self) -> "SymNode":
         return self._floor()  # type: ignore[attr-defined]
 
-    def is_integer(self) -> SymNode:
+    def is_integer(self) -> "SymNode":
         return self._is_integer()  # type: ignore[attr-defined]
 
-    def sym_float(self) -> SymNode:  # noqa: F811
+    def sym_float(self) -> "SymNode":  # noqa: F811
         return self._sym_float()  # type: ignore[attr-defined]
 
-    def sym_int(self) -> SymNode:
+    def sym_int(self) -> "SymNode":
         return self._sym_int()  # type: ignore[attr-defined]
 
-    def ceil(self) -> SymNode:
+    def ceil(self) -> "SymNode":
         return self._ceil()  # type: ignore[attr-defined]
 
-    def neg(self) -> SymNode:
+    def neg(self) -> "SymNode":
         return self._neg()  # type: ignore[attr-defined]
 
-    def sym_min(self, other) -> SymNode:  # noqa: F811
+    def sym_min(self, other) -> "SymNode":  # noqa: F811
         return self._sym_min(other)  # type: ignore[attr-defined]
 
-    def sym_max(self, other) -> SymNode:  # noqa: F811
+    def sym_max(self, other) -> "SymNode":  # noqa: F811
         return self._sym_max(other)  # type: ignore[attr-defined]
 
-    def sym_ite(self, then_val, else_val) -> SymNode:
+    def sym_ite(self, then_val, else_val) -> "SymNode":
         return self._sym_ite(then_val, else_val)  # type: ignore[attr-defined]
 
-    def is_contiguous(self, sizes, strides) -> SymNode:
+    def is_contiguous(self, sizes, strides) -> "SymNode":
         return self._is_contiguous(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_channels_last_contiguous_2d(self, sizes, strides) -> SymNode:
+    def is_channels_last_contiguous_2d(self, sizes, strides) -> "SymNode":
         return self._is_channels_last_contiguous_2d(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_channels_last_contiguous_3d(self, sizes, strides) -> SymNode:
+    def is_channels_last_contiguous_3d(self, sizes, strides) -> "SymNode":
         return self._is_channels_last_contiguous_3d(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_channels_last_strides_2d(self, sizes, strides) -> SymNode:
+    def is_channels_last_strides_2d(self, sizes, strides) -> "SymNode":
         return self._is_channels_last_strides_2d(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_channels_last_strides_3d(self, sizes, strides) -> SymNode:
+    def is_channels_last_strides_3d(self, sizes, strides) -> "SymNode":
         return self._is_channels_last_strides_3d(sizes, strides)  # type: ignore[attr-defined]
 
-    def is_non_overlapping_and_dense_indicator(self, sizes, strides) -> SymNode:
+    def is_non_overlapping_and_dense_indicator(self, sizes, strides) -> "SymNode":
         return self._is_non_overlapping_and_dense_indicator(sizes, strides)  # type: ignore[attr-defined]
 
     # Make C++ happy
@@ -413,7 +409,7 @@ def sym_and(self, other):
     def truediv(self, other):
         return self.float_truediv(other)
 
-    def floordiv(self, other) -> SymNode:
+    def floordiv(self, other) -> "SymNode":
         return self.int_floordiv(other)
 
     # We didn't bind integer pow in C++
@@ -633,7 +629,6 @@ def fn(self):
     "asin",
     "acos",
     "atan",
-    "log2",
 )
 for name in math_op_names:
     sym_name = f"sym_{name}"
@@ -661,7 +656,7 @@ def fn(self):
 bool_magic_methods = only_bool_magic_methods | also_bool_magic_methods
 
 # Methods that are only for float
-only_float_magic_methods = {"is_integer", "round", "sym_int", "sym_log2"}
+only_float_magic_methods = {"is_integer", "round", "sym_int"}
 
 
 magic_methods_on_operator_with_trailing_underscore = {"and", "or"}

torch/utils/_sympy/functions.py

@@ -1197,8 +1197,6 @@ def eval(cls, a):
                     a = sympy.oo
                 if a is -int_oo:
                     a = -sympy.oo
-                if name == "log2":
-                    return sympy.log(a, 2)
                 return getattr(sympy, name)(a)
             return None
 
@@ -1223,4 +1221,3 @@ def eval(cls, a):
 OpaqueUnaryFn_exp = make_opaque_unary_fn("exp")
 OpaqueUnaryFn_log = make_opaque_unary_fn("log")
 OpaqueUnaryFn_asinh = make_opaque_unary_fn("asinh")
-OpaqueUnaryFn_log2 = make_opaque_unary_fn("log2")

torch/utils/_sympy/interp.py

@@ -31,7 +31,6 @@
     Min,
     Mod,
     ModularIndexing,
-    OpaqueUnaryFn_log2,
     PowByNatural,
     PythonMod,
     RoundDecimal,
@@ -102,11 +101,7 @@ def handlers():
         Identity: "identity",
         IsNonOverlappingAndDenseIndicator: "is_non_overlapping_and_dense_indicator",
         RoundDecimal: "round_decimal",
-        # TODO: do the rest of the opaque unary functions...
-        OpaqueUnaryFn_log2: "log2",
     }
-    # TODO: This is kind of pointless, we shouldn't be generating sympy.sin
-    # for these functions, they should be Opaque instead
     for name in ["cos", "sin", "tan", "sinh", "cosh", "tanh", "asin", "acos", "atan"]:
         HANDLERS[getattr(sympy, name)] = name
 

torch/utils/_sympy/reference.py

@@ -17,7 +17,6 @@
     Mod,
     OpaqueUnaryFn_exp,
     OpaqueUnaryFn_log,
-    OpaqueUnaryFn_log2,
     OpaqueUnaryFn_sqrt,
     PowByNatural,
     RoundDecimal,
@@ -163,10 +162,6 @@ def exp(x):
     def log(x):
         return OpaqueUnaryFn_log(x)
 
-    @staticmethod
-    def log2(x):
-        return OpaqueUnaryFn_log2(x)
-
     @staticmethod
     def sqrt(x):
         return OpaqueUnaryFn_sqrt(x)
@@ -252,10 +247,6 @@ def exp(x):
     def log(x):
         raise AssertionError("log is not valid shape sympy expr")
 
-    @staticmethod
-    def log2(x):
-        return torch._sym_log2(x)  # type: ignore[attr-defined]
-
     @staticmethod
     def sqrt(x):
         return torch._sym_sqrt(x)  # type: ignore[attr-defined]
@@ -481,10 +472,6 @@ def exp(x):
     def log(x):
         return torch.ops.aten.log.default(x)
 
-    @staticmethod
-    def log2(x):
-        return torch.ops.aten.log2.default(x)
-
     @staticmethod
     def sqrt(x):
         return torch.ops.aten.sqrt.default(x)

torch/utils/_sympy/value_ranges.py

@@ -34,7 +34,6 @@
     IntTrueDiv,
     OpaqueUnaryFn_exp,
     OpaqueUnaryFn_log,
-    OpaqueUnaryFn_log2,
     OpaqueUnaryFn_sqrt,
     PowByNatural,
     RoundDecimal,
@@ -761,13 +760,6 @@ def log(x):
             return ValueRanges.unknown()
         return ValueRanges.increasing_map(x, OpaqueUnaryFn_log)
 
-    @staticmethod
-    def log2(x):
-        x = ValueRanges.wrap(x)
-        if x.lower <= 0:
-            return ValueRanges.unknown()
-        return ValueRanges.increasing_map(x, OpaqueUnaryFn_log2)
-
     @classmethod
     def minimum(cls, a, b):
         return cls.min_or_max(a, b, sympy.Min)