Nx defaults to Nx.LinAlg.cholesky method, and numpy is using svd. Everything works after changing the method.