Repository: google/jax

Issue:
Cannot invert zero-shape matrices
Hi all, 

As described, calling `jax.numpy.linalg.inv(jax.numpy.zeros((0, 0)))` fails with a `ZeroDivisionError` in the shape machinery
```python-traceback
---------------------------------------------------------------------------
FilteredStackTrace                        Traceback (most recent call last)
<ipython-input-1-551ceea3fd37> in <module>
      2 
----> 3 jax.numpy.linalg.inv(jax.numpy.zeros((0, 0)))

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/linalg.py in inv(a)
    334       .format(jnp.shape(a)))
--> 335   return solve(
    336     a, lax.broadcast(jnp.eye(a.shape[-1], dtype=lax.dtype(a)), a.shape[:-2]))

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/linalg.py in solve(a, b)
    449   a, b = _promote_arg_dtypes(jnp.asarray(a), jnp.asarray(b))
--> 450   return lax_linalg._solve(a, b)
    451 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in _solve(a, b)
    265     # b.shape == [..., m, k]
--> 266     return api.vmap(custom_solve, b.ndim - 1, max(a.ndim, b.ndim) - 1)(b)
    267 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py in custom_linear_solve(matvec, b, solve, transpose_solve, symmetric)
   2182 
-> 2183   solve_jaxpr, solve_consts, out_tree = _initial_style_jaxpr(
   2184       _shape_checked(partial(solve, matvec), "solve"), in_args_tree, b_avals)

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py in _initial_style_jaxpr(fun, in_tree, in_avals)
     71 def _initial_style_jaxpr(fun: Callable, in_tree, in_avals):
---> 72   jaxpr, out_avals, consts, out_tree = _initial_style_open_jaxpr(fun, in_tree, in_avals)
     73   closed_jaxpr = core.ClosedJaxpr(pe.convert_constvars_jaxpr(jaxpr), ())

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py in _initial_style_open_jaxpr(fun, in_tree, in_avals)
     66   wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
---> 67   jaxpr, out_avals, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals)
     68   return jaxpr, out_avals, consts, out_tree()

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py in f(x)
   2173     def f(x):
-> 2174       y = fun(x)
   2175       _check_shapes(name, "b", y, b_flat)

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in <lambda>(_, x)
    258       lambda x: _matvec_multiply(a, x),
--> 259       solve=lambda _, x: lu_solve(lu_, permutation, x, trans=0),
    260       transpose_solve=lambda _, x: lu_solve(lu_, permutation, x, trans=1))

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in lu_solve(lu, permutation, b, trans)
    987   """LU solve with broadcasting."""
--> 988   return _lu_solve(lu, permutation, b, trans)
    989 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in _lu_solve(lu, permutation, b, trans)
    981                        .format(lu.shape, b.shape))
--> 982   x = _lu_solve_core(lu, permutation, b, trans)
    983   return x[..., 0] if rhs_vector else x

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/vectorize.py in wrapped(*args)
    303       vectorized_func = api.vmap(vectorized_func, in_axes)
--> 304     return vectorized_func(*vec_args)
    305 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/vectorize.py in wrapped(*args)
    134   def wrapped(*args):
--> 135     out = func(*args)
    136     out_shapes = map(jnp.shape, out if isinstance(out, tuple) else [out])

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/vectorize.py in new_func(*args)
    175       args.insert(i, arg)
--> 176     return func(*args)
    177 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in _lu_solve_core(lu, permutation, b, trans)
    938   m = lu.shape[0]
--> 939   x = jnp.reshape(b, (m, -1))
    940   if trans == 0:

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in reshape(a, newshape, order)
   1246   try:
-> 1247     return a.reshape(newshape, order=order)  # forward to method for ndarrays
   1248   except AttributeError:

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in _reshape_method(a, *newshape, **kwargs)
   1291     newshape = newshape[0]
-> 1292   return _reshape(a, newshape, order=order)
   1293 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in _reshape(a, newshape, order)
   1267 def _reshape(a, newshape, order="C"):
-> 1268   computed_newshape = _compute_newshape(a, newshape)
   1269   if order == "C":

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in _compute_newshape(a, newshape)
   1261   if np.any(np.equal(newshape, -1)):
-> 1262     fix = -a.size // (newshape if type(newshape) is Poly else _prod(newshape))
   1263     return [d if d != -1 else fix for d in newshape]

FilteredStackTrace: ZeroDivisionError: integer division or modulo by zero

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

ZeroDivisionError                         Traceback (most recent call last)
<ipython-input-1-551ceea3fd37> in <module>
      1 import jax
      2 
----> 3 jax.numpy.linalg.inv(jax.numpy.zeros((0, 0)))

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/linalg.py in inv(a)
    333     raise ValueError("Argument to inv must have shape [..., n, n], got {}."
    334       .format(jnp.shape(a)))
--> 335   return solve(
    336     a, lax.broadcast(jnp.eye(a.shape[-1], dtype=lax.dtype(a)), a.shape[:-2]))
    337 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    137   def reraise_with_filtered_traceback(*args, **kwargs):
    138     try:
--> 139       return fun(*args, **kwargs)
    140     except Exception as e:
    141       if not is_under_reraiser(e):

/opt/texat-venv/lib/python3.8/site-packages/jax/api.py in f_jitted(*args, **kwargs)
    369         return cache_miss(*args, **kwargs)[0]  # probably won't return
    370     else:
--> 371       return cpp_jitted_f(*args, **kwargs)
    372   f_jitted._cpp_jitted_f = cpp_jitted_f
    373 

/opt/texat-venv/lib/python3.8/site-packages/jax/api.py in cache_miss(*args, **kwargs)
    276       _check_arg(arg)
    277     flat_fun, out_tree = flatten_fun(f, in_tree)
--> 278     out_flat = xla.xla_call(
    279         flat_fun,
    280         *args_flat,

/opt/texat-venv/lib/python3.8/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1227 
   1228   def bind(self, fun, *args, **params):
-> 1229     return call_bind(self, fun, *args, **params)
   1230 
   1231   def process(self, trace, fun, tracers, params):

/opt/texat-venv/lib/python3.8/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1218   tracers = map(top_trace.full_raise, args)
   1219   with maybe_new_sublevel(top_trace):
-> 1220     outs = primitive.process(top_trace, fun, tracers, params)
   1221   return map(full_lower, apply_todos(env_trace_todo(), outs))
   1222 

/opt/texat-venv/lib/python3.8/site-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1230 
   1231   def process(self, trace, fun, tracers, params):
-> 1232     return trace.process_call(self, fun, tracers, params)
   1233 
   1234   def post_process(self, trace, out_tracers, params):

/opt/texat-venv/lib/python3.8/site-packages/jax/core.py in process_call(self, primitive, f, tracers, params)
    596 
    597   def process_call(self, primitive, f, tracers, params):
--> 598     return primitive.impl(f, *tracers, **params)
    599   process_map = process_call
    600 

/opt/texat-venv/lib/python3.8/site-packages/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, donated_invars, *args)
    567 
    568 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name, donated_invars):
--> 569   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
    570                                *unsafe_map(arg_spec, args))
    571   try:

/opt/texat-venv/lib/python3.8/site-packages/jax/linear_util.py in memoized_fun(fun, *args)
    249       fun.populate_stores(stores)
    250     else:
--> 251       ans = call(fun, *args)
    252       cache[key] = (ans, fun.stores)
    253 

/opt/texat-venv/lib/python3.8/site-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    643   abstract_args, arg_devices = unzip2(arg_specs)
    644   if config.omnistaging_enabled:
--> 645     jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(fun, abstract_args)
    646     if any(isinstance(c, core.Tracer) for c in consts):
    647       raise core.UnexpectedTracerError("Encountered an unexpected tracer.")

/opt/texat-venv/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_final(fun, in_avals)
   1228     main.source_info = fun_sourceinfo(fun.f)  # type: ignore
   1229     main.jaxpr_stack = ()  # type: ignore
-> 1230     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1231     del main
   1232   return jaxpr, out_avals, consts

/opt/texat-venv/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1209     trace = DynamicJaxprTrace(main, core.cur_sublevel())
   1210     in_tracers = map(trace.new_arg, in_avals)
-> 1211     ans = fun.call_wrapped(*in_tracers)
   1212     out_tracers = map(trace.full_raise, ans)
   1213   jaxpr, out_avals, consts = frame.to_jaxpr(in_tracers, out_tracers)

/opt/texat-venv/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    158 
    159     try:
--> 160       ans = self.f(*args, **dict(self.params, **kwargs))
    161     except:
    162       # Some transformations yield from inside context managers, so we have to

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/linalg.py in solve(a, b)
    448 def solve(a, b):
    449   a, b = _promote_arg_dtypes(jnp.asarray(a), jnp.asarray(b))
--> 450   return lax_linalg._solve(a, b)
    451 
    452 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in _solve(a, b)
    264   else:
    265     # b.shape == [..., m, k]
--> 266     return api.vmap(custom_solve, b.ndim - 1, max(a.ndim, b.ndim) - 1)(b)
    267 
    268 def _T(x): return jnp.swapaxes(x, -1, -2)

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    137   def reraise_with_filtered_traceback(*args, **kwargs):
    138     try:
--> 139       return fun(*args, **kwargs)
    140     except Exception as e:
    141       if not is_under_reraiser(e):

/opt/texat-venv/lib/python3.8/site-packages/jax/api.py in batched_fun(*args)
   1184     in_axes_flat = flatten_axes("vmap in_axes", in_tree, in_axes)
   1185     _ = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "vmap")
-> 1186     out_flat = batching.batch(flat_fun, args_flat, in_axes_flat,
   1187                               lambda: flatten_axes("vmap out_axes", out_tree(),
   1188                                                    out_axes),

/opt/texat-venv/lib/python3.8/site-packages/jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dests, axis_name)
     33   # executes a batched version of `fun` following out_dim_dests
     34   batched_fun = batch_fun(fun, in_dims, out_dim_dests, axis_name=axis_name)
---> 35   return batched_fun.call_wrapped(*in_vals)
     36 
     37 @lu.transformation_with_aux

/opt/texat-venv/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    158 
    159     try:
--> 160       ans = self.f(*args, **dict(self.params, **kwargs))
    161     except:
    162       # Some transformations yield from inside context managers, so we have to

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py in custom_linear_solve(matvec, b, solve, transpose_solve, symmetric)
   2181   _check_tree("matvec", "b", out_tree, tree)
   2182 
-> 2183   solve_jaxpr, solve_consts, out_tree = _initial_style_jaxpr(
   2184       _shape_checked(partial(solve, matvec), "solve"), in_args_tree, b_avals)
   2185   _check_tree("solve", "b", out_tree, tree)

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py in _initial_style_jaxpr(fun, in_tree, in_avals)
     70 @cache()
     71 def _initial_style_jaxpr(fun: Callable, in_tree, in_avals):
---> 72   jaxpr, out_avals, consts, out_tree = _initial_style_open_jaxpr(fun, in_tree, in_avals)
     73   closed_jaxpr = core.ClosedJaxpr(pe.convert_constvars_jaxpr(jaxpr), ())
     74   return closed_jaxpr, consts, out_tree

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py in _initial_style_open_jaxpr(fun, in_tree, in_avals)
     65 def _initial_style_open_jaxpr(fun: Callable, in_tree, in_avals):
     66   wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
---> 67   jaxpr, out_avals, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals)
     68   return jaxpr, out_avals, consts, out_tree()
     69 

/opt/texat-venv/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_dynamic(fun, in_avals)
   1199     main.source_info = fun_sourceinfo(fun.f)  # type: ignore
   1200     main.jaxpr_stack = ()  # type: ignore
-> 1201     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1202     del main
   1203   return jaxpr, out_avals, consts

/opt/texat-venv/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1209     trace = DynamicJaxprTrace(main, core.cur_sublevel())
   1210     in_tracers = map(trace.new_arg, in_avals)
-> 1211     ans = fun.call_wrapped(*in_tracers)
   1212     out_tracers = map(trace.full_raise, ans)
   1213   jaxpr, out_avals, consts = frame.to_jaxpr(in_tracers, out_tracers)

/opt/texat-venv/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    158 
    159     try:
--> 160       ans = self.f(*args, **dict(self.params, **kwargs))
    161     except:
    162       # Some transformations yield from inside context managers, so we have to

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/control_flow.py in f(x)
   2172   def _shape_checked(fun, name):
   2173     def f(x):
-> 2174       y = fun(x)
   2175       _check_shapes(name, "b", y, b_flat)
   2176       return y

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in <lambda>(_, x)
    257       lax.custom_linear_solve,
    258       lambda x: _matvec_multiply(a, x),
--> 259       solve=lambda _, x: lu_solve(lu_, permutation, x, trans=0),
    260       transpose_solve=lambda _, x: lu_solve(lu_, permutation, x, trans=1))
    261   if a.ndim == b.ndim + 1:

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in lu_solve(lu, permutation, b, trans)
    986 def lu_solve(lu, permutation, b, trans=0):
    987   """LU solve with broadcasting."""
--> 988   return _lu_solve(lu, permutation, b, trans)
    989 
    990 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    137   def reraise_with_filtered_traceback(*args, **kwargs):
    138     try:
--> 139       return fun(*args, **kwargs)
    140     except Exception as e:
    141       if not is_under_reraiser(e):

/opt/texat-venv/lib/python3.8/site-packages/jax/api.py in f_jitted(*args, **kwargs)
    369         return cache_miss(*args, **kwargs)[0]  # probably won't return
    370     else:
--> 371       return cpp_jitted_f(*args, **kwargs)
    372   f_jitted._cpp_jitted_f = cpp_jitted_f
    373 

/opt/texat-venv/lib/python3.8/site-packages/jax/api.py in cache_miss(*args, **kwargs)
    276       _check_arg(arg)
    277     flat_fun, out_tree = flatten_fun(f, in_tree)
--> 278     out_flat = xla.xla_call(
    279         flat_fun,
    280         *args_flat,

/opt/texat-venv/lib/python3.8/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1227 
   1228   def bind(self, fun, *args, **params):
-> 1229     return call_bind(self, fun, *args, **params)
   1230 
   1231   def process(self, trace, fun, tracers, params):

/opt/texat-venv/lib/python3.8/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1218   tracers = map(top_trace.full_raise, args)
   1219   with maybe_new_sublevel(top_trace):
-> 1220     outs = primitive.process(top_trace, fun, tracers, params)
   1221   return map(full_lower, apply_todos(env_trace_todo(), outs))
   1222 

/opt/texat-venv/lib/python3.8/site-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1230 
   1231   def process(self, trace, fun, tracers, params):
-> 1232     return trace.process_call(self, fun, tracers, params)
   1233 
   1234   def post_process(self, trace, out_tracers, params):

/opt/texat-venv/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in process_call(self, call_primitive, f, tracers, params)
   1083   def process_call(self, call_primitive, f, tracers, params):
   1084     in_avals = [t.aval for t in tracers]
-> 1085     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(f, self.main, in_avals)
   1086     if not jaxpr.eqns:
   1087       return core.eval_jaxpr(jaxpr, consts, *tracers)

/opt/texat-venv/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1209     trace = DynamicJaxprTrace(main, core.cur_sublevel())
   1210     in_tracers = map(trace.new_arg, in_avals)
-> 1211     ans = fun.call_wrapped(*in_tracers)
   1212     out_tracers = map(trace.full_raise, ans)
   1213   jaxpr, out_avals, consts = frame.to_jaxpr(in_tracers, out_tracers)

/opt/texat-venv/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    158 
    159     try:
--> 160       ans = self.f(*args, **dict(self.params, **kwargs))
    161     except:
    162       # Some transformations yield from inside context managers, so we have to

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in _lu_solve(lu, permutation, b, trans)
    980                        "(shape {}) must match"
    981                        .format(lu.shape, b.shape))
--> 982   x = _lu_solve_core(lu, permutation, b, trans)
    983   return x[..., 0] if rhs_vector else x
    984 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/vectorize.py in wrapped(*args)
    302       vmap_counts = [max(c - 1, 0) for c in vmap_counts]
    303       vectorized_func = api.vmap(vectorized_func, in_axes)
--> 304     return vectorized_func(*vec_args)
    305 
    306   return wrapped

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/vectorize.py in wrapped(*args)
    133   """Check that output core dimensions match the signature."""
    134   def wrapped(*args):
--> 135     out = func(*args)
    136     out_shapes = map(jnp.shape, out if isinstance(out, tuple) else [out])
    137 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/vectorize.py in new_func(*args)
    174     for i, arg in static_args:
    175       args.insert(i, arg)
--> 176     return func(*args)
    177 
    178   return new_func, dynamic_args

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/lax/linalg.py in _lu_solve_core(lu, permutation, b, trans)
    937 def _lu_solve_core(lu, permutation, b, trans):
    938   m = lu.shape[0]
--> 939   x = jnp.reshape(b, (m, -1))
    940   if trans == 0:
    941     x = x[permutation, :]

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in reshape(a, newshape, order)
   1245 def reshape(a, newshape, order="C"):
   1246   try:
-> 1247     return a.reshape(newshape, order=order)  # forward to method for ndarrays
   1248   except AttributeError:
   1249     return _reshape(a, newshape, order=order)

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in _reshape_method(a, *newshape, **kwargs)
   1290           type(newshape[0]) is not Poly):
   1291     newshape = newshape[0]
-> 1292   return _reshape(a, newshape, order=order)
   1293 
   1294 

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in _reshape(a, newshape, order)
   1266 
   1267 def _reshape(a, newshape, order="C"):
-> 1268   computed_newshape = _compute_newshape(a, newshape)
   1269   if order == "C":
   1270     return lax.reshape(a, computed_newshape, None)

/opt/texat-venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in _compute_newshape(a, newshape)
   1260   newshape = [check(size) for size in newshape] if iterable else check(newshape)
   1261   if np.any(np.equal(newshape, -1)):
-> 1262     fix = -a.size // (newshape if type(newshape) is Poly else _prod(newshape))
   1263     return [d if d != -1 else fix for d in newshape]
   1264   else:

ZeroDivisionError: integer division or modulo by zero
```

Is there any way I can fix this without writing a new variant of the function that handles the zero shape case?



Discussion:
Thanks for the report! I believe this is a bug... the numpy equivalent returns an empty matrix:
```python
import numpy as np
np.linalg.inv(np.zeros((0, 0)))
# array([], shape=(0, 0), dtype=float64)
```
And there are good arguments to be made that an empty matrix is its own inverse.

Resolving patch:
diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
--- a/jax/_src/lax/linalg.py
+++ b/jax/_src/lax/linalg.py
@@ -244,10 +244,11 @@ def _matvec_multiply(a, b):
   return lax.dot(a, b, precision=lax.Precision.HIGHEST)
 
 def _check_solve_shapes(a, b):
-  if not (a.ndim >= 2 and a.shape[-1] == a.shape[-2] and b.ndim >= 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))
+  if not (a.ndim >= 2 and b.ndim in [a.ndim, a.ndim - 1] and
+          a.shape[-1] == a.shape[-2] == b.shape[a.ndim - 2]):
+    raise ValueError(
+        "The arguments to solve must have shapes a=[..., m, m] and "
+        f"b=[..., m, k] or b=[..., m]; got a={a.shape} and b={b.shape}")
 
 def _solve(a, b):
   _check_solve_shapes(a, b)
@@ -956,7 +957,7 @@ def lu_pivots_to_permutation(swaps, m):
 @partial(vectorize, excluded={3}, signature='(n,n),(n),(n,k)->(n,k)')
 def _lu_solve_core(lu, permutation, b, trans):
   m = lu.shape[0]
-  x = jnp.reshape(b, (m, -1))
+  x = jnp.reshape(b, (m, np.prod(b.shape[1:])))
   if trans == 0:
     x = x[permutation, :]
     x = triangular_solve(lu, x, left_side=True, lower=True, unit_diagonal=True)
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -330,8 +330,8 @@ def _pinv_jvp(rcond, primals, tangents):
 @_wraps(np.linalg.inv)
 def inv(a):
   if jnp.ndim(a) < 2 or a.shape[-1] != a.shape[-2]:
-    raise ValueError("Argument to inv must have shape [..., n, n], got {}."
-      .format(jnp.shape(a)))
+    raise ValueError(
+      f"Argument to inv must have shape [..., n, n], got {a.shape}.")
   return solve(
     a, lax.broadcast(jnp.eye(a.shape[-1], dtype=lax.dtype(a)), a.shape[:-2]))
 


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -707,6 +707,7 @@ def tensor_maker():
           ((8, 8), (8, 4)),
           ((1, 2, 2), (3, 2)),
           ((2, 1, 3, 3), (2, 4, 3, 4)),
+          ((1, 0, 0), (1, 0, 2)),
       ]
       for dtype in float_types + complex_types))
   def testSolve(self, lhs_shape, rhs_shape, dtype):
@@ -722,7 +723,7 @@ def testSolve(self, lhs_shape, rhs_shape, dtype):
       {"testcase_name":
        "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
        "shape": shape, "dtype": dtype}
-      for shape in [(1, 1), (4, 4), (2, 5, 5), (200, 200), (5, 5, 5)]
+      for shape in [(1, 1), (4, 4), (2, 5, 5), (200, 200), (5, 5, 5), (0, 0)]
       for dtype in float_types))
   def testInv(self, shape, dtype):
     rng = jtu.rand_default(self.rng())