torch: minimally pad tensor such that num elements divisible by x
Suppose I have a tensor t of arbitrary ndim
I want to pad (with zeroes) it such that
a) I introduce the fewest possible # elements
b) after padding, (t.numel() % x) == 0
Is there a better algorithm for doing this than
find the largest dimension and increase it by 1 until condition (b) is satisfied?
Maybe working code:
def pad_minimally(t, x):
    largest_dim = np.argmax(t.shape)
    buffer_shape = list(t.shape)
    new_t = t.clone()
    print(t.shape)
    for n_to_add in range(x):
        if new_t.numel() % x == 0:
            break
        buffer_shape[largest_dim] = n_to_add
        new_buffer = torch.zeros(*buffer_shape)
        new_t = torch.cat([t, new_buffer], axis=largest_dim)
    assert new_t.numel() % x == 0
    return new_t
assert pad_minimally(torch.rand(3,1), 7).shape == (7,1)
assert pad_minimally(torch.rand(3,2), 7).shape == (7,2)
assert pad_minimally(torch.rand(3,2, 6), 7).shape == (3,2,7)


First off, simply adding one to the largest dimension until numel is divisible by x doesn't work in all cases. For example if the shape of t is (3, 2) and x = 9 then we would want to pad t to be (3, 3), not (9, 2).
Even more concerning is that there's no guarantee that only one dimension needs to be padded. For example if t has shape (13, 17, 25) and x = 8 then the optimally padded t would be either (14, 18, 26) or (13, 18, 28).
Distilling this into the mathematics, the problem becomes

Given positive integers s[1], ..., s[D] find positive integers q[1], ..., q[D] that minimize prod(q[i], i=1 to D) subject to the constraints that prod(q[i], i=1 to D) is divisible by x and q[i] &gt;= s[i] for all i=1 to D.

I wasn't able to develop an efficient solution (see update for more efficient solution), though I'm not particularly well versed in non-linear integer programming. Perhaps an efficient solution to this problem exists. If it does I imagine it would involve the prime factors of x and q and/or better memoization. That said, it is possible to solve the problem using an exhaustive search, provided that x and D (i.e. len(t.shape)) are sufficiently small (otherwise the algorithm may run for a really really long time).
The brute force search algorithm I came up with iterates over each multiple of x greater-than or equal to t.numel() and uses depth-first search to see if a padding exists for that multiple. As soon as a valid padding is found the algorithm finishes. The python code for this algorithm is:
import numpy as np

def search(shape, target_numel, memory):
    numel = np.prod(shape)
    if numel == target_numel:
        return True
    elif numel &lt; target_numel:
        for idx in range(len(shape)):
            shape[idx] += 1
            if tuple(shape) not in memory:
                if search(shape, target_numel, memory):
                    return True
                memory.add(tuple(s for s in shape))
            shape[idx] -= 1
    return False

def minimal_shape(shape, target_multiple):
    shape = [s for s in shape]
    target_numel = target_multiple * int(np.ceil(max(1, np.prod(shape)) / target_multiple))
    while not search(shape, target_numel, set()):
        target_numel += target_multiple
    return shape

Once you have the minimal shape, the pad_minimal function can be implemented pretty succinctly as
def pad_minimally(t, x):
    new_shape = minimal_shape(t.shape, x)
    new_t = t.new_zeros(new_shape)
    new_t[[slice(0, s) for s in t.shape]] = t
    return new_t

I'm not sure if this will be fast enough for your needs. Hopefully someone else can come along with a more efficient version.

Some test cases for minimal_shape
assert minimal_shape([2, 2], 9) == [3, 3]
assert minimal_shape([2, 8], 6) == [2, 9]
assert minimal_shape([13, 17, 25], 8) in [[14, 18, 26], [13, 18, 28]]
assert minimal_shape([5, 13, 19], 6) == [5, 14, 21]


Update
I asked about this algorithm on CS.SE. Based on the answer I received there and the subsequent update to the question the following is a much more efficient implementation of minimal_shape.
from functools import reduce
from operator import mul
from copy import deepcopy

def prod(x):
    return reduce(mul, x, 1)

def argsort(x, reverse=False):
    return sorted(range(len(x)), key=lambda idx: x[idx], reverse=reverse)

def divisors(v):
    &quot;&quot;&quot; does not include 1 &quot;&quot;&quot;
    d = {v} if v &gt; 1 else set()
    for n in range(2, int(v**0.5) + 1):
        if v % n == 0:
            d.add(n)
            d.add(v // n)
    return d

def update_memory(b, c_rem, memory):
    tuple_m = tuple(b + [c_rem])
    if tuple_m in memory:
        return False
    memory.add(tuple_m)
    return True

def dfs(a, b, c, c_rem, memory, p_best=float('inf'), b_best=None):
    ab = [ai + bi for ai, bi in zip(a, b)]
    p = prod(ab)
    if p &gt;= p_best:
        return p_best, b_best
    elif p % c == 0:
        return p, deepcopy(b)

    dc = divisors(c_rem)
    for i in argsort(ab):
        for d in dc:
            db = (d - ab[i]) % d
            b[i] += db
            if update_memory(b, c_rem // d, memory):
                p_best, b_best = dfs(a, b, c, c_rem // d, memory, p_best, b_best)
            b[i] -= db

    return p_best, b_best

def minimal_shape(shape, target_multiple):
    a = list(shape)
    b = [0 for _ in range(len(a))]
    c = target_multiple
    _, b = dfs(a, b, c, c, set())
    return [ai + bi for a, b in zip(a, b)]

