


def right_inverse(self, Z):
    # Project Z onto the rank 1 matrices
    U, S, Vh = torch.linalg.svd(Z, full_matrices=False)
    # Return rescaled singular vectors
    s0_sqrt = S[0].sqrt().unsqueeze(-1)
    return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt

linear_rank_one = P.register_parametrization(nn.Linear(4, 4), "weight", RankOne())
print(torch.linalg.matrix_rank(linear_rank_one.weight).item())
   1


