Repository: google/jax

Issue:
unexpected behaviour of np.broadcast_to
As detected during discussions at https://github.com/google/jax/pull/1515, `np.broadcast_to(np.ones((2, 3)), (1, 3)).shape` should raise an error instead of returning `(1, 2, 3)`.


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
@@ -889,23 +889,20 @@ def broadcast_arrays(*args):
 def broadcast_to(arr, shape):
   """Like Numpy's broadcast_to but doesn't necessarily return views."""
   arr = arr if isinstance(arr, ndarray) or isscalar(arr) else array(arr)
-  shape = tuple(map(int, shape))
-  if _shape(arr) != shape:
-    # TODO(mattjj): revise this to call lax.broadcast_in_dim rather than
-    # lax.broadcast and lax.transpose
-    lax.broadcast_shapes(shape, _shape(arr))  # error checking
-    nlead = len(shape) - len(_shape(arr))
-    diff, = onp.where(onp.not_equal(shape[nlead:], _shape(arr)))
-
+  shape = tuple(map(int, shape))  # check that shape is concrete
+  arr_shape = _shape(arr)
+  if arr_shape == shape:
+    return arr
+  else:
+    nlead = len(shape) - len(arr_shape)
+    compatible = onp.equal(arr_shape, shape[nlead:]) | onp.equal(arr_shape, 1)
+    if nlead < 0 or not onp.all(compatible):
+      msg = "Incompatible shapes for broadcasting: {} and requested shape {}"
+      raise ValueError(msg.format(arr_shape, shape))
+    diff, = onp.where(onp.not_equal(shape[nlead:], arr_shape))
     new_dims = tuple(range(nlead)) + tuple(nlead + diff)
     kept_dims = tuple(onp.delete(onp.arange(len(shape)), new_dims))
-    perm = onp.argsort(new_dims + kept_dims)
-
-    broadcast_dims = onp.take(shape, new_dims)
-    squeezed_array = squeeze(arr, diff)
-    return lax.transpose(lax.broadcast(squeezed_array, broadcast_dims), perm)
-  else:
-    return arr
+    return lax.broadcast_in_dim(squeeze(arr, diff), shape, kept_dims)
 
 
 @_wraps(onp.split)
@@ -3144,4 +3141,4 @@ def _unstack(x):
   if x.ndim == 0:
     raise ValueError("Argument to _unstack must be non-scalar")
   return [lax.index_in_dim(x, i, keepdims=False) for i in range(x.shape[0])]
-setattr(DeviceArray, "_unstack", _unstack)
\ No newline at end of file
+setattr(DeviceArray, "_unstack", _unstack)


Test patch:
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
@@ -1230,7 +1230,6 @@ def f(x):
 
     f(arr)
 
-
   def testNonArrayErrorMessage(self):
     x = [1., 2.]
     y = onp.array([3., 4.])
@@ -1931,6 +1930,27 @@ def body(i, xy):
       any(onp.array_equal(x, onp.full((3, 4), 2., dtype=onp.float32))
           for x in consts))
 
+  @parameterized.named_parameters(
+      {"testcase_name": "_from={}_to={}".format(from_shape, to_shape),
+       "rng": rng, "from_shape": from_shape, "to_shape": to_shape}
+      for from_shape, to_shape in [
+          [(1, 3), (4, 3)],
+          [(3,), (2, 1, 3)],
+          [(3,), (3, 3)],
+          [(1,), (3,)],
+      ]
+      for rng in [jtu.rand_default()])
+  def testBroadcastTo(self, from_shape, to_shape, rng):
+    args_maker = self._GetArgsMaker(rng, [from_shape], [onp.float32])
+    onp_op = lambda x: onp.broadcast_to(x, to_shape)
+    lnp_op = lambda x: lnp.broadcast_to(x, to_shape)
+    self._CheckAgainstNumpy(onp_op, lnp_op, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lnp_op, args_maker, check_dtypes=True)
+
+  def testBroadcastToIssue1522(self):
+    self.assertRaisesRegex(
+        ValueError, "Incompatible shapes for broadcasting: .*",
+        lambda: lnp.broadcast_to(onp.ones((2, 3)), (1, 3)))
 
 # Most grad tests are at the lax level (see lax_test.py), but we add some here
 # as needed for e.g. particular compound ops of interest.