Add sym_log2 (pytorch#137980)

Internal xref: https://fb.workplace.com/groups/1075192433118967/permalink/1515595595745313/

Signed-off-by: Edward Z. Yang <[email protected]>

Pull Request resolved: pytorch#137980
Approved by: https://github.com/bobrenjc93

Author: ezyang

test/test_dynamic_shapes.py

@@ -499,14 +499,24 @@ 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(s0), 2.0)"""
+            str(shape_env.guards[0][0]), """Eq(OpaqueUnaryFn_sqrt(ToFloat(s0)), 2.0)"""
         )
 
     def test_sym_floor(self):
@@ -540,7 +550,8 @@ 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(s0)), 2)"""
+            str(shape_env.guards[1][0]),
+            """Eq(TruncToInt(OpaqueUnaryFn_sqrt(ToFloat(s0))), 2)""",
         )
 
     def test_sym_ceil(self):

torch/__init__.py

@@ -876,6 +876,8 @@ 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)
@@ -895,6 +897,7 @@ 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,4 +1,8 @@
 # mypy: allow-untyped-defs
+
+from __future__ import annotations
+
+
 """
 This file does three things:
 - Contains the definition of SymNode
@@ -145,12 +149,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
@@ -281,121 +285,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
@@ -409,7 +413,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++
@@ -426,7 +430,7 @@ def int_(self):
     # functions consider factoring it out to be metaprogrammed too.  Note that
     # some load bearing logic is directly in torch.sym_sum
 
-    def sym_sum(self, args) -> "SymNode":
+    def sym_sum(self, args) -> SymNode:
         import sympy
 
         # Inner impl
@@ -629,6 +633,7 @@ def fn(self):
     "asin",
     "acos",
     "atan",
+    "log2",
 )
 for name in math_op_names:
     sym_name = f"sym_{name}"
@@ -656,7 +661,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"}
+only_float_magic_methods = {"is_integer", "round", "sym_int", "sym_log2"}
 
 
 magic_methods_on_operator_with_trailing_underscore = {"and", "or"}

torch/utils/_sympy/functions.py

@@ -1197,6 +1197,8 @@ 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
 
@@ -1221,3 +1223,4 @@ 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,6 +31,7 @@
     Min,
     Mod,
     ModularIndexing,
+    OpaqueUnaryFn_log2,
     PowByNatural,
     PythonMod,
     RoundDecimal,
@@ -101,7 +102,11 @@ 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,6 +17,7 @@
     Mod,
     OpaqueUnaryFn_exp,
     OpaqueUnaryFn_log,
+    OpaqueUnaryFn_log2,
     OpaqueUnaryFn_sqrt,
     PowByNatural,
     RoundDecimal,
@@ -162,6 +163,10 @@ 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)
@@ -247,6 +252,10 @@ 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]
@@ -472,6 +481,10 @@ 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,6 +34,7 @@
     IntTrueDiv,
     OpaqueUnaryFn_exp,
     OpaqueUnaryFn_log,
+    OpaqueUnaryFn_log2,
     OpaqueUnaryFn_sqrt,
     PowByNatural,
     RoundDecimal,
@@ -760,6 +761,13 @@ 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)