# Buggy code
def _broadcasting_shape_rule(name, *avals):
    shapes = onp.array([aval.shape for aval in avals if aval.shape])
    if not shapes.size:
        return ()
    if len({len(shape) for shape in shapes}) != 1:
        msg = "{} got arrays of different rank: {}."
        raise TypeError(msg.format(name, ", ".join(map(str, map(tuple, shapes)))))
    is_zero = onp.any(shapes == 0, axis=0)
    max_shape = onp.max(shapes, axis=0)
    result_shape = onp.where(is_zero, 0, max_shape)
    if not onp.all((shapes == result_shape) | (shapes == 1)):
        msg = "{} got incompatible shapes for broadcasting: {}."
        raise TypeError(msg.format(name, ", ".join(map(str, map(tuple, shapes)))))
    return tuple(result_shape)


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-9-b3bfdb20b090> in <module>()
3
4 jax.shapecheck(["(m,n)", "n"], "(m,n)")(jnp.add)  # passes
----> 5 jax.shapecheck(["n", "(m,n)"], "(m,n)")(jnp.add)  # errors

google3/third_party/py/jax/api.py in shapecheck(in_shapes, out_shape, fun)
1263   flat_fun, out_tree_ = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
1264   avals = map(partial(ShapedArray, dtype=onp.float32), in_shapes)
-> 1265   out_shapes_ = [o.shape for o in pe.abstract_eval_fun(flat_fun.call_wrapped, *avals)]
1266   if out_tree != out_tree_(): raise TypeError("pytree mismatch")
1267   if not all(map(masking._shape_spec_consistent, out_shapes, out_shapes_)):

google3/third_party/py/jax/interpreters/partial_eval.py in abstract_eval_fun(fun, *avals, **params)
340   pvals_in = [PartialVal.unknown(a) for a in avals]
341   _, pvals_out, _ = trace_to_jaxpr(lu.wrap_init(fun, params), pvals_in,
--> 342                                   instantiate=True, stage_out=True)
343   avals_out, _ = unzip2(pvals_out)
344   for aval_out in avals_out:

google3/third_party/py/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out, bottom, trace_type)
436   with new_master(trace_type, bottom=bottom) as master:
437     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 438     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
439     assert not env
440     del master

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
148     gen = None
149
--> 150     ans = self.f(*args, **dict(self.params, **kwargs))
151     del args
152     while stack:

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
148     gen = None
149
--> 150     ans = self.f(*args, **dict(self.params, **kwargs))
151     del args
152     while stack:

google3/third_party/py/jax/numpy/lax_numpy.py in fn(x1, x2)
350 def _maybe_bool_binop(numpy_fn, lax_fn, bool_lax_fn):
351   def fn(x1, x2):
--> 352     x1, x2 = _promote_args(numpy_fn.__name__, x1, x2)
353     return lax_fn(x1, x2) if x1.dtype != bool_ else bool_lax_fn(x1, x2)
354   return _wraps(numpy_fn)(fn)

google3/third_party/py/jax/numpy/lax_numpy.py in _promote_args(fun_name, *args)
281   """Convenience function to apply Numpy argument shape and dtype promotion."""
282   _check_arraylike(fun_name, *args)
--> 283   return _promote_shapes(fun_name, *_promote_dtypes(*args))
284
285 def _promote_args_inexact(fun_name, *args):

google3/third_party/py/jax/numpy/lax_numpy.py in _promote_shapes(fun_name, *args)
218       if FLAGS.jax_numpy_rank_promotion != "allow":
219         _rank_promotion_warning_or_error(fun_name, shapes)
--> 220       result_rank = len(lax.broadcast_shapes(*shapes))
221       return [lax.reshape(arg, (1,) * (result_rank - len(shp)) + shp)
222               if shp and len(shp) != result_rank else arg

google3/third_party/py/jax/lax/lax.py in broadcast_shapes(*shapes)
75   shapes = onp.array([(1,) * (ndim - len(shape)) + shape for shape in shapes])
76   is_zero = onp.any(shapes == 0, axis=0)
---> 77   max_shape = onp.max(shapes, axis=0)
78   result_shape = onp.where(is_zero, 0, max_shape)
79   if not onp.all((shapes == result_shape) | (shapes == 1)):

google3/third_party/py/numpy/core/fromnumeric.py in amax(a, axis, out, keepdims, initial)
2503     """
2504     return _wrapreduction(a, np.maximum, 'max', axis, None, out, keepdims=keepdims,
-> 2505                           initial=initial)
2506
2507

google3/third_party/py/numpy/core/fromnumeric.py in _wrapreduction(obj, ufunc, method, axis, dtype, out, **kwargs)
84                 return reduction(axis=axis, out=out, **passkwargs)
85
---> 86     return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
87
88

google3/third_party/py/jax/interpreters/masking.py in __le__(self, other)
184
185   def __le__(self, other):
--> 186     return _ensure_poly(other) >= self
187
188   def __lt__(self, other):

google3/third_party/py/jax/interpreters/masking.py in __ge__(self, other)
181
182     raise ValueError('Polynomials comparison "{} >= {}" is inconclusive.'
--> 183                      .format(self, other))
184
185   def __le__(self, other):

ValueError: Polynomials comparison "1 >= m" is inconclusive.

# Fixed code
def _broadcasting_shape_rule(name, *avals):
    shapes = onp.array([aval.shape for aval in avals if aval.shape])
    if not shapes.size:
        return ()
    if len({len(shape) for shape in shapes}) != 1:
        msg = "{} got arrays of different rank: {}."
        raise TypeError(msg.format(name, ", ".join(map(str, map(tuple, shapes)))))
    result_shape = _try_broadcast_shapes(shapes)
    if result_shape is None:
        msg = "{} got incompatible shapes for broadcasting: {}."
        raise TypeError(msg.format(name, ", ".join(map(str, map(tuple, shapes)))))
    return result_shape