Enhance symbolic arithmetics

thunder/clang/__init__.py

@@ -1370,7 +1370,7 @@ def round(a: TensorLike | Number) -> TensorLike | Number:
     return _elementwise_unary_wrapper(
         a,
         prim=prims.round,
-        type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
+        type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.NUMBER_TO_INT,
     )
 
 
@@ -1453,7 +1453,7 @@ def tanh(a):
     )
 
 
-@clangop()
+@clangop(method_name="trunc")
 def trunc(a: TensorLike | Number) -> TensorLike | Number:
     # Short-circuits on unsigned inputs (which are already trivially truncated)
     if dtypes.is_exact_dtype(dtypes.to_dtype(a)):
@@ -1462,7 +1462,7 @@ def trunc(a: TensorLike | Number) -> TensorLike | Number:
     return _elementwise_unary_wrapper(
         a,
         prim=prims.trunc,
-        type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
+        type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.NUMBER_TO_INT,
     )
 
 
@@ -1793,7 +1793,7 @@ def zeta(a, b):
     )
 
 
-@clangop()
+@clangop(method_name="bitwise_left_shift")
 def bitwise_left_shift(a, b):
     return _elementwise_binary_wrapper(
         a,
@@ -1803,7 +1803,7 @@ def bitwise_left_shift(a, b):
     )
 
 
-@clangop()
+@clangop(method_name="bitwise_right_shift")
 def bitwise_right_shift(a, b):
     return _elementwise_binary_wrapper(
         a,

thunder/core/interpreter.py

@@ -398,7 +398,7 @@ def __init__(
         self._with_provenance_tracking = with_provenance_tracking
         if with_provenance_tracking:
             assert isinstance(uncacheable_classes, (list, tuple))
-            uncacheable_classes = tuple(set(uncacheable_classes) | {NoneType, int, str, float, bool})
+            uncacheable_classes = tuple(set(uncacheable_classes) | {NoneType, int, str, float, bool, complex})
 
         self._uncacheable_classes = uncacheable_classes
 

thunder/core/jit_ext.py

@@ -312,6 +312,7 @@ def proxify(self, value: WrappedValue) -> Any:
                     self.add_constraint((clang.check_number_type_and_value, p, uvalue))
             elif co is CACHE_OPTIONS.SYMBOLIC_VALUES:
                 if p is not uvalue:
+                    self.add_constraint((clang.check_instance, p, (type(uvalue),)))
                     value.register_proxy(p)
             elif co not in (CACHE_OPTIONS.SAME_INPUT, CACHE_OPTIONS.NO_CACHING):
                 raise NotImplementedError(f"Unsupported cache option {co}")

thunder/core/prims.py

@@ -1945,9 +1945,12 @@ def _numpy_array_to_torch_tensor_meta(a: TensorProxy, /) -> TensorProxy:
 #   usually produce an output with that same datatype (SAME).
 #   Sometimes, however, elementwise operations can produce an output with a different
 #   datatype than the inputs. For example, comparison operations like eq and lt always
-#   produce boolean results (ALWAYS_BOOL), math.ceil/math.floor produces integer outputs for number inputs while preserves datatype for tensor inputs, and other operations, like abs, map
-#   complex numbers to floats (COMPLEX_TO_FLOAT).
+#   produce boolean results (ALWAYS_BOOL), ceil/floor produces integer outputs for
+#   number inputs while preserves datatype for tensor inputs (INT_FOR_NUMBER), and
+#   other operations, like abs, map complex numbers to floats (COMPLEX_TO_FLOAT).
 #   The ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND enum describes these three behaviors so that
+#   operations, like abs, map complex numbers to floats (COMPLEX_TO_FLOAT).
+#   The ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND enum describes these four behaviors so that
 #   elementwise operations can rely on helper functions to implement this behavior.
 class ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND(Enum):
     SAME = auto()
@@ -2316,6 +2319,7 @@ def frexp_meta(a: TensorProxy, /) -> (TensorProxy, TensorProxy):
     "round",
     number_fn=builtins.round,
     supported_input_dtypes=fp_math_dtypes,
+    output_dtype_kind=ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND.INT_FOR_NUMBER,
 )
 
 rsqrt = _make_elementwise_unary_prim(
@@ -2385,12 +2389,12 @@ def _signbit_number(a: Number) -> bool:
     supported_input_dtypes=fp_math_dtypes,
 )
 
-# NOTE This trunc preserves the dtype of its input
 trunc = _make_elementwise_unary_prim(
     PrimIDs.TRUNC,
     "trunc",
     supported_input_dtypes=fp_math_dtypes,
     number_fn=math.trunc,
+    output_dtype_kind=ELEMENTWISE_PRIM_OUTPUT_DTYPE_KIND.INT_FOR_NUMBER,
 )
 
 
@@ -2804,21 +2808,13 @@ def _lerp_meta(start: TensorProxy, end: TensorProxy, weight: Number | TensorProx
 )
 
 
-# TODO Restore Number x Number x Number support
 def _where_meta(pred: Number | TensorProxy, a: Number | TensorProxy, b: Number | TensorProxy, /) -> TensorProxy:
     # Checks types
     # NOTE pred must be a bool tensor or bool (this is checked later)
     utils.check_type(pred, (TensorProxy, Number, NumberProxy))
     utils.check_type(a, (TensorProxy, Number, NumberProxy))
     utils.check_type(b, (TensorProxy, Number, NumberProxy))
 
-    if (
-        isinstance(pred, (Number, NumberProxy))
-        and isinstance(a, (Number, NumberProxy))
-        and isinstance(b, (Number, NumberProxy))
-    ):
-        raise NotImplementedError
-
     # Checks pred dtype (bool or bool tensor)
     if isinstance(pred, (Number, NumberProxy)):
         utils.check(
@@ -2843,11 +2839,19 @@ def _where_meta(pred: Number | TensorProxy, a: Number | TensorProxy, b: Number |
     numbertype, tensordtype = utils.check_same_dtype(a, b)
     dtype = tensordtype if tensordtype is not None else numbertype
 
+    # Returns a NumberProxy for all-Number inputs
+    if (
+        isinstance(pred, (Number, NumberProxy))
+        and isinstance(a, (Number, NumberProxy))
+        and isinstance(b, (Number, NumberProxy))
+    ):
+        result_value = pyval(a) if pyval(pred) else pyval(b)
+        return numberproxy(numbertype, result_value, constraint=utils.resolve_constraints(pred, a, b))
+
     # Checks shapes
     utils.check_same_shape(pred, a, b)
 
     # Determines output shape
-    # NOTE Assumes at least one of pred, a, and b is a TensorProxy because of prior check for Number x Number x Number
     shapes = tuple(x.shape for x in (pred, a, b) if isinstance(x, TensorProxy) and not utils.is_cpu_scalar_tensor(x))
     if not shapes:
         shapes = (pred.shape,)

thunder/core/proxies.py

@@ -30,7 +30,7 @@
     TagBase,
 )
 import thunder.core.baseutils as baseutils
-from thunder.core.langctxs import resolve_method, get_langctx
+from thunder.core.langctxs import LanguageContext, resolve_method, get_langctx
 import thunder.core.devices as devices
 import thunder.core.dtypes as dtypes
 
@@ -739,7 +739,7 @@ def _elementwise_unary_helper(a, name, fn, type_promotion_kind=None):
         vala = pyval(a)
 
         trace: None | TraceCtx = get_tracectx()
-        lang: None | LangCtx = None
+        lang: None | LanguageContext = None
         try:
             lang = get_langctx()
         except LookupError:
@@ -775,7 +775,7 @@ def __neg__(self):
         return self._elementwise_unary_helper(self, "neg", operator.neg)
 
     def __pos__(self):
-        return self._elementwise_unary_helper(self, "pos", operator.pos)
+        return self
 
     # See https://docs.python.org/3/reference/datamodel.html#object.__round__
     def __round__(self):
@@ -797,7 +797,7 @@ def _elementwise_binary_helper(a, b, name, fn, type_promotion_kind=None):
         valb = pyval(b) if isinstance(b, NumberProxy) else b
 
         trace: None | TraceCtx = get_tracectx()
-        lang: None | LangCtx = None
+        lang: None | LanguageContext = None
         try:
             lang = get_langctx()
         except LookupError:
@@ -954,16 +954,16 @@ def __rxor__(self, other):
     #   tracks implementing these
 
     def __lshift__(self, other):
-        raise NotImplementedError
+        return self._elementwise_binary_helper(self, other, "bitwise_left_shift", operator.lshift)
 
     def __rlshift__(self, other):
-        raise NotImplementedError
+        return self._elementwise_binary_helper(other, self, "bitwise_left_shift", operator.lshift)
 
     def __rshift__(self, other):
-        raise NotImplementedError
+        return self._elementwise_binary_helper(self, other, "bitwise_right_shift", operator.rshift)
 
     def __rrshift__(self, other):
-        raise NotImplementedError
+        return self._elementwise_binary_helper(other, self, "bitwise_right_shift", operator.rshift)
 
     #
     # Casts to Python numbers
@@ -1676,8 +1676,7 @@ def __neg__(self):
         return method(self)
 
     def __pos__(self):
-        method = resolve_method("pos", self)
-        return method(self)
+        return self
 
     def __round__(self):
         method = resolve_method("round", self)

thunder/core/utils.py

@@ -458,6 +458,9 @@ def elementwise_type_promotion(*args, type_promotion_kind: ELEMENTWISE_TYPE_PROM
 
     ALWAYS_BOOL is like PRESERVE, except the result dtype is always bool.
 
+    NUMBER_TO_INT is like DEFAULT, except float promotion dtypes *with no tensor inputs* use int
+    for their result dtypes. This absorbs the difference between e.g. math.ceil and torch.ceil.
+
     Example operators for each type promotion option:
 
       DEFAULT                 : add
@@ -504,7 +507,7 @@ def elementwise_type_promotion(*args, type_promotion_kind: ELEMENTWISE_TYPE_PROM
         and is_float_dtype(promotiontype)
         and all_number_type
     ):
-        return int, int
+        return promotiontype, int
 
     # Falls through to DEFAULT
     if is_low_precision_dtype(promotiontype):

thunder/executors/pythonex.py

@@ -257,6 +257,10 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None:
         a.clear()
 
 
+ceil = ex.register_operator("ceil", like=prims.ceil, module=math)
+floor = ex.register_operator("floor", like=prims.floor, module=math)
+trunc = ex.register_operator("trunc", like=prims.trunc, module=math)
+py_round = ex.register_operator("round", like=prims.round, module=builtins)
 acos = ex.register_operator("acos", like=prims.acos, module=math)
 acosh = ex.register_operator("acosh", like=prims.acosh, module=math)
 asin = ex.register_operator("asin", like=prims.asin, module=math)
@@ -273,6 +277,10 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None:
     "clear_mutable_collection", meta=_clear_mutable_collection_meta, fn=_clear_mutable_collection_prim_impl
 )
 
+ex.register_implementation(prims.ceil, ceil, checker=_elementwise_unary_checker)
+ex.register_implementation(prims.floor, floor, checker=_elementwise_unary_checker)
+ex.register_implementation(prims.trunc, trunc, checker=_elementwise_unary_checker)
+ex.register_implementation(prims.round, py_round, checker=_elementwise_unary_checker)
 ex.register_implementation(prims.acos, acos, checker=_elementwise_unary_checker)
 ex.register_implementation(prims.acosh, acosh, checker=_elementwise_unary_checker)
 ex.register_implementation(prims.asin, asin, checker=_elementwise_unary_checker)
@@ -287,8 +295,6 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None:
 ex.register_implementation(prims.signbit, signbit, checker=_elementwise_unary_checker)
 
 
-# # bitwise_not = _elementwise_unary_factory("invert", operator)
-# # ceil = _elementwise_unary_factory("ceil", math)
 # # cos = _elementwise_unary_factory("cos", math)
 # # cosh = _elementwise_unary_factory("cosh", math)
 # # erf = _elementwise_unary_factory("erf", math)
@@ -298,7 +304,6 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None:
 # # exp = _elementwise_unary_factory("exp", math)
 # # exp2 = None
 # # expm1 = _elementwise_unary_factory("expm1", math)
-# # floor = _elementwise_unary_factory("floor", math)
 # # isfinite = _elementwise_unary_factory("isfinite", cmath)
 # # lgamma = _elementwise_unary_factory("lgamma", math)
 # # log = _elementwise_unary_factory("log", math)
@@ -307,16 +312,13 @@ def _clear_mutable_collection_prim_impl(a: Collection) -> None:
 # # log2 = _elementwise_unary_factory("log2", math)
 # # ndtri = None
 # # reciprocal = None
-# # # NOTE pythonex_round to avoid a name conflict with the builtin round
-# # pythonex_round = _elementwise_unary_factory("round", builtins)
 # # rsqrt = None
 # # sign = None
 # # sin = _elementwise_unary_factory("sin", math)
 # # sinh = _elementwise_unary_factory("sinh", math)
 # # sqrt = _elementwise_unary_factory("sqrt", math)
 # # tan = _elementwise_unary_factory("tan", math)
 # # tanh = _elementwise_unary_factory("tanh", math)
-# # trunc = _elementwise_unary_factory("trunc", math)
 
 #
 # Elementwise binary primitives
@@ -332,6 +334,9 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten
 bitwise_and = ex.register_operator("bitwise_and", like=prims.bitwise_and, fn=operator.and_)
 bitwise_or = ex.register_operator("bitwise_or", like=prims.bitwise_or, fn=operator.or_)
 bitwise_xor = ex.register_operator("bitwise_xor", like=prims.bitwise_xor, fn=operator.xor)
+bitwise_not = ex.register_operator("bitwise_not", like=prims.bitwise_not, fn=operator.inv)
+bitwise_left_shift = ex.register_operator("bitwise_left_shift", like=prims.bitwise_left_shift, fn=operator.lshift)
+bitwise_right_shift = ex.register_operator("bitwise_right_shift", like=prims.bitwise_right_shift, fn=operator.rshift)
 eq = ex.register_operator("eq", like=prims.eq, module=operator)
 py_floordiv = ex.register_operator("floordiv", like=prims.py_floordiv, module=operator)
 fmod = ex.register_operator("fmod", like=prims.fmod, module=math)
@@ -356,6 +361,9 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten
 ex.register_implementation(prims.bitwise_and, bitwise_and, checker=_elementwise_binary_checker)
 ex.register_implementation(prims.bitwise_or, bitwise_or, checker=_elementwise_binary_checker)
 ex.register_implementation(prims.bitwise_xor, bitwise_xor, checker=_elementwise_binary_checker)
+ex.register_implementation(prims.bitwise_not, bitwise_not, checker=_elementwise_unary_checker)
+ex.register_implementation(prims.bitwise_left_shift, bitwise_left_shift, checker=_elementwise_binary_checker)
+ex.register_implementation(prims.bitwise_right_shift, bitwise_right_shift, checker=_elementwise_binary_checker)
 ex.register_implementation(prims.eq, eq, checker=_elementwise_binary_checker)
 ex.register_implementation(prims.py_floordiv, py_floordiv, checker=_elementwise_binary_checker)
 ex.register_implementation(prims.fmod, fmod, checker=_elementwise_binary_checker)
@@ -373,6 +381,24 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten
 ex.register_implementation(prims.shape, shape, checker=_always_executable)
 
 
+def _elementwise_ternary_checker(
+    a: NumberLike | TensorProxy, b: NumberLike | TensorProxy, c: NumberLike | TensorProxy
+) -> bool:
+    return (
+        isinstance(a, (Number, NumberProxy))
+        and isinstance(b, (Number, NumberProxy))
+        and isinstance(c, (Number, NumberProxy))
+    )
+
+
+def _where_prim_impl(pred, a, b):
+    return a if pred else b
+
+
+where = ex.register_operator("where", like=prims.where, fn=_where_prim_impl)
+ex.register_implementation(prims.where, where, checker=_elementwise_ternary_checker)
+
+
 def _sink(*args, **kwargs):
     return