jax.numpy.squeeze: error for nonexistent / nonsingleton axis
Numpy's squeeze() function fails with an error if the passed axis is either missing or has a size other than 1. jax.numpy.squeeze() should presumably have the same behaviour.
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 737d685bf..9856b3fb3 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -924,15 +924,19 @@ def ravel(a, order="C"):
 
 @_wraps(onp.squeeze)
 def squeeze(a, axis=None):
-  if 1 not in shape(a):
-    return a
+  shape_a = shape(a)
   if axis is None:
-    newshape = [d for d in shape(a) if d != 1]
+    if 1 not in shape_a:
+      return a
+    newshape = [d for d in shape_a if d != 1]
   else:
     if isinstance(axis, int):
       axis = (axis,)
     axis = frozenset(_canonicalize_axis(i, ndim(a)) for i in axis)
-    newshape = [d for i, d in enumerate(shape(a))
+    if _any(shape_a[a] != 1 for a in axis):
+      raise ValueError("cannot select an axis to squeeze out which has size "
+                       "not equal to one")
+    newshape = [d for i, d in enumerate(shape_a)
                 if d != 1 or i not in axis]
   return lax.reshape(a, newshape)
 
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index fd5b67404..cb5b276d4 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1478,6 +1478,23 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_inshape={}_axis={}".format(
+          jtu.format_shape_dtype_string(arg_shape, dtype), ax),
+       "arg_shape": arg_shape, "dtype": dtype, "ax": ax,
+       "rng_factory": jtu.rand_default}
+      for arg_shape, ax in [
+          ((3,), 0),
+          ((1, 3), 1),
+          ((1, 3, 1), (0, 1))]
+      for dtype in default_dtypes))
+  def testSqueezeFailsOnNonsingletonAxis(self, arg_shape, dtype, ax,
+                                         rng_factory):
+    rng = rng_factory()
+    x = jnp.zeros(arg_shape, dtype=dtype)
+    fun = lambda: jnp.squeeze(x, ax)
+    self.assertRaisesRegex(ValueError, "cannot select an axis to squeeze", fun)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_axis={}_weights={}_returned={}".format(
           jtu.format_shape_dtype_string(shape, dtype),
['jax/numpy/lax_numpy.py', 'tests/lax_numpy_test.py']