[Semi Auto] Revise spmd rule static mode API

[JZ-LIANG]

PR types

Function optimization

PR changes

Others

Description

Pcard-70448

Refactor the static mode python api to distinguish input between tensor and [tensor].

Now the python api signature is exact the same as its C++ implement (which is like the format in Op's Python API).

Before:

# matmul rule
results = rule.infer_forward( [x, y], [True, False] ) # need to wrap tensor and attr into list
results = rule.infer_backward( [x, y], [out], [True, False] )

# concate rule
results = rule.infer_forward([2], [v1, v2], [axis]) # ranges, list of tensors, list of attrs
results = rule.infer_backward([2, 1], [v1, v2, out], [axis])

# replicate rule
results = rule.infer_forward([2, 2], [x, y, out1, out2], [])  # placeholder for attrs
results = rule.infer_backward([2, 2], [x, y, out1, out2], [])

After:

# matmul rule
result_dist_attrs = self.rule.infer_forward( x, y, True, False ) # exactly like matmul Python API
result_dist_attrs = self.rule.infer_backward( x, y, out, True, False)

# concate rule
results = rule.infer_forward( [v1, v2], axis) # exactly like concat Python API
results = rule.infer_backward([v1, v2], out, axis) 

# replicate rule
results = rule.infer_forward([x, y], [out1, out2])
results = rule.infer_backward([x, y], [out1, out2])

TODO：

1. exactly match the Python API (filter the non-relative attributes internally)
2. supporting more ops using codegen.
3. replace dist_tensor_spec as tensor(dist_tensor)

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[paddle-bot]

✅ This PR's description meets the template requirements!
Please wait for other CI results.

[zhiqiu]

LGTM overall

[zhiqiu]

py::args& can be const?

[JZ-LIANG]

it work

[zhiqiu]

I suggest we can support tuple and list both.

[JZ-LIANG]

there is no tuple argument in Phi API yaml right now, we would support it once there is tuple input defined by Phi.

[zhiqiu]

Maybe we can consider unify DistMetaTensor and DistTensorSpec in the future.

[JZ-LIANG]

yes, that is on the list