Support scipy special function with tuple output #3139

RandomY-2 · 2022-06-12T17:48:12Z

What do these changes do?

We can use ExecutableTuple and TensorTupleElementOp to support mars execution of scipy special functions with tuple returns. Specifically, the following syntaxes can be supported:

from mars.tensor import special
import mars.tensor as mt
x = mt.linspace(-3, 3)

S, C = special.fresnel(x)
S.execute()
C.execute()
special.fresnel(x).execute()
(S + C).execute()

Also since we invoke scipy.special.fresnel during tensor executions, a cache is added in core.py so if we execute multiple elements of the same function during one execute() call, we don't need to do the scipy computation for each element. Later executes can just use the cached scipy outputs.

Related issue number

Check code requirements

tests added / passed (if needed)
Ensure all linting tests pass, see here for how to run them

wjsi · 2022-06-13T02:53:25Z

Given your implementation, it is not hard to use one single operand for two outputs. You may take a look at TensorSVD for an example, where we use output_limit property to define the number of outputs:

mars/mars/tensor/linalg/svd.py

Lines 44 to 46 in 424cfb9

    
           @property 
        
           def output_limit(self): 
        
               return 3

use new_tensors(..., kws[kw1, kw2]) to create multiple output objects:

mars/mars/tensor/linalg/svd.py

Lines 69 to 78 in 424cfb9

    
           U, s, V = self.new_tensors( 
        
               [a], 
        
               order=TensorOrder.C_ORDER, 
        
               kws=[ 
        
                   {"side": "U", "dtype": tiny_U.dtype, "shape": U_shape}, 
        
                   {"side": "s", "dtype": tiny_s.dtype, "shape": s_shape}, 
        
                   {"side": "V", "dtype": tiny_V.dtype, "shape": V_shape}, 
        
               ], 
        
           ) 
        
           return ExecutableTuple([U, s, V])

and use ctx[out.key] = xxx to store execute results.

mars/mars/tensor/linalg/svd.py

Lines 150 to 155 in 424cfb9

    
           with device(device_id): 
        
               u, s, v = xp.linalg.svd(a, full_matrices=False) 
        
               uc, sc, vc = op.outputs 
        
               ctx[uc.key] = u 
        
               ctx[sc.key] = s 
        
               ctx[vc.key] = v

RandomY-2 · 2022-06-13T03:19:13Z

Given your implementation, it is not hard to use one single operand for two outputs. You may take a look at TensorSVD for an example, where we use output_limit property to define the number of outputs:

mars/mars/tensor/linalg/svd.py

Lines 44 to 46 in 424cfb9

@property

def output_limit(self):

return 3

use new_tensors(..., kws[kw1, kw2]) to create multiple output objects:

mars/mars/tensor/linalg/svd.py

Lines 69 to 78 in 424cfb9

U, s, V = self.new_tensors(

[a],

order=TensorOrder.C_ORDER,

kws=[

{"side": "U", "dtype": tiny_U.dtype, "shape": U_shape},

{"side": "s", "dtype": tiny_s.dtype, "shape": s_shape},

{"side": "V", "dtype": tiny_V.dtype, "shape": V_shape},

],

)

return ExecutableTuple([U, s, V])

and use ctx[out.key] = xxx to store execute results.

mars/mars/tensor/linalg/svd.py

Lines 150 to 155 in 424cfb9

with device(device_id):

u, s, v = xp.linalg.svd(a, full_matrices=False)

uc, sc, vc = op.outputs

ctx[uc.key] = u

ctx[sc.key] = s

ctx[vc.key] = v

Thank you for your feedback! Will take a look.

wjsi

Code structure LGTM. Do not use for i in range(len(list_obj)). Using for i, obj in enumerate(obj_list) when indices are needed or for obj in obj_list is more pythonic.

wjsi · 2022-06-15T02:17:10Z

mars/tensor/special/core.py

+                for i in range(len(op.outputs)):
+                    ctx[op.outputs[i].key] = ret[i]


Loop over op.outputs itself, not indices.

Suggested change

for i in range(len(op.outputs)):

ctx[op.outputs[i].key] = ret[i]

for output in op.outputs:

ctx[output] = ret[i]

wjsi · 2022-06-15T02:17:41Z

mars/tensor/special/core.py

+            return ExecutableTuple(res_tensors)
+
+        for i in range(len(res_tensors)):
+            out[i].data = res_tensors[i].data


wjsi · 2022-06-15T02:20:23Z

mars/tensor/special/core.py

+                    "dtype": res[i].dtype,
+                    "shape": res[i].shape,
+                }
+                for i in range(len(res))


wjsi · 2022-06-15T02:20:38Z

mars/tensor/special/tests/test_special.py

+    assert isinstance(r, ExecutableTuple)
+    assert len(r) == 2
+
+    for i in range(len(r)):


wjsi · 2022-06-15T02:20:46Z

mars/tensor/special/tests/test_special.py

+    assert len(out) == 2
+    assert len(r_out) == 2
+
+    for i in range(len(r_out)):


wjsi · 2022-06-15T02:20:57Z

mars/tensor/special/tests/test_special_execution.py

+    result = r.execute().fetch()
+    expected = sp_func(raw)
+
+    for i in range(len(result)):


wjsi · 2022-06-15T02:21:23Z

mars/tensor/special/err_fresnel.py

+@_register_special_op
+class TensorFresnel(TensorTupleOp):
+    _func_name = "fresnel"
+    _func_outputs = 2


Rename to _n_outputs might be better.

wjsi

LGTM

hekaisheng

LGTM

…3151) Co-authored-by: RandomY-2 <[email protected]>

RandomY-2 added 3 commits June 10, 2022 10:13

Support scipy functions with tuple outputs (#26)

1e3c925

test for fresnel

2f3c654

update input keys declarations

1e6865a

RandomY-2 requested a review from a team as a code owner June 12, 2022 17:48

RandomY-2 added 4 commits June 12, 2022 10:50

Update err_fresnel.py

2826457

test execution

ea715e8

Update test_special.py

896c63e

Update test_special.py

f2c5eae

RandomY-2 added 4 commits June 14, 2022 14:44

fresnel using one operand (#27)

900eb80

Update test_special.py

f438a89

update imports

a0cea7e

Update test_special.py

16b55c2

wjsi reviewed Jun 15, 2022

View reviewed changes

wjsi added mod: tensor type: feature New feature to be backported Indicate that the PR need to be backported to stable branch labels Jun 15, 2022

wjsi added this to the v0.10.0a2 milestone Jun 15, 2022

RandomY-2 added 2 commits June 14, 2022 20:51

update object enumerating

a37c5eb

Update core.py

723f2ff

wjsi approved these changes Jun 15, 2022

View reviewed changes

hekaisheng approved these changes Jun 16, 2022

View reviewed changes

hekaisheng merged commit a8cb40e into mars-project:master Jun 16, 2022

wjsi pushed a commit to wjsi/mars that referenced this pull request Jun 16, 2022

Support scipy special function with tuple output (mars-project#3139)

64273cd

wjsi mentioned this pull request Jun 16, 2022

[BACKPORT] Support scipy special function with tuple output (#3139) #3151

Merged

wjsi added a commit that referenced this pull request Jun 20, 2022

[BACKPORT] Support scipy special function with tuple output (#3139) (#…

e56fdbb

…3151) Co-authored-by: RandomY-2 <[email protected]>

wjsi added backported already PR has been backported and removed to be backported Indicate that the PR need to be backported to stable branch labels Jun 20, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support scipy special function with tuple output #3139

Support scipy special function with tuple output #3139

RandomY-2 commented Jun 12, 2022 •

edited

Loading

wjsi commented Jun 13, 2022 •

edited

Loading

RandomY-2 commented Jun 13, 2022

wjsi left a comment •

edited

Loading

wjsi Jun 15, 2022

wjsi Jun 15, 2022

wjsi Jun 15, 2022

wjsi Jun 15, 2022

wjsi Jun 15, 2022

wjsi Jun 15, 2022

wjsi Jun 15, 2022

wjsi left a comment

hekaisheng left a comment

		for i in range(len(op.outputs)):
		ctx[op.outputs[i].key] = ret[i]

Support scipy special function with tuple output #3139

Support scipy special function with tuple output #3139

Conversation

RandomY-2 commented Jun 12, 2022 • edited Loading

What do these changes do?

Related issue number

Check code requirements

wjsi commented Jun 13, 2022 • edited Loading

RandomY-2 commented Jun 13, 2022

wjsi left a comment • edited Loading

Choose a reason for hiding this comment

wjsi Jun 15, 2022

Choose a reason for hiding this comment

wjsi Jun 15, 2022

Choose a reason for hiding this comment

wjsi Jun 15, 2022

Choose a reason for hiding this comment

wjsi Jun 15, 2022

Choose a reason for hiding this comment

wjsi Jun 15, 2022

Choose a reason for hiding this comment

wjsi Jun 15, 2022

Choose a reason for hiding this comment

wjsi Jun 15, 2022

Choose a reason for hiding this comment

wjsi left a comment

Choose a reason for hiding this comment

hekaisheng left a comment

Choose a reason for hiding this comment

RandomY-2 commented Jun 12, 2022 •

edited

Loading

wjsi commented Jun 13, 2022 •

edited

Loading

wjsi left a comment •

edited

Loading