Revert div_exact primitive and fix div grad to use prims.div

Co-authored-by: IvanYashchuk <19621411+IvanYashchuk@users.noreply.github.com>

Author: Copilot

thunder/clang/__init__.py

@@ -1601,12 +1601,6 @@ def eq(a, b):
 # NOTE This is distinct from true_divide, which also wraps prims.div, because it doesn't promote
 #   integers to floating point values
 def _c_div(a: TensorProxy | Number, b: TensorProxy | Number) -> TensorProxy | Number:
-    # Uses non-differentiable version of div for exact types
-    if utils.is_exact_dtype(utils.to_dtype(a)) and utils.is_exact_dtype(utils.to_dtype(b)):
-        return _elementwise_binary_wrapper(
-            a, b, prim=prims.div_exact, type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
-        )
-    # OTOH trunc_div, which uses _c_div, is differentiable for inexact types
     return _elementwise_binary_wrapper(
         a, b, prim=prims.div, type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
     )

thunder/core/prims.py

@@ -223,7 +223,6 @@ class PrimIDs(Enum):
     BITWISE_OR = auto()
     BITWISE_XOR = auto()
     DIV = auto()
-    DIV_EXACT = auto()
     EQ = auto()
     PY_FLOORDIV = auto()
     FMOD = auto()
@@ -2619,14 +2618,6 @@ def _div_numbers(a: Number, b: Number) -> Number:
     supported_input_dtypes=math_dtypes,
 )
 
-# The non-differentiable version of div
-div_exact = _make_elementwise_binary_prim(
-    PrimIDs.DIV_EXACT,
-    "div_exact",
-    number_fn=_div_numbers,
-    supported_input_dtypes=dtypes.exact_dtypes,
-)
-
 # We currently do not support floordiv on tensors.
 py_floordiv = _make_elementwise_binary_prim(
     PrimIDs.PY_FLOORDIV,

thunder/core/transforms.py

@@ -1135,14 +1135,12 @@ def _add_prim_grad(a: Number | TensorProxy, b: Number | TensorProxy, /) -> Numbe
 register_grad(pids.ADD, _add_prim_grad)
 
 
-# NOTE The following grad definition relies on the fact that only inexact dtypes are differentiable,
-#   and torch's true division operator and the division primitive agree on those types
 def _div_prim_grad(a: Number | TensorProxy, b: Number | TensorProxy, /) -> Number | TensorProxy:
-    fwd = a / b
+    fwd = prims.div(a, b)
 
     g = get_grad(fwd)
-    a_grad = g / b
-    b_grad = -g * ((a / b) / b)
+    a_grad = prims.div(g, b)
+    b_grad = -g * prims.div(prims.div(a, b), b)
     put_grads((a, b), (a_grad, b_grad))
 
     return fwd
@@ -2590,7 +2588,6 @@ def uniform_backward(primal, minval, maxval, g):
     prims.PrimIDs.FULL,
     prims.PrimIDs.FLOOR,
     prims.PrimIDs.CEIL,
-    prims.PrimIDs.DIV_EXACT,
 }
 
 

thunder/executors/nvfuserex_impl.py

@@ -1861,7 +1861,6 @@ def div(a: TensorProxy | Number, b: TensorProxy | Number, *, fd: FusionDefinitio
 
 
 register_supported(PrimIDs.DIV, div, _elementwise_binary_check)
-register_supported(PrimIDs.DIV_EXACT, div, _elementwise_binary_check)
 
 
 def eq(a: TensorProxy | Number, b: TensorProxy | Number, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any:

thunder/executors/pythonex.py

@@ -346,7 +346,6 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten
 pythonex_pow = ex.register_operator("pow", like=prims.pow, module=operator)
 sub = ex.register_operator("sub", like=prims.sub, module=operator)
 div = ex.register_operator("div", like=prims.div, fn=_div_prim_impl)
-div_exact = ex.register_operator("div_exact", like=prims.div_exact, fn=_div_prim_impl)
 shape = ex.register_operator("shape", like=prims.shape, fn=lambda x: x.shape)
 
 # TODO: Restore truediv once we find it...
@@ -371,7 +370,6 @@ def _elementwise_binary_checker(a: NumberLike | TensorProxy, b: NumberLike | Ten
 ex.register_implementation(prims.pow, pythonex_pow, checker=_elementwise_binary_checker)
 ex.register_implementation(prims.sub, sub, checker=_elementwise_binary_checker)
 ex.register_implementation(prims.div, div, checker=_elementwise_binary_checker)
-ex.register_implementation(prims.div_exact, div_exact, checker=_elementwise_binary_checker)
 ex.register_implementation(prims.shape, shape, checker=_always_executable)
 
 

thunder/executors/torchex.py

@@ -1024,7 +1024,6 @@ def _div_transform(
 _register_elementwise_binary_implementation(prims.bitwise_xor, bitwise_xor)
 div_prim_impl = ex.register_operator("torch_div_prim_impl", meta=prims.div.meta, fn=_div_prim_impl)
 _register_elementwise_binary_implementation(prims.div, div_prim_impl)
-_register_elementwise_binary_implementation(prims.div_exact, div_prim_impl)
 _register_elementwise_binary_implementation(prims.eq, eq)
 _register_elementwise_binary_implementation(prims.fmod, fmod)
 _register_elementwise_binary_implementation(prims.ge, ge)

thunder/tests/test_ops.py

@@ -528,6 +528,7 @@ def fn(a):
 
 
 def test_div_exact():
+    # Test that division with integer inputs and requires_grad tensor output works correctly
     def fn(a, b, c):
         indices = torch.div(a, b, rounding_mode="trunc")
         # this would throw an error if indices are not ints
@@ -536,9 +537,13 @@ def fn(a, b, c):
     jfn = thunder.jit(fn)
     a = torch.randint(1, 5, (5,))
     b = torch.ones(5, dtype=torch.int32)
-    c = torch.randn(5, 5)
-    assert_close(fn(a, b, c), jfn(a, b, c))
+    c = torch.randn(5, 5, requires_grad=True)
+    result_eager = fn(a, b, c)
+    result_jit = jfn(a, b, c)
+    assert_close(result_eager, result_jit)
+    # Ensure the division primitive is used (no div_exact)
     trc = thunder.last_traces(jfn)[-1]
     for bsym in trc.bound_symbols:
         if bsym.sym.id == "div":
-            assert "div_exact" in [ssym.sym.name for ssym in bsym.subsymbols[0].subsymbols]
+            # Verify that prims.div is used, not a separate div_exact primitive
+            assert "div" in [ssym.sym.name for ssym in bsym.subsymbols[0].subsymbols]