def broadcast_params(self, optim_result):
"""
Apply Broadcast operations in the sequential order of parameter groups.
Returns:
bool, the status flag.
"""
param_group = []
key_group = []
for _ in range(self.dev_num):
param_group.append(F.make_tuple())
key_group.append(F.make_tuple())
for i in range(self.param_length):
param_group[self.param_rank[i]] = param_group[
self.param_rank[i]] + (self.parameters[i], )
key = P.MakeRefKey(self.param_names[i])()
key_group[
self.param_rank[i]] = key_group[self.param_rank[i]] + (key, )
new_param_group = []
for root in range(self.dev_num):
ops = P.Broadcast(root)
next_params = ops(param_group[root])
new_param_group.append(next_params)
for i in range(F.tuple_len(next_params)):
F.assign(key_group[root][i], next_params[i])
status = F.control_depend(optim_result, new_param_group[0][0])
for i in range(self.dev_num - 1):
status = F.depend(
F.control_depend(new_param_group[i],
new_param_group[i + 1][0]), status)
return status
def construct(self, x):
key = P.MakeRefKey("y")()
P.Assign()(key, x)
return x
