sp.linalg.solve SEGFAULT when matrix total size is > int32.max

The below code works for a matrix of size n=46340:

import time
import jax.numpy as np
import jax.random as random
import jax.scipy.linalg as linalg


def get_data(n, c):
    A = random.normal(random.PRNGKey(1), (n, n)).astype(np.float32)
    A = A @ A.T + np.identity(n, np.float32)
    b = random.normal(random.PRNGKey(1), (n, c)).astype(np.float32)
    return A, b


def solve(A, b):
    x = linalg.solve(A,
                     b,
                     sym_pos=True,
                     overwrite_a=True,
                     overwrite_b=True,
                     check_finite=False)
    return x


A, b = get_data(n=46340, c=10)

start = time.time()
x = solve(A, b)
x.block_until_ready()
print('shapes = %s, %s,' % (str(A.shape), str(b.shape)),
      'time = %s,' % (time.time() - start),
      'max error = %s.' % np.max(np.abs(A @ x - b)))

But not for n=46341:

/Users/romann/.conda/envs/conda/bin/python /Users/romann/PycharmProjects/untitled2/jax_main.py
/Users/romann/.conda/envs/conda/lib/python3.7/site-packages/jax/lib/xla_bridge.py:114: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
 
Process finished with exit code 139 (interrupted by signal 11: SIGSEGV)

Numpy version works fine for n=50000. Note that the square root of 2147483647 is 46340.95.

It would be very useful to have it work for larger sizes since the training set size for many classic datasets is >50,000, and kernel methods require linalg.solve called on these matrices.

Thanks!

[shoyer]

@romanngg I wrote a full description of this potential optimization in https://github.com/google/jax/issues/1747 (with code that should work in simpler cases)

[hawkinsp]

The segfault that was the subject of this issue is long fixed at this point.