[Semi-Auto] Add reshape infer_backward rule #56443
Conversation
|
你的PR提交成功，感谢你对开源项目的贡献! |
|
Sorry to inform you that 6e45db7's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually. |
pkuzyc
force-pushed
the
reshape_backward
branch
from
6e45db7
to
367f522
Compare
pkuzyc
force-pushed
the
reshape_backward
branch
from
367f522
to
56077ad
Compare
| @@ -248,6 +248,171 @@ def test_reshape_infer_forward(self): | |||
| with self.assertRaises(BaseException): | |||
| self.rule.infer_forward([self.x_dist_tensor_spec], self.attrs) | |||
|
|
|||
| def test_reshape_infer_backward(self): | |||
There was a problem hiding this comment.
unitest lack some cases where the axis of output is shard and need to be recover to unsharded before reshard：
case1: last axis of output
shape: [6, 12, 48, 24] --> [1, 72, 48, 4, 6] (input --> output)
dims_mapping: [-1, 0, -1, -1, 1] --> [0, -1, -1, -1], [-1, 0, 1, -1, -1] (output --> input, output)
case2: 2nd axis of output
shape: [6, 12, 48, 24] --> [3, 24, 6, 8, 24] (input --> output)
dims_mapping: [-1, 1, -1, -1, 0] --> [-1, -1, -1, 0], [-1, -1, -1, -1, 0] (output --> input, output)
PR types
New features
PR changes
Others
Description
Pcard-70448
Add infer backward SPMD rule for reshape. It infers inputs' dims mappings as following: