parameters: (threshold_size, device_grads)
return_statement: return small_grads, large_grads
docstring: 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 <= threshold_size
         elements.
      large_grads: Subset of device_grads where shape is > threshold_size
         elements.
docstring_summary: Break gradients into two sets according to tensor size.
docstring_tokens: Break gradients into two sets according to tensor size .
function: def split_grads_by_size(threshold_size, device_grads):
    """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 <= threshold_size
         elements.
      large_grads: Subset of device_grads where shape is > threshold_size
         elements.
    """
    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
function_tokens: 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