Trying to implement gradient for SVD of low-rank matrix

[Jakob-Unfried]

I am trying to implement a function that performs a "reduced" SVD, i.e. for a (m, n) matrix of rank k < min(m,n), only the first k singular values are nonzero and thus only the first k columns of U and the first k rows of Vh are relevant.

the code for the forward pass is

@custom_transforms
def svd_reduced(X: array):
    assert len(X.shape) == 2

    U, S, Vh = np.linalg.svd(X, full_matrices=False)
    k = np.sum(S > 0.).astype(int)
    U = dynamic_slice_in_dim(U, 0, k, axis=1)  # U[:,:k]
    S = dynamic_slice_in_dim(S, 0, k, axis=0)  # S[:k]
    Vh = dynamic_slice_in_dim(Vh, 0, k, axis=0)  # Vh[:k,:]

    return U, S, Vh

I worked out the VJP formula and implemented it (not pasting it all, its rather long)
defvjp(svd_reduced, _vjp_rule_svd)

However I run into a
TypeError: Abstract value passed to int, which requires a concrete value. Try using value.astype(int) instead.
This is because k, which is the rank of the matrix X depends on its values and thus is not available during tracing.

I realise that this means that I will not be able to jit this function, but IIUC it should in principle be possible to calculate its grad and that is all I want.

How would I go about defining this reduced_svd and its vjp_rule ?

[shoyer]

You can't do Python control flow inside custom_transforms currently, but this should be possible after the rewrite of custom transforms lands: #2026

[Jakob-Unfried]

That looks like it will be awesome! Thanks