# Buggy code
def _gather_shape_rule(operand, start_indices, *, dimension_numbers, slice_sizes):
    if len(operand.shape) != len(slice_sizes):
        msg = (
            "slice_sizes must have rank equal to the gather operand; "
            "operand.shape={}, slice_sizes={}".format(operand.shape, slice_sizes)
        )
        raise ValueError(msg)
    result_rank = len(dimension_numbers.offset_dims) + start_indices.ndim - 1
    start_indices_shape = iter(start_indices.shape[:-1])
    slice_sizes = iter(np.delete(slice_sizes, dimension_numbers.collapsed_slice_dims))
    return tuple(
        next(slice_sizes)
        if i in dimension_numbers.offset_dims
        else next(start_indices_shape)
        for i in range(result_rank)
    )


# Traceback (LayerError)
LayerError                                Traceback (most recent call last)
<ipython-input-91-cc4d1abeffd2> in <module>
----> 1 embed_out = embed(x)

/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 Embedding_9088_256 (in pure_fn):
layer created in file [...]/<ipython-input-89-4dbaf3019a8b>, line 1
layer input shapes: ShapeDtype{shape:(16, 18), dtype:int32}

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!

# Fixed code
def _gather_shape_rule(operand, start_indices, *, dimension_numbers, slice_sizes):
    """Validates the well-formedness of the arguments to Gather.

    The code implements the checks based on the detailed operation semantics of
    XLA's `Gather <https://www.tensorflow.org/xla/operation_semantics#gather>`_
    operator and following the outline of the implementation of
    ShapeInference::InferGatherShape in TensorFlow.
    """

    offset_dims = dimension_numbers.offset_dims
    collapsed_slice_dims = dimension_numbers.collapsed_slice_dims
    start_index_map = dimension_numbers.start_index_map

    # Note: in JAX, index_vector_dim is always computed as below, cf. the
    # documentation of the GatherDimensionNumbers class.
    index_vector_dim = _rank(start_indices) - 1

    # This case should never happen in JAX, due to the implicit construction of
    # index_vector_dim, but is included for completeness.
    if _rank(start_indices) < index_vector_dim or index_vector_dim < 0:
        raise TypeError(
            f"Gather index leaf dimension must be within [0, rank("
            f"start_indices) + 1). rank(start_indices) is "
            f"{_rank(start_indices)} and gather index leaf dimension "
            f"is {index_vector_dim}."
        )

    expanded_start_indices_shape = list(start_indices.shape)

    # This case should never happen in JAX, due to the implicit construction of
    # index_vector_dim, but is included for completeness.
    if len(expanded_start_indices_shape) == index_vector_dim:
        expanded_start_indices_shape.append(1)

    # Start ValidateGatherDimensions
    # In the error messages output by XLA, "offset_dims" is called "Output window
    # dimensions" in error messages. For consistency's sake, our error messages
    # stick to "offset_dims".
    _is_sorted(offset_dims, "gather", "offset_dims")
    _no_duplicate_dims(offset_dims, "gather", "offset_dims")

    output_offset_dim_count = len(offset_dims)
    output_shape_rank = len(offset_dims) + _rank(start_indices) - 1

    for i in range(output_offset_dim_count):
        offset_dim = offset_dims[i]
        if offset_dim < 0 or offset_dim >= output_shape_rank:
            raise TypeError(
                f"Offset dimension {i} in gather op is out of bounds; "
                f"got {offset_dim}, but should have been in "
                f"[0, {output_shape_rank})"
            )

    if len(start_index_map) != start_indices.shape[index_vector_dim]:
        raise TypeError(
            f"Gather op has {len(start_index_map)} elements in "
            f"start_index_map and the bound of dimension "
            f"index_vector_dim={index_vector_dim} of start_indices is "
            f"{start_indices.shape[index_vector_dim]}. These two "
            f"numbers must be equal."
        )

    for i in range(len(start_index_map)):
        operand_dim_for_start_index_i = start_index_map[i]
        if operand_dim_for_start_index_i < 0 or operand_dim_for_start_index_i >= _rank(
            operand
        ):
            raise TypeError(
                f"Invalid start_index_map; domain is "
                f"[0, {_rank(operand)}), got: "
                f"{i}->{operand_dim_for_start_index_i}."
            )

    _no_duplicate_dims(start_index_map, "gather", "start_index_map")

    # _is_sorted and _sorted_dims_in_range are checked in the opposite order
    # compared to the XLA implementation. In cases when the input is not sorted
    # AND there are problematic collapsed_slice_dims, the error message will thus
    # be different.
    _is_sorted(collapsed_slice_dims, "gather", "collapsed_slice_dims")
    _sorted_dims_in_range(
        collapsed_slice_dims, _rank(operand), "gather", "collapsed_slice_dims"
    )
    _no_duplicate_dims(collapsed_slice_dims, "gather", "collapsed_slice_dims")
    # End ValidateGatherDimensions

    if _rank(operand) != len(slice_sizes):
        raise TypeError(
            f"Gather op must have one slice size for every input "
            f"dimension; got: len(slice_sizes)={len(slice_sizes)}, "
            f"input_shape.rank={_rank(operand)}"
        )

    if len(slice_sizes) != len(offset_dims) + len(collapsed_slice_dims):
        raise TypeError(
            f"All components of the offset index in a gather op must "
            f"either be a offset dimension or explicitly collapsed; "
            f"got len(slice_sizes)={len(slice_sizes)}, "
            f"output_slice_sizes={offset_dims}, collapsed_slice_dims="
            f"{collapsed_slice_dims}."
        )

    for i in range(len(slice_sizes)):
        slice_size = slice_sizes[i]
        corresponding_input_size = operand.shape[i]

        if slice_size < 0 or slice_size > corresponding_input_size:
            raise TypeError(
                f"Slice size at index {i} in gather op is out of range, "
                f"must be within [0, {corresponding_input_size + 1}), "
                f"got {slice_size}."
            )

    for i in range(len(collapsed_slice_dims)):
        bound = slice_sizes[collapsed_slice_dims[i]]
        if bound > 1:
            raise TypeError(
                f"Gather op can only collapse slice dims with bound 1 "
                f"or 0, but bound is {bound} for index "
                f"{collapsed_slice_dims[i]} at position {i}."
            )

    expanded_start_indices_shape.pop(index_vector_dim)
    start_indices_shape = iter(expanded_start_indices_shape)

    slice_sizes = iter(np.delete(slice_sizes, collapsed_slice_dims))
    return tuple(
        next(slice_sizes) if i in offset_dims else next(start_indices_shape)
        for i in range(output_shape_rank)
    )