cusolver bug (maybe SVD?)

[clemisch]

While rerunning the repro from #1259 (batched SVD) on jax version 0.1.45 I got a new error.

Repro:

import jax
import jax.numpy as np
import numpy as onp

x_host = onp.random.rand(100000, 3, 3).astype(onp.float32)
x_gpu = np.array(x_host)

svd_batch = jax.jit(jax.vmap(np.linalg.svd, 0, 0))

u1, s1, v1 = onp.linalg.svd(x_host)
u2, s2, v2 = np.linalg.svd(x_gpu) # Error
u3, s3, v3 = svd_batch(x_gpu)     # Error

Error:

[...]
~/.local/lib/python3.5/site-packages/jaxlib/cusolver.py in gesvd(c, a, full_matrices, compute_uv)
    195   dtype = a_shape.element_type()
    196   b = 1
--> 197   m, n = a_shape.dimensions()
    198   singular_vals_dtype = _real_type(dtype)
    199 

ValueError: too many values to unpack (expected 2)

This repro worked fine when I posted the issue. I guess this is a new bug?

[clemisch]

I compiled and installed the very newest jaxlib and now it works 🤔