Shampoo Factor Shapes

Problem

Compute Shampoo factor shapes and storage for a tensor shape.

A weight of shape (d1, d2, ...) keeps one d_i x d_i preconditioner factor per dimension. Return [factors, storage, ratio], where factors is the list of [d_i, d_i] pairs, storage is the total number of factor entries, and ratio is the dense Kronecker size prod(d_i)^2 divided by storage.

Hint

Sum d squared over dimensions; the ratio grows quickly with tensor order.

Test cases

case 1

in: [[4,3]]

expected: [[[4,4],[3,3]],25,5.76]

case 2

in: [[10,5]]

expected: [[[10,10],[5,5]],125,20]

case 3

in: [[3,3]]

expected: [[[3,3],[3,3]],18,4.5]

case 4

in: [[2,3,4]]

expected: [[[2,2],[3,3],[4,4]],29,19.862068965517242]

Starter codePython

def shampoo_factor_shapes(shape):
    # Returns [factors, storage, ratio]
    # Your code here
    pass

4 test cases · run your code against them in the editor