Concatenate gradients together more intelligently.

  Does binpacking
  Args:
    tower_grads: List of lists of (gradient, variable) tuples.
    max_bytes: Int giving max number of bytes in a tensor that
      may be considered small.

def pack_small_tensors(tower_grads, max_bytes=0):
    """Concatenate gradients together more intelligently.

  Does binpacking
  Args:
    tower_grads: List of lists of (gradient, variable) tuples.
    max_bytes: Int giving max number of bytes in a tensor that
      may be considered small.
  """
    assert max_bytes >= 0
    orig_grads = [g for g, _ in tower_grads[0]]
    # Check to make sure sizes are accurate; not entirely important
    assert all(g.dtype == tf.float32 for g in orig_grads)
    sizes = [4 * g.shape.num_elements() for g in orig_grads]
    print_stats(sizes)
    small_ranges = []
    large_indices = []
    new_sizes = []

    def end_interval(indices, small_ranges, large_indices):
        if len(indices) > 1:
            small_ranges.insert(0, [indices[0], indices[-1]])
        else:
            large_indices.insert(0, indices[0])

    cur_range = []
    cur_size = 0
    for i, s in reversed(list(enumerate(sizes))):
        if cur_size > max_bytes:
            end_interval(cur_range, small_ranges, large_indices)
            new_sizes.insert(0, cur_size)
            cur_range = []
            cur_size = 0
        cur_range.insert(0, i)
        cur_size += s
    end_interval(cur_range, small_ranges, large_indices)
    new_sizes.insert(0, cur_size)

    print_stats(new_sizes)
    num_gv = len(orig_grads)
    packing = {}
    if len(small_ranges):
        new_tower_grads = []
        for dev_idx, gv_list in enumerate(tower_grads):
            assert len(gv_list) == num_gv, (
                "Possible cause: "
                "Networks constructed on different workers "
                "don't have the same number of variables. "
                "If you use tf.GraphKeys or tf.global_variables() "
                "with multiple graphs per worker during network "
                "construction, you need to use "
                "appropriate scopes, see "
                "https://github.com/ray-project/ray/issues/3136")
            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_tower_grads.append(new_gv_list)
        return new_tower_grads, packing
    else:
        return tower_grads, None