jnp.broadcast_to: better error for invalid shape

Author: jakevdp

jax/_src/numpy/util.py

@@ -418,6 +418,8 @@ def _broadcast_to(arr: ArrayLike, shape: DimSize | Shape) -> Array:
   arr_shape = np.shape(arr)
   if core.definitely_equal_shape(arr_shape, shape):
     return arr
+  elif len(shape) < len(arr_shape):
+    raise ValueError(f"Cannot broadcast to shape with fewer dimensions: {arr_shape=} {shape=}")
   else:
     nlead = len(shape) - len(arr_shape)
     shape_tail = shape[nlead:]

tests/lax_numpy_test.py

@@ -5177,6 +5177,13 @@ def testBroadcastTo(self, from_shape, to_shape):
     self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
     self._CompileAndCheck(jnp_op, args_maker)
 
+  def testBroadcastToInvalidShape(self):
+    # Regression test for https://github.com/google/jax/issues/20533
+    x = jnp.zeros((3, 4, 5))
+    with self.assertRaisesRegex(
+        ValueError, "Cannot broadcast to shape with fewer dimensions"):
+      jnp.broadcast_to(x, (4, 5))
+
   @jtu.sample_product(
     [dict(shapes=shapes, broadcasted_shape=broadcasted_shape)
       for shapes, broadcasted_shape in [