Repository: google/jax

Issue:
throw an error (or warning) when an array is silently broadcast to higher dimension
Feature request submitted based on in person discussion with @mattjj .  Silent overly-generous broadcasting of arrays often masks bugs. (It might make sense for the behavior -- error, warn, silent, where silent corresponds to numpy behavior --  to be controlled by a global variable.) e.g., I believe the following should result in an error or warning:
```
a = np.arange(5) # shape (5,)
b = np.zeros((5,5)) # shape (5,5)
c = a + b # should throw an error, since a and b have different numbers of dimensions
```


Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -27,12 +27,14 @@
 from __future__ import division
 from __future__ import print_function
 
+from distutils.util import strtobool
 import collections
 import itertools
+import os
 import re
 import string
-import warnings
 import types
+import warnings
 
 import numpy as onp
 import opt_einsum
@@ -42,12 +44,21 @@
 from jax import jit, device_put
 from .. import core
 from ..abstract_arrays import UnshapedArray, ShapedArray, ConcreteArray
+from ..config import flags
 from ..interpreters.xla import DeviceArray
 from .. import lax
 from ..util import partial, get_module_functions, unzip2, prod as _prod
 from ..lib import xla_bridge
 from ..lib import xla_client
 
+FLAGS = flags.FLAGS
+flags.DEFINE_enum(
+    'jax_numpy_rank_promotion', os.getenv('JAX_NUMPY_RANK_PROMOTION', 'allow'),
+    enum_values=['allow', 'warn', 'raise'],
+    help=
+    'Control NumPy-style automatic rank promotion broadcasting '
+    '("allow", "warn", or "raise").')
+
 if six.PY3:
   def removechars(s, chars):
     return s.translate(str.maketrans(dict.fromkeys(chars)))
@@ -158,9 +169,20 @@ def _promote_shapes(*args):
     return args
   else:
     shapes = [shape(arg) for arg in args]
-    nd = len(lax.broadcast_shapes(*shapes))
-    return [lax.reshape(arg, (1,) * (nd - len(shp)) + shp)
-            if shp and len(shp) != nd else arg for arg, shp in zip(args, shapes)]
+    ranks = [len(shp) for shp in shapes]
+    if len(set(ranks)) == 1:
+      return args
+    elif FLAGS.jax_numpy_rank_promotion != "raise":
+      if FLAGS.jax_numpy_rank_promotion == "warn":
+        msg = "following NumPy automatic rank promotion behavior for {}."
+        warnings.warn(msg.format(' '.join(map(str, shapes))))
+      nd = len(lax.broadcast_shapes(*shapes))
+      return [lax.reshape(arg, (1,) * (nd - len(shp)) + shp)
+              if shp and len(shp) != nd else arg for arg, shp in zip(args, shapes)]
+    else:
+      msg = ("operands could not be broadcast together with shapes {} "
+             "and with the config option jax_numpy_rank_promotion='raise'.")
+      raise ValueError(msg.format(' '.join(map(str, shapes))))
 
 def _promote_dtypes(*args):
   """Convenience function to apply Numpy argument dtype promotion."""


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -186,9 +186,15 @@ def test_binop_mismatch(self):
     def f(x, y):
       return x + y
 
-    jtu.check_raises(lambda: grad(f)(onp.zeros(3), onp.zeros(4)),
-                     ValueError,
-                     "Incompatible shapes for broadcasting: ((3,), (4,))")
+    jtu.check_raises(
+        lambda: f(np.zeros(3), np.zeros(4)),
+        TypeError,
+        "add got incompatible shapes for broadcasting: (3,), (4,).")
+
+    jtu.check_raises(
+        lambda: grad(f)(onp.zeros(3), onp.zeros(4)),
+        TypeError,
+        "add got incompatible shapes for broadcasting: (3,), (4,).")
 
   def test_dot_mismatch(self):
     def f(x, y):
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -23,6 +23,7 @@
 import operator
 import unittest
 from unittest import SkipTest
+import warnings
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -1829,5 +1830,35 @@ def testMeshGrid(self, shapes, dtype, indexing, sparse, rng):
     lnp_fun = partial(lnp.meshgrid, indexing=indexing, sparse=sparse)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
+  def testDisableNumpyRankPromotionBroadcasting(self):
+    try:
+      prev_flag = FLAGS.jax_numpy_rank_promotion
+      FLAGS.jax_numpy_rank_promotion = "allow"
+      lnp.ones(2) + lnp.ones((1, 2))  # works just fine
+    finally:
+      FLAGS.jax_numpy_rank_promotion = prev_flag
+
+    try:
+      prev_flag = FLAGS.jax_numpy_rank_promotion
+      FLAGS.jax_numpy_rank_promotion = "raise"
+      self.assertRaises(ValueError, lambda: lnp.ones(2) + lnp.ones((1, 2)))
+    finally:
+      FLAGS.jax_numpy_rank_promotion = prev_flag
+
+    try:
+      prev_flag = FLAGS.jax_numpy_rank_promotion
+      FLAGS.jax_numpy_rank_promotion = "warn"
+      with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        lnp.ones(2) + lnp.ones((1, 2))
+        assert len(w) > 0
+        msg = str(w[-1].message)
+        self.assertEqual(
+            msg,
+            "following NumPy automatic rank promotion behavior for (2,) (1, 2).")
+    finally:
+      FLAGS.jax_numpy_rank_promotion = prev_flag
+
+
 if __name__ == "__main__":
   absltest.main()