(self, grad_packs)
| 436 | return grad_packs |
| 437 | |
| 438 | def unpack(self, grad_packs): |
| 439 | with cached_name_scope("GradientPacker", top_level=False): |
| 440 | concat_grads = tf.concat(grad_packs, 0, name='concatenated_packs') |
| 441 | # concat_grads = tf.cast(concat_grads, self._grad_dtype) |
| 442 | flattened_grads = tf.split(concat_grads, self._sizes) |
| 443 | grads = [tf.reshape(g, shape) for g, shape in zip(flattened_grads, self._shapes)] |
| 444 | return grads |
| 445 | |
| 446 | def pack_all(self, all_grads, devices): |
| 447 | """ |
no test coverage detected