add nonzero in ops_infer_shape_in_runtime

[fxfxfxfxfxfxfxfx]

PR Category

Auto Parallel

PR Types

Bug fixes

Description

In a distributed scenario, the output shape of the nonzero method is incorrect (the first dimension is -1), which is abnormal. This issue arises because the shape of nonzero is generated after the computation, but the framework does not place it into ops_infer_shape_in_runtime. The following code can reproduce the corresponding problem.

import paddle
import paddle.distributed as dist
x = paddle.to_tensor([1, 0, 1, 1])
mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
x = dist.shard_tensor(x,mesh,[dist.Replicate()])
out = paddle.nonzero(x)
print(out.shape)

[paddle-bot[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.

[CLAassistant]

All committers have signed the CLA.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}