Repository: google/jax
Issue: jax.numpy.full((2,2),[1,2]) fails

The last example at https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.full.html

`jax.numpy.full((2,2),[1,2])`

couldn't pass. It gives 

`TypeError: full must be called with scalar fill_value, got fill_value.shape (2,).`

Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_test.py, jax/test_util.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index d7b79bdc3..b82619b85 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1665,6 +1665,7 @@ The JAX version does not necessarily return a view of the input.
 """)
 def broadcast_to(arr, shape):
   arr = arr if isinstance(arr, ndarray) else array(arr)
+  shape = (shape,) if ndim(shape) == 0 else shape
   shape = canonicalize_shape(shape)  # check that shape is concrete
   arr_shape = _shape(arr)
   if arr_shape == shape:
@@ -2879,17 +2880,26 @@ def ones_like(a, dtype=None, shape=None):
 @_wraps(np.full)
 def full(shape, fill_value, dtype=None):
   lax._check_user_dtype_supported(dtype, "full")
-  shape = (shape,) if ndim(shape) == 0 else shape
-  return lax.full(shape, fill_value, dtype)
+  _check_arraylike("full", fill_value)
+  if ndim(fill_value) == 0:
+    shape = (shape,) if ndim(shape) == 0 else shape
+    return lax.full(shape, fill_value, dtype)
+  else:
+    return broadcast_to(asarray(fill_value, dtype=dtype), shape)
 
 
 @_wraps(np.full_like)
 def full_like(a, fill_value, dtype=None, shape=None):
-  _check_arraylike("full_like", a)
   lax._check_user_dtype_supported(dtype, "full_like")
-  if np.isscalar(shape):
-    shape = (shape,)
-  return lax.full_like(a, fill_value, dtype, shape)
+  _check_arraylike("full_like", a, fill_value)
+  if shape is not None:
+    shape = (shape,) if ndim(shape) == 0 else shape
+  if ndim(fill_value) == 0:
+    return lax.full_like(a, fill_value, dtype, shape)
+  else:
+    shape = np.shape(a) if shape is None else shape
+    dtype = _dtype(a) if dtype is None else dtype
+    return broadcast_to(asarray(fill_value, dtype=dtype), shape)
 
 
 @_wraps(np.zeros)
diff --git a/jax/test_util.py b/jax/test_util.py
index 6b1d4ab4f..0a9b7b81f 100644
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -452,6 +452,8 @@ def _dims_of_shape(shape):
     return shape
   elif isinstance(shape, ScalarShape):
     return ()
+  elif np.ndim(shape) == 0:
+    return (shape,)
   else:
     raise TypeError(type(shape))
 
@@ -467,6 +469,9 @@ def _cast_to_shape(value, shape, dtype):
   elif type(shape) in (list, tuple):
     assert np.shape(value) == tuple(shape)
     return value
+  elif np.ndim(shape) == 0:
+    assert np.shape(value) == (shape,)
+    return value
   else:
     raise TypeError(type(shape))
 
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index d3adb3730..f7c665dd6 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -61,7 +61,7 @@ scalar_shapes = [jtu.NUMPY_SCALAR_SHAPE, jtu.PYTHON_SCALAR_SHAPE]
 array_shapes = nonempty_array_shapes + empty_array_shapes
 nonzerodim_shapes = nonempty_nonscalar_array_shapes + empty_array_shapes
 nonempty_shapes = scalar_shapes + nonempty_array_shapes
-all_shapes =  scalar_shapes + array_shapes
+all_shapes = scalar_shapes + array_shapes
 
 float_dtypes = jtu.dtypes.all_floating
 complex_dtypes = jtu.dtypes.complex
@@ -92,7 +92,7 @@ def _shape_and_dtypes(shapes, dtypes):
       yield (shape, dtype)
 
 def _compatible_shapes(shape):
-  if shape in scalar_shapes:
+  if shape in scalar_shapes or np.ndim(shape) == 0:
     return [shape]
   return (shape[n:] for n in range(len(shape) + 1))
 
@@ -2483,19 +2483,21 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CompileAndCheck(jnp_fun, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_inshape={}_outdtype={}".format(
+      {"testcase_name": "_inshape={}_outdtype={}_fillshape={}".format(
           jtu.format_shape_dtype_string(shape, fill_value_dtype),
-          np.dtype(out_dtype).name if out_dtype else "None"),
-       "shape": shape, "fill_value_dtype": fill_value_dtype,
-       "out_dtype": out_dtype}
+          np.dtype(out_dtype).name if out_dtype else "None",
+          fill_value_shape),
+       "fill_value_dtype": fill_value_dtype, "fill_value_shape": fill_value_shape,
+       "shape": shape, "out_dtype": out_dtype}
       for shape in array_shapes + [3, np.array(7, dtype=np.int32)]
       for fill_value_dtype in default_dtypes
+      for fill_value_shape in _compatible_shapes(shape)
       for out_dtype in [None] + default_dtypes))
-  def testFull(self, shape, fill_value_dtype, out_dtype):
+  def testFull(self, shape, fill_value_dtype, fill_value_shape, out_dtype):
     rng = jtu.rand_default(self.rng())
     np_fun = lambda fill_value: np.full(shape, fill_value, dtype=out_dtype)
     jnp_fun = lambda fill_value: jnp.full(shape, fill_value, dtype=out_dtype)
-    args_maker = lambda: [rng((), fill_value_dtype)]
+    args_maker = lambda: [rng(fill_value_shape, fill_value_dtype)]
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
@@ -2558,20 +2560,20 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
 
   @unittest.skipIf(numpy_version < (1, 17), "shape parameter not supported in older numpy")
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_inshape={}_filldtype={}_outdtype={}_outshape={}".format(
+      {"testcase_name": "_inshape={}_filldtype={}_fillshape={}_outdtype={}_outshape={}".format(
           jtu.format_shape_dtype_string(shape, in_dtype),
-          np.dtype(fill_value_dtype).name,
-          np.dtype(out_dtype).name,
-          out_shape),
+          np.dtype(fill_value_dtype).name, fill_value_shape,
+          np.dtype(out_dtype).name, out_shape),
        "shape": shape, "in_dtype": in_dtype,
-       "fill_value_dtype": fill_value_dtype, "out_dtype": out_dtype,
-       "out_shape": out_shape}
+       "fill_value_dtype": fill_value_dtype, "fill_value_shape": fill_value_shape,
+       "out_dtype": out_dtype, "out_shape": out_shape}
       for shape in array_shapes
       for out_shape in [None] + array_shapes
       for in_dtype in default_dtypes
       for fill_value_dtype in default_dtypes
+      for fill_value_shape in _compatible_shapes(shape if out_shape is None else out_shape)
       for out_dtype in default_dtypes))
-  def testFullLike(self, shape, in_dtype, fill_value_dtype, out_dtype, out_shape):
+  def testFullLike(self, shape, in_dtype, fill_value_dtype, fill_value_shape, out_dtype, out_shape):
     if numpy_version < (1, 19) and out_shape == ():
       raise SkipTest("Numpy < 1.19 treats out_shape=() like out_shape=None")
     rng = jtu.rand_default(self.rng())
@@ -2579,7 +2581,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       x, fill_value, dtype=out_dtype, shape=out_shape)
     jnp_fun = lambda x, fill_value: jnp.full_like(
       x, fill_value, dtype=out_dtype, shape=out_shape)
-    args_maker = lambda: [rng(shape, in_dtype), rng((), fill_value_dtype)]
+    args_maker = lambda: [rng(shape, in_dtype), rng(fill_value_shape, fill_value_dtype)]
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
@@ -4546,6 +4548,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
           [(3,), (2, 1, 3)],
           [(3,), (3, 3)],
           [(1,), (3,)],
+          [(1,), 3],
       ])
   def testBroadcastTo(self, from_shape, to_shape):
     rng = jtu.rand_default(self.rng())