Repository: google/jax

Issue:
shapecheck of arithmetic is non-commutative
```python
import jax.numpy as jnp
import jax

jax.shapecheck(["(m,n)", "n"], "(m,n)")(jnp.add)  # passes
jax.shapecheck(["n", "(m,n)"], "(m,n)")(jnp.add)  # errors
```
This results in the error:
```
---------------------------------------------------------------------------
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.
```


Discussion:
ping @JuliusKunze

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -66,6 +66,14 @@
 DType = Any
 Shape = Sequence[int]
 
+def _try_broadcast_shapes(shapes):
+  # Replace 1 with 0 to avoid inconclusive comparisons for polymorphic dims:
+  out_shape = onp.max(onp.where(shapes == 1, 0, shapes), axis=0)
+  out_shape = onp.where(onp.all(shapes == 1, axis=0), 1, out_shape)
+  if not onp.all((shapes == out_shape) | (shapes == 1)):
+    return None
+  return canonicalize_shape(out_shape)
+
 @cache()
 def broadcast_shapes(*shapes):
   """Returns the shape that results from NumPy broadcasting of `shapes`."""
@@ -73,13 +81,11 @@ def broadcast_shapes(*shapes):
     return shapes[0]
   ndim = _max(len(shape) for shape in shapes)
   shapes = onp.array([(1,) * (ndim - len(shape)) + shape for shape in 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)):
+  result_shape = _try_broadcast_shapes(shapes)
+  if result_shape is None:
     raise ValueError("Incompatible shapes for broadcasting: {}"
                      .format(tuple(map(tuple, shapes))))
-  return canonicalize_shape(result_shape)
+  return result_shape
 
 def _identity(x): return x
 
@@ -1781,13 +1787,11 @@ def _broadcasting_shape_rule(name, *avals):
   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)):
+  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 tuple(result_shape)
+  return result_shape
 
 
 def naryop(result_dtype, accepted_dtypes, name, translation_rule=None):


Test patch:
diff --git a/tests/masking_test.py b/tests/masking_test.py
--- a/tests/masking_test.py
+++ b/tests/masking_test.py
@@ -101,6 +101,7 @@ def test_Poly_rsub(self):
     assert -1 - n == -n - 1
 
   def test_add_broadcast(self):
+    @shapecheck(['n', '(m, n)'], '(m, n)')
     @shapecheck(['(m, n)', 'n'], '(m, n)')
     @shapecheck(['n', ''], 'n')
     def add(a, b):