class DemoNet(nn.Layer):
    def __init__(self, np_w0, np_w1):
        super().__init__()
        self.w0 = self.create_parameter(
            shape=[IMAGE_SIZE, IMAGE_SIZE],
            attr=paddle.framework.ParamAttr(
                name="demo_weight_1",
                initializer=paddle.nn.initializer.Assign(np_w0),
            ),
        )
        self.w1 = self.create_parameter(
            shape=[IMAGE_SIZE, CLASS_NUM],
            attr=paddle.framework.ParamAttr(
                name="nemo_weight_2",
                initializer=paddle.nn.initializer.Assign(np_w1),
            ),
        )
    def forward(self, x):
        y = paddle.matmul(x, self.w0)
        z = paddle.matmul(y, self.w1)
        return z
def run_dynamic(self, layer, parallel=False):
    # create loss
    loss_fn = nn.MSELoss()
    # run forward and backward
    image = paddle.to_tensor(self.image)
    out = layer(image)
    label = (
        dist.shard_tensor(
            self.label,
            dist_attr=dist.DistAttr(
                mesh=self._mesh, sharding_specs=[None, None]
            ),
        )
        if parallel is True
        else paddle.to_tensor(self.label)
    )
    loss = loss_fn(out, label)
    loss.backward()

class DPDemoNet(nn.Layer):
    def __init__(self, np_w0, np_w1, mesh):
        super().__init__()
        self.replicate_dist_attr = dist.DistAttr(
            mesh=mesh, sharding_specs=[None, None]
        )
        self.shard_axis0_dist_attr = dist.DistAttr(
            mesh=mesh, sharding_specs=['x', None]
        )
        self.w0 = dist.shard_tensor(
            self.create_parameter(
                shape=[IMAGE_SIZE, IMAGE_SIZE],
                attr=paddle.framework.ParamAttr(
                    name="dp_demo_weight_1",
                    initializer=paddle.nn.initializer.Assign(np_w0),
                ),
            ),
            dist_attr=self.replicate_dist_attr,
        )
        self.w1 = dist.shard_tensor(
            self.create_parameter(
                shape=[IMAGE_SIZE, CLASS_NUM],
                attr=paddle.framework.ParamAttr(
                    name="dp_nemo_weight_2",
                    initializer=paddle.nn.initializer.Assign(np_w1),
                ),
            ),
            dist_attr=self.replicate_dist_attr,
        )
    def forward(self, x):
       # 切分输入x
        y = paddle.matmul(
            dist.shard_tensor(x, dist_attr=self.shard_axis0_dist_attr),
            self.w0,
        )
        z = paddle.matmul(y, self.w1)
        return z

class MPDemoNet(nn.Layer):
    def __init__(self, np_w0, np_w1, mesh):
        super().__init__()
        self.replicate_dist_attr = dist.DistAttr(
            mesh=mesh, sharding_specs=[None, None]
        )
        self.shard_axis0_dist_attr = dist.DistAttr(
            mesh=mesh, sharding_specs=['x', None]
        )
        self.shard_axis1_dist_attr = dist.DistAttr(
            mesh=mesh, sharding_specs=['x', None]
        )
        # 切分参数w0
        self.w0 = dist.shard_tensor(
            self.create_parameter(
                shape=[IMAGE_SIZE, IMAGE_SIZE],
                attr=paddle.framework.ParamAttr(
                    name="mp_demo_weight_1",
                    initializer=paddle.nn.initializer.Assign(np_w0),
                ),
            ),
            dist_attr=self.shard_axis1_dist_attr,
        )
        # 切分参数w1
        self.w1 = dist.shard_tensor(
            self.create_parameter(
                shape=[IMAGE_SIZE, CLASS_NUM],
                attr=paddle.framework.ParamAttr(
                    name="mp_nemo_weight_2",
                    initializer=paddle.nn.initializer.Assign(np_w1),
                ),
            ),
            dist_attr=self.shard_axis0_dist_attr,
        )
    def forward(self, x):
        y = paddle.matmul(
            dist.shard_tensor(x, dist_attr=self.replicate_dist_attr), self.w0
        )
        z = paddle.matmul(y, self.w1)
        return z