def pack_small_tensors(replica_grads, max_bytes=0, max_group=0):
    "Concatenate small gradient tensors together for reduction.\n\n  Args:\n    replica_grads: List of lists of (gradient, variable) tuples.\n    max_bytes: Int giving max number of bytes in a tensor that\n      may be considered small.\n    max_group: Int giving max number of small tensors that may be\n      concatenated into one new tensor.\n\n  Returns:\n    new_replica_grads, packing where new_replica_grads is identical to\n      replica_grads except that all feasible small_tensors have been removed\n      from their places and concatenated into larger tensors that are\n      now in the front of the list for each replica, and packing contains\n      the data necessary to restore the replica_grads structure.\n\n  Look through the first replica for gradients of the same type (float),\n  and small size, that are all sequential.  For each such group,\n  replace by a new tensor that is a flattened concatenation.  Note\n  that the corresponding variable will be absent, which doesn't matter\n  because it isn't used during all-reduce.\n\n  Requires:\n    Every gv_list in replicas must have isomorphic structure including identical\n      tensor sizes and types.\n  "
    small_indices = []
    large_indices = []
    for (idx, (g, _)) in enumerate(replica_grads[0]):
        if ((g.dtype == dtypes.float32) and ((4 * g.shape.num_elements()) <= max_bytes)):
            small_indices.append(idx)
        else:
            large_indices.append(idx)
    (small_ranges, small_singles) = extract_ranges(small_indices, range_size_limit=max_group)
    large_indices = sorted((large_indices + small_singles))
    num_gv = len(replica_grads[0])
    packing = {}
    if small_ranges:
        new_replica_grads = []
        for (dev_idx, gv_list) in enumerate(replica_grads):
            assert (len(gv_list) == num_gv)
            new_gv_list = []
            for r in small_ranges:
                key = ('%d:%d' % (dev_idx, len(new_gv_list)))
                new_gv_list.append((pack_range(key, packing, gv_list, r), 'packing_var_placeholder'))
            for i in large_indices:
                new_gv_list.append(gv_list[i])
            new_replica_grads.append(new_gv_list)
        return (new_replica_grads, packing)
    else:
        return (replica_grads, None)
Concatenate small gradient tensors together for reduction.

Args:
  replica_grads: List of lists of (gradient, variable) tuples.
  max_bytes: Int giving max number of bytes in a tensor that
    may be considered small.
  max_group: Int giving max number of small tensors that may be
    concatenated into one new tensor.

Returns:
  new_replica_grads, packing where new_replica_grads is identical to
    replica_grads except that all feasible small_tensors have been removed
    from their places and concatenated into larger tensors that are
    now in the front of the list for each replica, and packing contains
    the data necessary to restore the replica_grads structure.

Look through the first replica for gradients of the same type (float),
and small size, that are all sequential.  For each such group,
replace by a new tensor that is a flattened concatenation.  Note
that the corresponding variable will be absent, which doesn't matter
because it isn't used during all-reduce.

Requires:
  Every gv_list in replicas must have isomorphic structure including identical
    tensor sizes and types.
def pack_small_tensors(replica_grads, max_bytes=0, max_group=0):
    "Concatenate small gradient tensors together for reduction.\n\n  Args:\n    replica_grads: List of lists of (gradient, variable) tuples.\n    max_bytes: Int giving max number of bytes in a tensor that\n      may be considered small.\n    max_group: Int giving max number of small tensors that may be\n      concatenated into one new tensor.\n\n  Returns:\n    new_replica_grads, packing where new_replica_grads is identical to\n      replica_grads except that all feasible small_tensors have been removed\n      from their places and concatenated into larger tensors that are\n      now in the front of the list for each replica, and packing contains\n      the data necessary to restore the replica_grads structure.\n\n  Look through the first replica for gradients of the same type (float),\n  and small size, that are all sequential.  For each such group,\n  replace by a new tensor that is a flattened concatenation.  Note\n  that the corresponding variable will be absent, which doesn't matter\n  because it isn't used during all-reduce.\n\n  Requires:\n    Every gv_list in replicas must have isomorphic structure including identical\n      tensor sizes and types.\n  "
    small_indices = []
    large_indices = []
    for (idx, (g, _)) in enumerate(replica_grads[0]):
        if ((g.dtype == dtypes.float32) and ((4 * g.shape.num_elements()) <= max_bytes)):
            small_indices.append(idx)
        else:
            large_indices.append(idx)
    (small_ranges, small_singles) = extract_ranges(small_indices, range_size_limit=max_group)
    large_indices = sorted((large_indices + small_singles))
    num_gv = len(replica_grads[0])
    packing = {}
    if small_ranges:
        new_replica_grads = []
        for (dev_idx, gv_list) in enumerate(replica_grads):
            assert (len(gv_list) == num_gv)
            new_gv_list = []
            for r in small_ranges:
                key = ('%d:%d' % (dev_idx, len(new_gv_list)))
                new_gv_list.append((pack_range(key, packing, gv_list, r), 'packing_var_placeholder'))
            for i in large_indices:
                new_gv_list.append(gv_list[i])
            new_replica_grads.append(new_gv_list)
        return (new_replica_grads, packing)
    else:
        return (replica_grads, None)