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)

Added a unit test for MoE layers with shared experts, and that multiple experts can be saved on a device

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

[pkuzyc]

PR 的介绍改一下，现在的介绍和pr内容不一致

[pkuzyc]

split 可以放循环外面，不用每次循环都 split

[fxfxfxfxfxfxfxfx]

感谢您的代码审查，已经按照意见修改

[pkuzyc]

LGTM

[luotao1]

hi, @fxfxfxfxfxfxfxfx

• 非常感谢你对飞桨的贡献，我们正在运营一个PFCC组织，会通过定期分享技术知识与发布开发者主导任务的形式持续为飞桨做贡献，详情可见 https://github.com/luotao1 主页说明。
• 如果你对PFCC有兴趣，请发送邮件至 ext_paddle_oss@baidu.com，我们会邀请你加入~