K-FAC Factor Shapes

Problem

Report the Kronecker factor shapes used by K-FAC for a dense layer with in_dim inputs and out_dim outputs.

The activation covariance has shape in_dim x in_dim and the gradient covariance has shape out_dim x out_dim. Return [a_shape, g_shape, overhead] where overhead is the factor storage (in_dim^2 + out_dim^2) divided by the weight storage in_dim * out_dim. Return overhead 0.0 for non-positive dimensions.

Hint

Factor storage is quadratic in each dimension rather than the product.

Test cases

case 1

in: [4,3]

expected: [[4,4],[3,3],2.0833333333333335]

case 2

in: [10,5]

expected: [[10,10],[5,5],2.5]

case 3

in: [1,1]

expected: [[1,1],[1,1],2]

case 4

in: [8,2]

expected: [[8,8],[2,2],4.25]

Starter codePython

def kfac_factor_shapes(in_dim, out_dim):
    # Returns [a_shape, g_shape, overhead]
    # Your code here
    pass

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