Integer partition N into M parts in parallel

Question

I am trying to randomly generate integer partitions (N into M parts) in pytorch with a minimum partition size of 1.

For example, (3, 1, 1) and (4, 1, 0) are both partitions of 5 into 3 parts but (4, 1, 0)'s minimum partition size is 0 so this should not be allowed

I would like to use this to generate my dataset on demand, so would nice if there was a pytorch (parrallel/gpgpu) solution.

See other questions about generating integer partitions:

• Generate restricted weak integer compositions (or partitions) of an integer n into k parts in Python
• Elegant Python code for Integer Partitioning

Answer 1

Here's a solution that works in limited cases, M hast to divide N, 2 has to divide M, and the maximium is limited, but this is the behaviour I wanted.

You start of with the equal partition then calculate a delta that sums to zero...

M = 4
N = 16
MINIMUM = 1
assert N % M == 0
assert M % 2 == 0
avg = N // M

equal_partition = torch.full((M,), avg)
half_delta = torch.randint(-(avg-MINIMUM), avg-MINIMUM, (M // 2,))
delta = torch.cat((half_delta, -half_delta), dim=-1)[torch.randperm(M)]
partition = equal_partition + delta
print(partition, partition.sum())