Fix tensor frexp (#3259)

* Fix tensor frexp

* Test for Ray DAG

* Fix

Co-authored-by: 刘宝 <po.lb@antgroup.com>

mars/tensor/arithmetic/frexp.py

@@ -57,17 +57,10 @@ def execute(cls, ctx, op):
                 where = None
             kw["order"] = op.order
 
-            try:
-                args = [input]
-                if out1 is not None:
-                    args.append(out1)
-                if out2 is not None:
-                    args.append(out2)
-                mantissa, exponent = xp.frexp(*args, **kw)
-            except TypeError:
-                if where is None:
-                    raise
-                mantissa, exponent = xp.frexp(input)
+            # The out1 out2 are immutable because they are got from
+            # the shared memory.
+            mantissa, exponent = xp.frexp(input)
+            if where is not None:
                 mantissa, exponent = (
                     xp.where(where, mantissa, out1),
                     xp.where(where, exponent, out2),

mars/tensor/arithmetic/tests/test_arithmetic_execution.py

@@ -25,7 +25,7 @@
 from ....session import execute, fetch
 from ....tests.core import require_cupy
 from ....utils import ignore_warning
-from ...datasource import ones, tensor, zeros
+from ...datasource import ones, tensor, zeros, arange
 from .. import (
     add,
     cos,
@@ -349,6 +349,7 @@ def test_arctan2_execution(setup):
     np.testing.assert_equal(result, np.arctan2(0, raw2.A))
 
 
+@pytest.mark.ray_dag
 def test_frexp_execution(setup):
     data1 = np.random.RandomState(0).randint(0, 100, (5, 9, 6))
 
@@ -381,6 +382,17 @@ def test_frexp_execution(setup):
     expected = sum(np.frexp(data1.toarray()))
     np.testing.assert_equal(res.toarray(), expected)
 
+    x = np.arange(9)
+    a = np.zeros(9)
+    b = np.zeros(9)
+    mx = arange(9)
+    ma = zeros(9)
+    mb = zeros(9)
+    res = frexp(mx, ma, mb, where=mx > 5).execute()
+    expected = np.frexp(x, a, b, where=x > 5)
+    np.testing.assert_equal(res[0], expected[0])
+    np.testing.assert_equal(res[1], expected[1])
+
 
 def test_frexp_order_execution(setup):
     data1 = np.random.RandomState(0).random((5, 9))