def test_auto_parallel(): context.set_context(save_graphs=True) context.set_auto_parallel_context(parallel_mode="auto_parallel", device_num=16, global_rank=0) net = Full(_w1, 3) net.set_auto_parallel() net.set_train() _executor.compile(net, _x, phase='train') num_ops = _executor._get_num_parallel_ops(net) expected_num = 16 assert num_ops == expected_num
def test_double_subgraphs(): context.set_context(save_graphs=False) context.set_auto_parallel_context(parallel_mode="auto_parallel", device_num=8, global_rank=0) net = TrainStepWarp(NetWithLoss(Net())) _set_multi_subgraphs() net.set_auto_parallel() x = Tensor(np.ones([8, 8, 8, 8]), dtype=ms.float32) reset_op_id() net.set_train() _executor.compile(net, x, phase='train') num_ops = _executor._get_num_parallel_ops(net) expected_num = 7 assert expected_num == num_ops