| 134 | # This variable controls number of iterations. |
| 135 | loop_vars.append(variables.VariableV1(0.)) |
| 136 | def loop_body(dev0_tensors, dev1_tensors, loop_tensor): |
| 137 | return_ops = [] |
| 138 | for i in range(len(devices)): |
| 139 | device = devices[i] |
| 140 | device_tensors = dev0_tensors if i == 0 else dev1_tensors |
| 141 | with ops.device(device): |
| 142 | device_collectives = [] |
| 143 | for j in range(num_vars): |
| 144 | # NOTE(ayushd): we need the `identity` here to ensure that the |
| 145 | # input to `all_reduce` has an explicit device string. |
| 146 | input_tensor = array_ops.identity(device_tensors[j]) |
| 147 | collective_op = collective_ops.all_reduce( |
| 148 | input_tensor, group_size, group_key, instances[j], |
| 149 | 'Add', 'Id') |
| 150 | device_collectives.append(collective_op) |
| 151 | return_ops.append(device_collectives) |
| 152 | return_ops.append(math_ops.add(loop_tensor, 1.)) |
| 153 | return return_ops |
| 154 | # Run until last variable exceeds number of iterations. |
| 155 | loop_cond = lambda d0, d1, i: math_ops.less(i, num_iterations) |
| 156 | sess.run(variables.global_variables_initializer()) |