# Buggy code
def solve(a, b):
    a, b = _promote_arg_dtypes(np.asarray(a), np.asarray(b))
    a_shape = np.shape(a)
    b_shape = np.shape(b)
    a_ndims = len(a_shape)
    b_ndims = len(b_shape)
    if not (
        a_ndims >= 2
        and a_shape[-1] == a_shape[-2]
        and (a_ndims == b_ndims or a_ndims == b_ndims + 1)
    ):
        msg = (
            "The arguments to solve must have shapes a=[..., m, m] and "
            "b=[..., m, k] or b=[..., m]; got a={} and b={}"
        )
        raise ValueError(msg.format(a_shape, b_shape))
    lu, pivots = lax_linalg.lu(a)
    dtype = lax.dtype(a)

    # TODO(phawkins): add unit_diagonal support to solve_triangular, use it here
    # instead of explicit masking of l.
    m = a_shape[-1]
    l = np.tril(lu, -1)[:, :m] + np.eye(m, m, dtype=dtype)

    # TODO(phawkins): triangular_solve only supports matrices on the RHS, so we
    # add a dummy dimension. Extend it to support vectors and simplify this.
    x = b if a_ndims == b_ndims else b[..., None]

    permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
    x = x[..., permutation, :]

    x = lax_linalg.triangular_solve(l, x, left_side=True, lower=True)
    x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=False)

    return x[..., 0] if a_ndims != b_ndims else x


# Traceback (ValueError)
(1, 2, 2) (3, 2)
[[1. 1.]
[1. 1.]
[1. 1.]]
/Users/skainswo/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/mle_normal.py in <module>
31 print(A.shape, b.shape)
32 print(np.linalg.solve(A, b))
---> 33 print(jp.linalg.solve(A, b))

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/numpy/linalg.py in solve(a, b)
247   x = b if a_ndims == b_ndims else b[..., None]
248
--> 249   permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
250   x = x[..., permutation, :]
251

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax_linalg.py in lu_pivots_to_permutation(swaps, k)
462     return swaps, permutation
463
--> 464   n, = np.shape(swaps)
465   permutation = np.arange(k)
466   _, permutation = lax.fori_loop(

ValueError: too many values to unpack (expected 1)

# Fixed code
def solve(a, b):
    a, b = _promote_arg_dtypes(np.asarray(a), np.asarray(b))
    a_shape = np.shape(a)
    b_shape = np.shape(b)
    a_ndims = len(a_shape)
    b_ndims = len(b_shape)
    if not (a_ndims >= 2 and a_shape[-1] == a_shape[-2] and b_ndims >= 1):
        msg = (
            "The arguments to solve must have shapes a=[..., m, m] and "
            "b=[..., m, k] or b=[..., m]; got a={} and b={}"
        )
        raise ValueError(msg.format(a_shape, b_shape))
    lu, pivots = lax_linalg.lu(a)
    dtype = lax.dtype(a)

    m = a_shape[-1]

    # Numpy treats the RHS as a (batched) vector if the number of dimensions
    # differ by 1. Otherwise, broadcasting rules apply.
    x = b[..., None] if a_ndims == b_ndims + 1 else b

    batch_dims = lax.broadcast_shapes(lu.shape[:-2], x.shape[:-2])
    x = np.broadcast_to(x, batch_dims + x.shape[-2:])
    lu = np.broadcast_to(lu, batch_dims + lu.shape[-2:])

    permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
    permutation = np.broadcast_to(permutation, batch_dims + (m,))
    iotas = np.ix_(*(lax.iota(np.int32, b) for b in batch_dims + (1,)))
    x = x[iotas[:-1] + (permutation, slice(None))]

    # TODO(phawkins): add unit_diagonal support to triangular_solve, use it here
    # instead of explicit masking of l.
    l = np.tril(lu, -1)[..., :, :m] + np.eye(m, m, dtype=dtype)
    x = lax_linalg.triangular_solve(l, x, left_side=True, lower=True)
    x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=False)

    return x[..., 0] if a_ndims == b_ndims + 1 else x