def __init__(self, group, transpose_a=False, transpose_b=False):
super(MatmulAllgather, self).__init__()
self.allgather = P.AllGather(group=group)
self.matmul = P.MatMul(transpose_a, transpose_b)
self.pow = P.Pow()
self.reduce_sum = P.ReduceSum()
self.allreduce = P.AllReduce(group=group)
def __init__(self):
super(Net, self).__init__()
self.all_gather = P.AllGather(group=NCCL_WORLD_COMM_GROUP)
self.x = Parameter(initializer(Tensor(x), x.shape), name='x')