# Buggy code
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))
        )
    return solve(
        a, lax.broadcast(jnp.eye(a.shape[-1], dtype=lax.dtype(a)), a.shape[:-2])
    )


# Traceback (ZeroDivisionError)
---------------------------------------------------------------------------
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

# Fixed code
def inv(a):
    if jnp.ndim(a) < 2 or a.shape[-1] != a.shape[-2]:
        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])
    )