parameters: (self, grouped_grads_and_vars)
return_statement: return device_grad_packs
docstring: Pack tensors.
docstring_summary: Pack tensors.
docstring_tokens: Pack tensors .
function: def pack(self, grouped_grads_and_vars):
      """Pack tensors."""
      self.grouped_grads_and_vars = grouped_grads_and_vars
      self.all_device_shapes = []
      self.all_device_sizes = []
  
      device_grad_packs = []
      for device_grads_and_vars in grouped_grads_and_vars:
        with ops.colocate_with(device_grads_and_vars[0][0]):
          # Flatten all the grads.
          flat_grads = [
              array_ops.reshape(g, [-1]) for g, _ in device_grads_and_vars
          ]
          # Remember the original shape of all the grads.
          device_shapes = [array_ops.shape(g) for g, _ in device_grads_and_vars]
          # Remember the original sizes of all the grads.
          device_sizes = [array_ops.size(g) for g, _ in device_grads_and_vars]
          # Concat all the flat grads into a big flat tensor.
          concat_grads = array_ops.concat(flat_grads, 0)
  
          # Split the big tensor into num_splits packs. In cases where the
          # total size is not divisible num_splits, the last pack gets
          # more elements.
          # TODO(zhengxq): it is also possible to optimize away all the concat
          # as well.
          num_splits = self.num_packs
  
          # The array_ops.size function will sometimes remove static shapes. So if
          # all gradient shapes are defined, we use another method to get the
          # total size.
          # TODO(yuefengz): move this logic to array_ops.size.
          if all(g.shape.is_fully_defined() for g, _ in device_grads_and_vars):
            total_grad_size = sum(
                [g.shape.num_elements() for g, _ in device_grads_and_vars])
          else:
            total_grad_size = array_ops.size(concat_grads)
  
          split_size = total_grad_size // num_splits
          split_size_last = total_grad_size - split_size * (num_splits - 1)
          split_sizes = [split_size] * (num_splits - 1) + [split_size_last]
          grad_packs = array_ops.split(concat_grads, split_sizes)
  
          # Ready to aggregate the repacked gradients, with fake variables.
          # TODO(zhengxq): It is hacky to have to use fake variables.
          # We should remove the need for variables in
          # aggregate_gradients_using*.
          device_grad_packs.append(zip(grad_packs, [None] * num_splits))
          self.all_device_shapes.append(device_shapes)
          self.all_device_sizes.append(device_sizes)
  
      return device_grad_packs
function_tokens: def pack ( self , grouped_grads_and_vars ) : self . grouped_grads_and_vars = grouped_grads_and_vars self . all_device_shapes = [ ] self . all_device_sizes = [ ] device_grad_packs = [ ] for device_grads_and_vars in grouped_grads_and_vars : with ops . colocate_with ( device_grads_and_vars [ 0 ] [ 0 ] ) : # Flatten all the grads. flat_grads = [ array_ops . reshape ( g , [ - 1 ] ) for g , _ in device_grads_and_vars ] # Remember the original shape of all the grads. device_shapes = [ array_ops . shape ( g ) for g , _ in device_grads_and_vars ] # Remember the original sizes of all the grads. device_sizes = [ array_ops . size ( g ) for g , _ in device_grads_and_vars ] # Concat all the flat grads into a big flat tensor. concat_grads = array_ops . concat ( flat_grads , 0 ) # Split the big tensor into num_splits packs. In cases where the # total size is not divisible num_splits, the last pack gets # more elements. # TODO(zhengxq): it is also possible to optimize away all the concat # as well. num_splits = self . num_packs # The array_ops.size function will sometimes remove static shapes. So if # all gradient shapes are defined, we use another method to get the # total size. # TODO(yuefengz): move this logic to array_ops.size. if all ( g . shape . is_fully_defined ( ) for g , _ in device_grads_and_vars ) : total_grad_size = sum ( [ g . shape . num_elements ( ) for g , _ in device_grads_and_vars ] ) else : total_grad_size = array_ops . size ( concat_grads ) split_size = total_grad_size // num_splits split_size_last = total_grad_size - split_size * ( num_splits - 1 ) split_sizes = [ split_size ] * ( num_splits - 1 ) + [ split_size_last ] grad_packs = array_ops . split ( concat_grads , split_sizes ) # Ready to aggregate the repacked gradients, with fake variables. # TODO(zhengxq): It is hacky to have to use fake variables. # We should remove the need for variables in # aggregate_gradients_using*. device_grad_packs . append ( zip ( grad_packs , [ None ] * num_splits ) ) self . all_device_shapes . append ( device_shapes ) self . all_device_sizes . append ( device_sizes ) return device_grad_packs