def merge_small_dims(var_shape, reshape_size):
    """Computes the shape of the variable for preconditioning.

    If the variable has several small dimensions, we can reshape it so
    that there are fewer big ones. e.g for a convolution (512, 5, 5, 1024)
    we could reshape it into (512, 25, 1024).

    Args:
      var_shape: the shape of the variable
      reshape_size: maximum size of a reshaped dimension
    Returns:
      shape: a list of integers. Product(shape) = number of elements in var.
    """
    shape = []
    product = 1
    for d in var_shape:
        if product * d <= reshape_size:
            product *= d
        else:
            if product > 1:
                shape.append(product)
            product = d
    if product > 1:
        shape.append(product)
    return shape