def split_grads_by_size(threshold_size, device_grads):
    'Break gradients into two sets according to tensor size.\n\n  Args:\n    threshold_size: int size cutoff for small vs large tensor.\n    device_grads: List of lists of (gradient, variable) tuples.  The outer\n        list is over devices. The inner list is over individual gradients.\n\n  Returns:\n    small_grads: Subset of device_grads where shape is <= threshold_size\n       elements.\n    large_grads: Subset of device_grads where shape is > threshold_size\n       elements.\n  '
    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)
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.
def split_grads_by_size(threshold_size, device_grads):
    'Break gradients into two sets according to tensor size.\n\n  Args:\n    threshold_size: int size cutoff for small vs large tensor.\n    device_grads: List of lists of (gradient, variable) tuples.  The outer\n        list is over devices. The inner list is over individual gradients.\n\n  Returns:\n    small_grads: Subset of device_grads where shape is <= threshold_size\n       elements.\n    large_grads: Subset of device_grads where shape is > threshold_size\n       elements.\n  '
    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)