Add out and where args for ht.div

heat/core/arithmetics.py

@@ -8,6 +8,7 @@
 from typing import Optional, Union, Tuple
 
 from . import factories
+from . import indexing
 from . import manipulations
 from . import _operations
 from . import sanitation
@@ -427,7 +428,12 @@ def diff(
     return ret
 
 
-def div(t1: Union[DNDarray, float], t2: Union[DNDarray, float]) -> DNDarray:
+def div(
+    t1: Union[DNDarray, float],
+    t2: Union[DNDarray, float],
+    out: Optional[DNDarray] = None,
+    where: Optional[DNDarray] = None,
+) -> DNDarray:
     """
     Element-wise true division of values of operand ``t1`` by values of operands ``t2`` (i.e ``t1/t2``).
     Operation is not commutative.
@@ -438,6 +444,10 @@ def div(t1: Union[DNDarray, float], t2: Union[DNDarray, float]) -> DNDarray:
         The first operand whose values are divided
     t2: DNDarray or scalar
         The second operand by whose values is divided
+    out: DNDarray, optional
+        The output array. It must have a shape that the inputs broadcast to
+    where: DNDarray, optional
+        Condition of interest, where true yield divided value else yield original value in t1
 
     Example
     ---------
@@ -453,7 +463,10 @@ def div(t1: Union[DNDarray, float], t2: Union[DNDarray, float]) -> DNDarray:
     DNDarray([[2.0000, 1.0000],
               [0.6667, 0.5000]], dtype=ht.float32, device=cpu:0, split=None)
     """
-    return _operations.__binary_op(torch.true_divide, t1, t2)
+    if where is not None:
+        t2 = indexing.where(where, t2, 1)
+
+    return _operations.__binary_op(torch.true_divide, t1, t2, out)
 
 
 DNDarray.__truediv__ = lambda self, other: div(self, other)

heat/core/tests/test_arithmetics.py

@@ -362,6 +362,24 @@ def test_div(self):
         self.assertTrue(ht.equal(ht.div(self.a_tensor, self.an_int_scalar), result))
         self.assertTrue(ht.equal(ht.div(self.a_split_tensor, self.a_tensor), commutated_result))
 
+        a = out = ht.empty((2, 2))
+        ht.div(self.a_tensor, self.a_scalar, out=out)
+        self.assertTrue(ht.equal(out, result))
+        self.assertIs(a, out)
+        b = ht.array([[1.0, 2.0], [3.0, 4.0]])
+        ht.div(b, self.another_tensor, out=b)
+        self.assertTrue(ht.equal(b, result))
+
+        result_where = ht.array([[1.0, 2.0], [1.5, 2.0]])
+        self.assertTrue(
+            ht.equal(ht.div(self.a_tensor, self.a_scalar, where=self.a_tensor > 2), result_where)
+        )
+        ht.div(self.a_tensor, self.a_scalar)
+
+        a = self.a_tensor.copy()
+        ht.div(a, self.a_scalar, out=a, where=a > 2)
+        self.assertTrue(ht.equal(a, result_where))
+
         with self.assertRaises(ValueError):
             ht.div(self.a_tensor, self.another_vector)
         with self.assertRaises(TypeError):