# Buggy code
def _scatter_shape_rule(
    operand,
    scatter_indices,
    updates,
    *,
    update_jaxpr,
    update_consts,
    dimension_numbers,
    indices_are_sorted,
    unique_indices,
):
    """Validates the well-formedness of the ``dimension_numbers`` argument to
    Scatter.

    The code implements the checks based on the detailed operation semantics of
    XLA's `Scatter <https://www.tensorflow.org/xla/operation_semantics#scatter>`_
    operator and following the outline of the implementation of
    ShapeInference::InferScatterShape in TensorFlow.
    """
    rank = lambda arr: len(arr.shape)

    def _is_sorted(dims, name):
        for i in range(1, len(dims)):
            if dims[i] < dims[i - 1]:
                raise TypeError(f"{name} in scatter op must be sorted; got {dims}")

    def _sorted_dims_in_range(dims, rank, name):
        if len(dims) == 0:
            return
        invalid_dim = None
        if dims[0] < 0:
            invalid_dim = dims[0]
        elif dims[-1] >= rank:
            invalid_dim = dims[-1]
        if invalid_dim:
            raise TypeError(
                f"Invalid {name} set in scatter op; valid range is "
                f"[0, {rank}); got: {invalid_dim}."
            )

    def _no_duplicate_dims(dims, name):
        if len(set(dims)) != len(dims):
            raise TypeError(f"{name} in scatter op must not repeat; got: {dims}.")

    update_window_dims = dimension_numbers.update_window_dims
    inserted_window_dims = dimension_numbers.inserted_window_dims
    scatter_dims_to_operand_dims = dimension_numbers.scatter_dims_to_operand_dims
    # Note: in JAX, index_vector_dim is always computed as below, cf. the
    # documentation of the ScatterDimensionNumbers class.
    index_vector_dim = rank(scatter_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(scatter_indices) < index_vector_dim or index_vector_dim < 0:
        raise TypeError(
            f"Scatter index leaf dimension must be within [0, "
            f"rank(scatter_indices) + 1). rank(scatter_indices) is "
            f"{rank(scatter_indices)} and scatter index leaf "
            f"dimension is {index_vector_dim}."
        )

    expanded_scatter_indices_shape = list(scatter_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_scatter_indices_shape) == index_vector_dim:
        expanded_scatter_indices_shape.append(1)

    expected_updates_rank = (
        len(expanded_scatter_indices_shape) - 1 + len(update_window_dims)
    )

    if rank(updates) != expected_updates_rank:
        raise TypeError(
            f"Updates tensor must be of rank {expected_updates_rank}; "
            f"got {rank(updates)}."
        )

    # Validate update_window_dims
    _is_sorted(update_window_dims, "update_window_dims")
    _no_duplicate_dims(update_window_dims, "update_window_dims")
    _sorted_dims_in_range(update_window_dims, rank(updates), "update_window_dims")

    # Validate inserted_window_dims
    _is_sorted(inserted_window_dims, "inserted_window_dims")
    _no_duplicate_dims(inserted_window_dims, "inserted_window_dims")
    _sorted_dims_in_range(inserted_window_dims, rank(operand), "inserted_window_dims")

    # Validate window_size
    window_size = len(update_window_dims) + len(inserted_window_dims)
    if rank(operand) != window_size:
        raise TypeError(
            f"Scatter op has window of size {window_size}; doesn't "
            f"match operand of rank {rank(operand)}."
        )

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

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

    _no_duplicate_dims(scatter_dims_to_operand_dims, "scatter_dims_to_operand_dims")

    max_update_slice_sizes = [
        operand.shape[i]
        for i in range(len(operand.shape))
        if not i in set(inserted_window_dims)
    ]

    for i in range(len(update_window_dims)):
        update_window_dim = update_window_dims[i]
        if updates.shape[update_window_dim] > max_update_slice_sizes[i]:
            raise TypeError(
                f"Bounds of the window dimensions of updates must not "
                f"exceed the bounds of the corresponding dimensions of "
                f"operand. For dimension {update_window_dim}, updates "
                f"bound is {updates.shape[update_window_dim]}, operand "
                f"bound is {max_update_slice_sizes[i]}."
            )

    update_scatter_dims = [
        dim for dim in range(rank(updates)) if dim not in set(update_window_dims)
    ]

    scatter_dims_seen = 0
    for i in update_scatter_dims:
        if scatter_dims_seen == index_vector_dim:
            scatter_dims_seen += 1
        if updates.shape[i] != expanded_scatter_indices_shape[scatter_dims_seen]:
            raise TypeError(
                f"Bounds of the scatter dimensions of updates must be "
                f"the same as the bounds of the corresponding dimensions "
                f"of scatter indices. For scatter dimension {i}, updates "
                f"bound is {updates.shape[i]}, scatter_indices bound is "
                f"{expanded_scatter_indices_shape[scatter_dims_seen]}."
            )
        scatter_dims_seen += 1

    return operand.shape


# 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 _scatter_shape_rule(
    operand,
    scatter_indices,
    updates,
    *,
    update_jaxpr,
    update_consts,
    dimension_numbers,
    indices_are_sorted,
    unique_indices,
):
    """Validates the well-formedness of the ``dimension_numbers`` argument to
    Scatter.

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

    update_window_dims = dimension_numbers.update_window_dims
    inserted_window_dims = dimension_numbers.inserted_window_dims
    scatter_dims_to_operand_dims = dimension_numbers.scatter_dims_to_operand_dims
    # Note: in JAX, index_vector_dim is always computed as below, cf. the
    # documentation of the ScatterDimensionNumbers class.
    index_vector_dim = _rank(scatter_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(scatter_indices) < index_vector_dim or index_vector_dim < 0:
        raise TypeError(
            f"Scatter index leaf dimension must be within [0, "
            f"rank(scatter_indices) + 1). rank(scatter_indices) is "
            f"{_rank(scatter_indices)} and scatter index leaf "
            f"dimension is {index_vector_dim}."
        )

    expanded_scatter_indices_shape = list(scatter_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_scatter_indices_shape) == index_vector_dim:
        expanded_scatter_indices_shape.append(1)

    expected_updates_rank = (
        len(expanded_scatter_indices_shape) - 1 + len(update_window_dims)
    )

    if _rank(updates) != expected_updates_rank:
        raise TypeError(
            f"Updates tensor must be of rank {expected_updates_rank}; "
            f"got {_rank(updates)}."
        )

    # Validate update_window_dims
    _is_sorted(update_window_dims, "scatter", "update_window_dims")
    _no_duplicate_dims(update_window_dims, "scatter", "update_window_dims")
    _sorted_dims_in_range(
        update_window_dims, _rank(updates), "scatter", "update_window_dims"
    )

    # Validate inserted_window_dims
    _is_sorted(inserted_window_dims, "scatter", "inserted_window_dims")
    _no_duplicate_dims(inserted_window_dims, "scatter", "inserted_window_dims")
    _sorted_dims_in_range(
        inserted_window_dims, _rank(operand), "scatter", "inserted_window_dims"
    )

    # Validate window_size
    window_size = len(update_window_dims) + len(inserted_window_dims)
    if _rank(operand) != window_size:
        raise TypeError(
            f"Scatter op has window of size {window_size}; doesn't "
            f"match operand of rank {_rank(operand)}."
        )

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

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

    _no_duplicate_dims(
        scatter_dims_to_operand_dims, "scatter", "scatter_dims_to_operand_dims"
    )

    max_update_slice_sizes = [
        operand.shape[i]
        for i in range(len(operand.shape))
        if not i in set(inserted_window_dims)
    ]

    for i in range(len(update_window_dims)):
        update_window_dim = update_window_dims[i]
        if updates.shape[update_window_dim] > max_update_slice_sizes[i]:
            raise TypeError(
                f"Bounds of the window dimensions of updates must not "
                f"exceed the bounds of the corresponding dimensions of "
                f"operand. For dimension {update_window_dim}, updates "
                f"bound is {updates.shape[update_window_dim]}, operand "
                f"bound is {max_update_slice_sizes[i]}."
            )

    update_scatter_dims = [
        dim for dim in range(_rank(updates)) if dim not in set(update_window_dims)
    ]

    scatter_dims_seen = 0
    for i in update_scatter_dims:
        if scatter_dims_seen == index_vector_dim:
            scatter_dims_seen += 1
        if updates.shape[i] != expanded_scatter_indices_shape[scatter_dims_seen]:
            raise TypeError(
                f"Bounds of the scatter dimensions of updates must be "
                f"the same as the bounds of the corresponding dimensions "
                f"of scatter indices. For scatter dimension {i}, updates "
                f"bound is {updates.shape[i]}, scatter_indices bound is "
                f"{expanded_scatter_indices_shape[scatter_dims_seen]}."
            )
        scatter_dims_seen += 1

    return operand.shape