Muon Newton-Schulz Iteration

Problem

Orthogonalize a small square matrix with the Muon Newton-Schulz iteration.

Normalize g by its Frobenius norm to get X, then repeat steps times: A = X X^T followed by X = 1.5 * X - 0.5 * A X. Return the final matrix as a list of rows. If the input has zero norm, return a zero matrix.

Hint

The iteration pushes every singular value of the normalized matrix toward 1.

Test cases

case 1

in: [[[1,0],[0,2]],2]

expected: [[0.8164331399447233,0],[0,0.9996118286615059]]

case 2

in: [[[0.5,-0.5],[1,0]],1]

expected: [[0.408248290463863,-0.5443310539518174],[0.8845379626717031,0.0680413817439772]]

case 3

in: [[[0,0],[0,0]],3]

expected: [[0,0],[0,0]]

case 4

in: [[[3,4],[0,0]],1]

expected: [[0.5999999999999999,0.8000000000000002],[0,0]]

Starter codePython

def muon_newton_schulz(g, steps):
    # Your code here
    pass

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