Repository: google/jax

Issue:
Shape checking rule issue
```pytb
---------------------------------------------------------------------------
LayerError                                Traceback (most recent call last)
<ipython-input-75-a4586cbd092a> in <module>
      1 # sample prediction
----> 2 tmp_pred = model(x)
      3 print(type(tmp_pred))
      4 print(f"tmp_pred has shape: {tmp_pred.shape}")

/opt/conda/lib/python3.7/site-packages/trax/layers/base.py in __call__(self, x, weights, state, rng)
    165       self.state = state  # Needed if the model wasn't fully initialized.
    166     state = self.state
--> 167     outputs, new_state = self.pure_fn(x, weights, state, rng)
    168     self.state = new_state
    169     self.weights = weights

/opt/conda/lib/python3.7/site-packages/trax/layers/base.py in pure_fn(self, x, weights, state, rng, use_cache)
    448       name, trace = self._name, _short_traceback(skip=3)
    449       raise LayerError(name, 'pure_fn',
--> 450                        self._caller, signature(x), trace) from None
    451 
    452   def output_signature(self, input_signature):
`
LayerError: Exception passing through layer Serial (in pure_fn):
  layer created in file [...]/<ipython-input-68-50afc500007a>, line 16
  layer input shapes: ShapeDtype{shape:(10, 27), dtype:int64}

  File [...]/trax/layers/combinators.py, line 88, in forward
    outputs, s = layer.pure_fn(inputs, w, s, rng, use_cache=True)

LayerError: Exception passing through layer Embedding_369_15 (in pure_fn):
  layer created in file [...]/<ipython-input-68-50afc500007a>, line 13
  layer input shapes: ShapeDtype{shape:(10, 27), dtype:int64}

  File [...]/trax/layers/core.py, line 150, in forward
    return jnp.take(self.weights, x, axis=0)

  File [...]/jax/numpy/lax_numpy.py, line 3298, in take
    slice_sizes=tuple(slice_sizes))

  File [...]/jax/lax/lax.py, line 835, in gather
    slice_sizes=canonicalize_shape(slice_sizes))

  File [...]/site-packages/jax/core.py, line 273, in bind
    return self.impl(*args, **kwargs)

  File [...]/jax/interpreters/xla.py, line 228, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)

  File [...]/jax/interpreters/xla.py, line 262, in xla_primitive_callable
    *avals, **params)

  File [...]/jax/interpreters/xla.py, line 320, in primitive_computation
    raise RuntimeError(msg) from e

RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 1), got 1.: 
This is a bug in JAX's shape-checking rules; please report it!
```


Discussion:
Thanks! The traceback on its own is not particularly helpful for diagnosing the problem. Can you provide a minimal code snippet that reproduces this error?

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -4676,6 +4676,11 @@ def take(a, indices, axis: Optional[int] = None, out=None, mode=None):
 
   index_dims = len(shape(indices))
   slice_sizes = list(shape(a))
+  if slice_sizes[axis_idx] == 0:
+    if indices.size != 0:
+      raise IndexError("Cannot do a non-empty jnp.take() from an empty axis.")
+    return a
+
   slice_sizes[axis_idx] = _min(indices.size, 1)
   dnums = lax.GatherDimensionNumbers(
     offset_dims=tuple(


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
@@ -3933,6 +3933,14 @@ def testTakeEmpty(self):
       jnp.array([], dtype=jnp.float32),
       jnp.take(jnp.array([], jnp.float32), jnp.array([], jnp.int32)))
 
+    np.testing.assert_array_equal(
+      jnp.ones((2, 0, 4), dtype=jnp.float32),
+      jnp.take(jnp.ones((2, 0, 4), dtype=jnp.float32), jnp.array([], jnp.int32),
+               axis=1))
+
+    with self.assertRaisesRegex(IndexError, "non-empty jnp.take"):
+      jnp.take(jnp.ones((2, 0, 4), dtype=jnp.float32),
+               jnp.array([0], jnp.int32), axis=1)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_index={}_axis={}".format(