code: def split_grads_by_size(threshold_size, device_grads):
      
      small_grads = []
      large_grads = []
      for dl in device_grads:
          small_dl = []
          large_dl = []
          for (g, v) in dl:
              tensor_size = g.get_shape().num_elements()
              if tensor_size <= threshold_size:
                  small_dl.append([g, v])
              else:
                  large_dl.append([g, v])
          if small_dl:
              small_grads.append(small_dl)
          if large_dl:
              large_grads.append(large_dl)
      return small_grads, large_grads
docs: Break gradients into two sets according to tensor size.
  
    Args:
      threshold_size: int size cutoff for small vs large tensor.
      device_grads: List of lists of (gradient, variable) tuples.  The outer
          list is over devices. The inner list is over individual gradients.
  
    Returns:
      small_grads: Subset of device_grads where shape is <= theshold_size
         elements.
      large_grads: Subset of device_grads where shape is > threshold_size
         elements.