Linear solvers (jax.numpy.linalg.solve, jax.scipy.linalg.solve) bugged when allocating large memory about jax HOT 3 CLOSED

import jax
import jax.numpy as jnp
device = jax.local_devices()[0]
print('on device:', device)

m = 10

import time
for n in [1e5, 1e6, 1e7]:
    n = int(n)
    A = jnp.repeat(jnp.identity(m)[None], n, axis = 0).block_until_ready()
    print("n=", n)
    
    st_time = time.time()
    U,S,Vh = jax.scipy.linalg.svd(A)
    A2 = jax.lax.batch_matmul(U * S[...,None,:], Vh)
    print(f"SVD error {jnp.mean(jnp.linalg.norm(A2 - A, axis=(-1, -2))/jnp.linalg.norm(A, axis=(-1,-2)))}, time = {time.time() - st_time} ")

    st_time = time.time()
    L = jax.scipy.linalg.cholesky(A)
    A2 = jax.lax.batch_matmul(L, L.swapaxes(-1,-2)).block_until_ready()
    print(f"Cholesky error {jnp.mean(jnp.linalg.norm(A2 - A, axis=(-1, -2))/jnp.linalg.norm(A, axis=(-1,-2)))}, time = {time.time() - st_time} ")

    
    if n <= 1e6:
        st_time = time.time()
        Q,R = jnp.linalg.qr(A)
        A2 = jax.lax.batch_matmul(Q,R).block_until_ready()
        print(f"QR error {jnp.mean(jnp.linalg.norm(A2 - A, axis=(-1, -2))/jnp.linalg.norm(A, axis=(-1,-2)))}, time = {time.time() - st_time} ")

on device: cuda:0
n= 100000
SVD error 0.0, time = 0.6100842952728271 
Cholesky error 0.0, time = 0.15522980690002441 
QR error 0.0, time = 3.7535462379455566 
n= 1000000
SVD error 0.0, time = 0.5560173988342285 
Cholesky error 0.0, time = 0.1310713291168213 
QR error 0.0, time = 35.838584184646606 
n= 10000000
SVD error 0.0, time = 2.056659460067749 
Cholesky error nan, time = 0.27480244636535645 

def solve_by_SVD(A,b):
    U,S,Vh = jax.scipy.linalg.svd(A)

    if b.ndim == A.ndim -1:
        expand = True
        b = b[...,None]
    else:
        expand = False
    
    Uhb = jax.lax.batch_matmul(jnp.conj(U.swapaxes(-1,-2)),b)/ S[...,None]
    x = jax.lax.batch_matmul(jnp.conj(Vh.swapaxes(-1,-2)),Uhb)

    if expand:
        x = x[...,0]

    return x