gurudesh/pytracebugs huggingface dataset:
# Buggy code
def result_type(*args):
    """Convenience function to apply Numpy argument dtype promotion."""
    return canonicalize_dtype(_result_type_raw(*args))


# Traceback (TypeError)
import objax
import jax
m = objax.nn.Sequential([objax.nn.Linear(3, 4)])

v = objax.random.uniform(((32, 3)))
jax.eval_shape(m, v)  # ShapeDtypeStruct(shape=(32, 4), dtype=float32)

d = objax.util.EasyDict(shape=(32, 3), dtype=v.dtype)
print(repr(d.shape), repr(d.dtype)) # (32, 3) dtype('float32')
print(hasattr(d, 'shape'), hasattr(d, 'dtype')) # True True

jax.eval_shape(m, d)
#
Traceback (most recent call last):
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3417, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "<ipython-input-24-d48e201250f2>", line 1, in <module>
jax.eval_shape(m, d)
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/api.py", line 2283, in eval_shape
*map(abstractify, args_flat))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/util.py", line 36, in safe_map
return list(map(f, *args))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/api.py", line 2279, in abstractify
return ShapedArray(np.shape(x), dtypes.result_type(x))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/dtypes.py", line 294, in result_type
return canonicalize_dtype(dtype(args[0]))
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/jax/dtypes.py", line 288, in dtype
return np.result_type(x)
File "<__array_function__ internals>", line 5, in result_type
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/numpy/core/_internal.py", line 64, in _usefields
names, formats, offsets, titles = _makenames_list(adict, align)
File "/usr/local/google/home/dberth/jax3/lib/python3.8/site-packages/numpy/core/_internal.py", line 40, in _makenames_list
format = dtype(obj[0], align=align)
TypeError: data type not understood

# Fixed code
def result_type(*args):
    """Convenience function to apply Numpy argument dtype promotion."""
    if len(args) == 0:
        raise ValueError("at least one array or dtype is required")
    return canonicalize_dtype(_result_type_raw(*args))
# Buggy code
def dynamic_slice_in_dim(operand, start_index, slice_size, axis=0):
    """Convenience wrapper around dynamic_slice applying to one dimension."""
    start_indices = [onp.array([0])] * operand.ndim
    slice_sizes = list(operand.shape)

    axis = int(axis)
    axis_size = _const(start_index, operand.shape[axis])
    start_indices[axis] = reshape(rem(start_index, axis_size), [1])
    slice_sizes[axis] = int(slice_size)

    start_indices = concatenate(start_indices, 0)
    return dynamic_slice(operand, start_indices, slice_sizes)


# Traceback (ValueError)
(1, 2, 2) (3, 2)
[[1. 1.]
[1. 1.]
[1. 1.]]
/Users/skainswo/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/mle_normal.py in <module>
31 print(A.shape, b.shape)
32 print(np.linalg.solve(A, b))
---> 33 print(jp.linalg.solve(A, b))

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/numpy/linalg.py in solve(a, b)
247   x = b if a_ndims == b_ndims else b[..., None]
248
--> 249   permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
250   x = x[..., permutation, :]
251

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax_linalg.py in lu_pivots_to_permutation(swaps, k)
462     return swaps, permutation
463
--> 464   n, = np.shape(swaps)
465   permutation = np.arange(k)
466   _, permutation = lax.fori_loop(

ValueError: too many values to unpack (expected 1)

# Fixed code
def dynamic_slice_in_dim(operand, start_index, slice_size, axis=0):
    """Convenience wrapper around dynamic_slice applying to one dimension."""
    start_indices = [onp.array([0], dtype=_dtype(start_index))] * operand.ndim
    slice_sizes = list(operand.shape)

    axis = int(axis)
    axis_size = _const(start_index, operand.shape[axis])
    start_indices[axis] = reshape(rem(start_index, axis_size), [1])
    slice_sizes[axis] = int(slice_size)

    start_indices = concatenate(start_indices, 0)
    return dynamic_slice(operand, start_indices, slice_sizes)
# Buggy code
def lu_pivots_to_permutation(swaps, k):
    """Converts the pivots (row swaps) returned by LU to a permutation."""

    def body_fn(i, loop_carry):
        swaps, permutation = loop_carry
        j = swaps[i]
        x, y = np.ravel(permutation[i]), np.ravel(permutation[j])
        permutation = lax.dynamic_update_index_in_dim(permutation, y, i, axis=0)
        permutation = lax.dynamic_update_index_in_dim(permutation, x, j, axis=0)
        return swaps, permutation

    (n,) = np.shape(swaps)
    permutation = np.arange(k)
    _, permutation = lax.fori_loop(
        onp.array(0, onp.int32), onp.array(n, onp.int32), body_fn, (swaps, permutation)
    )
    return permutation


# Traceback (ValueError)
(1, 2, 2) (3, 2)
[[1. 1.]
[1. 1.]
[1. 1.]]
/Users/skainswo/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/mle_normal.py in <module>
31 print(A.shape, b.shape)
32 print(np.linalg.solve(A, b))
---> 33 print(jp.linalg.solve(A, b))

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/numpy/linalg.py in solve(a, b)
247   x = b if a_ndims == b_ndims else b[..., None]
248
--> 249   permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
250   x = x[..., permutation, :]
251

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax_linalg.py in lu_pivots_to_permutation(swaps, k)
462     return swaps, permutation
463
--> 464   n, = np.shape(swaps)
465   permutation = np.arange(k)
466   _, permutation = lax.fori_loop(

ValueError: too many values to unpack (expected 1)

# Fixed code
def lu_pivots_to_permutation(swaps, m):
    """Converts the pivots (row swaps) returned by LU to a permutation.

    We build a permutation rather than applying `swaps` directly to the rows
    of a matrix because lax loops aren't differentiable.

    Args:
      swaps: an array of shape (..., k) of row swaps to perform
      m: the size of the output permutation. m should be >= k.
    Returns:
      An int32 array of shape (..., m).
    """
    assert len(swaps.shape) >= 1
    batch_dims = swaps.shape[:-1]
    k = swaps.shape[-1]

    def body_fn(i, permutation):
        j = swaps[..., i]
        iotas = np.ix_(*(lax.iota(np.int32, b) for b in batch_dims))
        x = permutation[..., i]
        y = permutation[iotas + (j,)]
        permutation = ops.index_update(permutation, ops.index[..., i], y)
        return ops.index_update(permutation, ops.index[iotas + (j,)], x)

    permutation = lax.broadcasted_iota(np.int32, batch_dims + (m,), len(batch_dims))
    return lax.fori_loop(
        onp.array(0, onp.int32), onp.array(k, onp.int32), body_fn, permutation
    )
# Buggy code
def body_fn(i, loop_carry):
    swaps, permutation = loop_carry
    j = swaps[i]
    x, y = np.ravel(permutation[i]), np.ravel(permutation[j])
    permutation = lax.dynamic_update_index_in_dim(permutation, y, i, axis=0)
    permutation = lax.dynamic_update_index_in_dim(permutation, x, j, axis=0)
    return swaps, permutation


# Traceback (ValueError)
(1, 2, 2) (3, 2)
[[1. 1.]
[1. 1.]
[1. 1.]]
/Users/skainswo/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/mle_normal.py in <module>
31 print(A.shape, b.shape)
32 print(np.linalg.solve(A, b))
---> 33 print(jp.linalg.solve(A, b))

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/numpy/linalg.py in solve(a, b)
247   x = b if a_ndims == b_ndims else b[..., None]
248
--> 249   permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
250   x = x[..., permutation, :]
251

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax_linalg.py in lu_pivots_to_permutation(swaps, k)
462     return swaps, permutation
463
--> 464   n, = np.shape(swaps)
465   permutation = np.arange(k)
466   _, permutation = lax.fori_loop(

ValueError: too many values to unpack (expected 1)

# Fixed code
def body_fn(i, permutation):
    j = swaps[..., i]
    iotas = np.ix_(*(lax.iota(np.int32, b) for b in batch_dims))
    x = permutation[..., i]
    y = permutation[iotas + (j,)]
    permutation = ops.index_update(permutation, ops.index[..., i], y)
    return ops.index_update(permutation, ops.index[iotas + (j,)], x)
# Buggy code
def _solve(a, b):
    _check_solve_shapes(a, b)

    # With custom_linear_solve, we can reuse the same factorization when
    # computing sensitivities. This is considerably faster.
    lu_, _, permutation = lu(lax.stop_gradient(a))
    custom_solve = partial(
        lax.custom_linear_solve,
        lambda x: _matvec_multiply(a, x),
        solve=lambda _, x: lu_solve(lu_, permutation, x, trans=0),
        transpose_solve=lambda _, x: lu_solve(lu_, permutation, x, trans=1),
    )
    if a.ndim == b.ndim + 1:
        # b.shape == [..., m]
        return custom_solve(b)
    else:
        # b.shape == [..., m, k]
        return api.vmap(custom_solve, b.ndim - 1, max(a.ndim, b.ndim) - 1)(b)


# Traceback (ValueError)
(1, 2, 2) (3, 2)
[[1. 1.]
[1. 1.]
[1. 1.]]
/Users/skainswo/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/mle_normal.py in <module>
31 print(A.shape, b.shape)
32 print(np.linalg.solve(A, b))
---> 33 print(jp.linalg.solve(A, b))

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/numpy/linalg.py in solve(a, b)
247   x = b if a_ndims == b_ndims else b[..., None]
248
--> 249   permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
250   x = x[..., permutation, :]
251

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax_linalg.py in lu_pivots_to_permutation(swaps, k)
462     return swaps, permutation
463
--> 464   n, = np.shape(swaps)
465   permutation = np.arange(k)
466   _, permutation = lax.fori_loop(

ValueError: too many values to unpack (expected 1)

# Fixed code
def _solve(a, b):
    _check_solve_shapes(a, b)

    # Broadcast leading dimensions of b to the shape of a, as is required by
    # custom_linear_solve.
    out_shape = tuple(
        d_a if d_b == 1 else d_b for d_a, d_b in zip(a.shape[:-1] + (1,), b.shape)
    )
    b = jnp.broadcast_to(b, out_shape)

    # With custom_linear_solve, we can reuse the same factorization when
    # computing sensitivities. This is considerably faster.
    lu_, _, permutation = lu(lax.stop_gradient(a))
    custom_solve = partial(
        lax.custom_linear_solve,
        lambda x: _matvec_multiply(a, x),
        solve=lambda _, x: lu_solve(lu_, permutation, x, trans=0),
        transpose_solve=lambda _, x: lu_solve(lu_, permutation, x, trans=1),
    )
    if a.ndim == b.ndim + 1:
        # b.shape == [..., m]
        return custom_solve(b)
    else:
        # b.shape == [..., m, k]
        return api.vmap(custom_solve, b.ndim - 1, max(a.ndim, b.ndim) - 1)(b)
# Buggy code
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}")


# 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 _is_sorted(dims, op_name, name):
    for i in range(1, len(dims)):
        if dims[i] < dims[i - 1]:
            raise TypeError(f"{name} in {op_name} op must be sorted; got {dims}")
# Buggy code
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}."
        )


# 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 _sorted_dims_in_range(dims, rank, op_name, 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 {op_name} op; valid range is "
            f"[0, {rank}); got: {invalid_dim}."
        )
# Buggy code
def _no_duplicate_dims(dims, name):
    if len(set(dims)) != len(dims):
        raise TypeError(f"{name} in scatter op must not repeat; got: {dims}.")


# 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 _no_duplicate_dims(dims, op_name, name):
    if len(set(dims)) != len(dims):
        raise TypeError(f"{name} in {op_name} op must not repeat; got: {dims}.")
# Buggy code
def _sort_translation_rule(c, *operands, dimension):
    out = xops.Sort(c, operands, dimension=dimension, is_stable=True)
    return out if len(operands) != 1 else xops.Tuple(c, [out])


# Traceback (RuntimeError)
❯ pipenv run python -m research.lqr
/Users/skainswo/dev/jax/jax/lib/xla_bridge.py:116: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
File "/usr/local/Cellar/python/3.7.6_1/Frameworks/Python.framework/Versions/3.7/lib/python3.7/runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "/usr/local/Cellar/python/3.7.6_1/Frameworks/Python.framework/Versions/3.7/lib/python3.7/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/Users/skainswo/dev/research/research/lqr.py", line 50, in <module>
_test_lqr1(2)
File "/Users/skainswo/dev/research/research/lqr.py", line 45, in _test_lqr1
actual = lqr_continuous_time_infinite_horizon(A, B, Q, R, N)
File "/Users/skainswo/dev/research/research/lqr.py", line 26, in lqr_continuous_time_infinite_horizon
argsort = jp.argsort(eigvals)
File "/Users/skainswo/dev/jax/jax/numpy/lax_numpy.py", line 2886, in argsort
_, perm = lax.sort_key_val(a, iota, dimension=axis)
File "/Users/skainswo/dev/jax/jax/lax/lax.py", line 1190, in sort_key_val
result = sort_key_val_p.bind(keys, values, dimension=dimension)
File "/Users/skainswo/dev/jax/jax/core.py", line 211, in bind
return self.impl(*args, **kwargs)
File "/Users/skainswo/dev/jax/jax/interpreters/xla.py", line 217, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *map(arg_spec, args), **params)
File "/Users/skainswo/dev/jax/jax/interpreters/xla.py", line 254, in xla_primitive_callable
compiled = backend.compile(built_c, compile_options=options)
RuntimeError: Unimplemented: complex comparison 'LT'

# Fixed code
def _sort_translation_rule(c, *operands, dimension):
    types = [c.get_shape(x).xla_element_type() for x in operands]
    subc = xla_bridge.make_computation_builder("sort_lt_comparator")
    params = [
        xb.parameter(subc, 2 * i + j, xc.Shape.array_shape(typ, ()))
        for i, typ in enumerate(types)
        for j in range(2)
    ]
    result = xla.lower_fun(_sort_lt_comparator, multiple_results=False)(subc, *params)
    comparator = subc.build(result)
    out = xops.Sort(
        c, operands, dimension=dimension, is_stable=True, comparator=comparator
    )
    return out if len(operands) != 1 else xops.Tuple(c, [out])
# Buggy code
def wrapped_fun(args, logical_env):
    args_flat, in_tree = tree_flatten(args)
    if in_tree != in_shapes_tree:
        raise TypeError("pytree mismatch")
    logical_env = {unique_ids[name]: val for name, val in logical_env.items()}
    in_shapes = map(masking.finalize_spec, in_specs, map(onp.shape, args_flat))
    padded_env = _bind_shapes(in_shapes, [x.shape for x in args_flat])
    f = lu.wrap_init(fun)
    flat_fun, out_tree = flatten_fun_nokwargs(f, in_tree)
    outs, out_shapes_ = masking.mask_fun(
        flat_fun, logical_env, padded_env, args_flat, in_shapes
    )
    if not out_tree() == out_shapes_tree:
        raise TypeError("pytree mismatch")
    out_shapes = map(masking.finalize_spec, out_specs, map(onp.shape, outs))
    if not out_shapes == list(out_shapes_):
        raise masking.ShapeError
    if not all(
        onp.shape(out) == masking.eval_shape_expr(padded_env, expr)
        for out, expr in zip(outs, out_shapes)
    ):
        raise masking.ShapeError
    return tree_unflatten(out_tree(), outs)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def wrapped_fun(args, logical_env):
    args_flat, in_tree = tree_flatten(args)
    if in_tree != in_shapes_tree:
        raise TypeError("pytree mismatch")
    logical_env = {unique_ids[name]: val for name, val in logical_env.items()}
    in_shapes = map(masking.finalize_spec, in_specs, map(onp.shape, args_flat))
    padded_env = _bind_shapes(in_shapes, [x.shape for x in args_flat])
    f = lu.wrap_init(fun)
    flat_fun, out_tree = flatten_fun_nokwargs(f, in_tree)
    outs, out_shapes_ = masking.mask_fun(
        flat_fun, logical_env, padded_env, args_flat, in_shapes
    )
    if not out_tree() == out_shapes_tree:
        raise TypeError("pytree mismatch")
    out_shapes = map(masking.finalize_spec, out_specs, map(onp.shape, outs))
    if not out_shapes == list(out_shapes_):
        raise masking.ShapeError
    if not all(
        onp.shape(out) == eval_polymorphic_shape(shape, padded_env)
        for out, shape in zip(outs, out_shapes)
    ):
        raise masking.ShapeError
    return tree_unflatten(out_tree(), outs)
# Buggy code
def _remap_ids(names, shape_spec):
    ShapeSpec, Poly, Mon = masking.ShapeSpec, masking.Poly, masking.Mon
    mdim = masking.monomorphic_dim
    return ShapeSpec(
        Poly(
            {
                Mon({names[id]: deg for id, deg in mon.items()}): coeff
                for mon, coeff in poly.items()
            }
        )
        if poly is not mdim
        else mdim
        for poly in shape_spec
    )


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def _remap_ids(names, shape_spec):
    return masking.ShapeSpec(
        Poly(
            {
                Mon({names[id]: deg for id, deg in mon.items()}): coeff
                for mon, coeff in poly.items()
            }
        )
        if poly is not masking._monomorphic_dim
        else masking._monomorphic_dim
        for poly in shape_spec
    )
# Buggy code
def _bind_shapes(shape_exprs, shapes):
    env = {}
    for shape_expr, shape in zip(shape_exprs, shapes):
        for poly, d in zip(shape_expr, shape):
            if ensure_poly(poly).is_constant:
                continue
            else:
                ((binder,),) = poly  # TODO generalize to handle striding
                if env.setdefault(binder, d) != d:
                    raise masking.ShapeError
    return env


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def _bind_shapes(shape_exprs, shapes):
    env = {}
    for shape_expr, shape in zip(shape_exprs, shapes):
        for poly, d in zip(shape_expr, shape):
            if type(poly) is not Poly or poly.is_constant:
                continue
            else:
                ((binder,),) = poly  # TODO generalize to handle striding
                if env.setdefault(binder, d) != d:
                    raise masking.ShapeError
    return env
# Buggy code
def shapecheck(in_shapes, out_shape, fun):
    _check_callable(fun)
    in_shapes, in_tree = tree_flatten(in_shapes)
    in_shapes = map(masking.parse_spec, in_shapes)
    out_shapes, out_tree = tree_flatten(out_shape)
    out_shapes = map(masking.parse_spec, out_shapes)
    flat_fun, out_tree_ = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
    out_shapes_ = masking.shapecheck(flat_fun, in_shapes)
    if out_tree != out_tree_():
        raise TypeError("pytree mismatch")
    if not all(map(_shape_spec_consistent, out_shapes, out_shapes_)):
        raise masking.ShapeError
    return fun


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def shapecheck(in_shapes, out_shape, fun: Callable):
    _check_callable(fun)
    in_shapes, in_tree = tree_flatten(in_shapes)
    in_shapes = map(masking.parse_spec, in_shapes)
    out_shapes, out_tree = tree_flatten(out_shape)
    out_shapes = map(masking.parse_spec, out_shapes)
    flat_fun, out_tree_ = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
    avals = map(partial(ShapedArray, dtype=onp.float32), in_shapes)
    out_shapes_ = [o.shape for o in pe.abstract_eval_fun(flat_fun.call_wrapped, *avals)]
    if out_tree != out_tree_():
        raise TypeError("pytree mismatch")
    if not all(map(masking._shape_spec_consistent, out_shapes, out_shapes_)):
        raise masking.ShapeError
    return fun
# Buggy code
def extend_shape_envs(logical_env, padded_env):
    global shape_envs
    new_logical = dict(it.chain(shape_envs.logical.items(), logical_env.items()))
    new_padded = dict(it.chain(shape_envs.padded.items(), padded_env.items()))
    shape_envs, prev = ShapeEnvs(new_logical, new_padded), shape_envs
    yield
    shape_envs = prev


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def extend_shape_envs(logical_env, padded_env):
    global shape_envs
    new_logical = dict(chain(shape_envs.logical.items(), logical_env.items()))
    new_padded = dict(chain(shape_envs.padded.items(), padded_env.items()))
    shape_envs, prev = ShapeEnvs(new_logical, new_padded), shape_envs
    yield
    shape_envs = prev
# Buggy code
def shape_as_value(expr):
    if type(expr) is tuple and is_polymorphic(expr):
        return tuple(
            eval_dim_expr(shape_envs.logical, d) if type(d) is Poly else d for d in expr
        )
    else:
        return expr


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def shape_as_value(shape):
    return eval_polymorphic_shape(shape, shape_envs.logical)
# Buggy code
def padded_shape_as_value(expr):
    if type(expr) is tuple and is_polymorphic(expr):
        return tuple(
            eval_dim_expr(shape_envs.padded, d) if type(d) is Poly else d for d in expr
        )
    else:
        return expr


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def padded_shape_as_value(shape):
    return eval_polymorphic_shape(shape, shape_envs.padded)
# Buggy code
def __init__(self, coeffs):
    # Makes sure Polynomials are always in canonical form to simplify operators:
    coeffs = {mon: coeff for mon, coeff in coeffs.items() if coeff != 0}
    coeffs = {Mon(): 0} if len(coeffs) == 0 else coeffs
    super().__init__(coeffs)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __init__(self, coeffs):
    # Makes sure Polynomials are always in canonical form
    coeffs = {mon: op.index(coeff) for mon, coeff in coeffs.items() if coeff != 0}
    coeffs = coeffs or {Mon(): 0}
    super().__init__(coeffs)
# Buggy code
def __add__(self, other):
    coeffs = self.copy()

    for mon, coeff in ensure_poly(other).items():
        coeffs[mon] = coeffs.get(mon, 0) + coeff

    return Poly(coeffs)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __add__(self, other):
    coeffs = self.copy()
    for mon, coeff in _ensure_poly(other).items():
        coeffs[mon] = coeffs.get(mon, 0) + coeff
    return Poly(coeffs)
# Buggy code
def __mul__(self, other):
    coeffs = dict()
    for (mon1, coeff1), (mon2, coeff2) in it.product(
        self.items(), ensure_poly(other).items()
    ):
        mon = Mon(mon1 + mon2)  # add monomials' id degrees
        coeff = coeff1 * coeff2  # multiply integer coeffs
        coeffs[mon] = coeffs.get(mon, 0) + coeff  # accumulate coeffs

    return Poly(coeffs)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __mul__(self, other):
    other = _ensure_poly(other)
    coeffs = {}
    for (mon1, coeff1), (mon2, coeff2) in product(self.items(), other.items()):
        mon = mon1 * mon2
        coeffs[mon] = coeffs.get(mon, 0) + coeff1 * coeff2
    return Poly(coeffs)
# Buggy code
def __rmul__(self, other):
    return self * other


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __rmul__(self, other):
    return self * other  # multiplication commutes
# Buggy code
def __radd__(self, other):
    return self + other


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __radd__(self, other):
    return self + other  # addition commutes
# Buggy code
def __rsub__(self, other):
    return self + -other


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __rsub__(self, other):
    return _ensure_poly(other) - self
# Buggy code
def __divmod__(self, divisor):
    if self.is_constant:
        q, r = divmod(int(self), divisor)

        return constant_poly(q), r

    def divided(count):
        q, r = divmod(count, divisor)
        if r != 0:
            raise ValueError(
                "shapecheck currently only supports strides "
                "that exactly divide the strided axis length."
            )
        return q

    return Poly(
        {
            k: coeff // divisor if k.degree == 0 else divided(coeff)
            for k, coeff in self.items()
        }
    ), self[Mon()] % divisor


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __divmod__(self, divisor):
    if self.is_constant:
        return divmod(int(self), divisor)
    else:

        def divided(count):
            q, r = divmod(count, divisor)
            if r != 0:
                raise ValueError(
                    "shapecheck currently only supports strides "
                    "that exactly divide the strided axis length."
                )
            return q

        return Poly(
            {
                k: coeff // divisor if k.degree == 0 else divided(coeff)
                for k, coeff in self.items()
            }
        ), self.get(Mon(), 0) % divisor
# Buggy code
def __hash__(self):
    return hash(super())


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __hash__(self):
    return hash(tuple(sorted(self.items())))
# Buggy code
def __eq__(self, other):
    return super().__eq__(ensure_poly(other))


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __eq__(self, other):
    return super().__eq__(_ensure_poly(other))
# Buggy code
def __ge__(self, other):
    other = ensure_poly(other)

    if other.is_constant and self.is_constant:
        return int(self) >= int(other)

    if other.is_constant and int(other) <= 1:
        # Assume polynomials > 0, allowing to use shape rules of binops, conv:
        return True

    if self.is_constant and int(self) <= 0:
        return False  # See above.

    if self == other:
        return True

    raise ValueError(
        'Polynomials comparison "{} >= {}" is inconclusive.'.format(self, other)
    )


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __ge__(self, other):
    other = _ensure_poly(other)

    if other.is_constant and self.is_constant:
        return int(self) >= int(other)
    elif other.is_constant and int(other) <= 1:
        # Assume nonzero polynomials are positive, allows use in shape rules
        return True
    elif self.is_constant and int(self) <= 0:
        return False  # See above.
    elif self == other:
        return True

    raise ValueError(
        'Polynomials comparison "{} >= {}" is inconclusive.'.format(self, other)
    )
# Buggy code
def __le__(self, other):
    return ensure_poly(other) >= self


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __le__(self, other):
    return _ensure_poly(other) >= self
# Buggy code
def __gt__(self, other):
    return not (ensure_poly(other) >= self)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __gt__(self, other):
    return not (_ensure_poly(other) >= self)
# Buggy code
def __int__(self):
    assert self.is_constant

    return int(next(iter(self.values())))


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __int__(self):
    assert self.is_constant
    return op.index(next(iter(self.values())))
# Buggy code
def finalize_spec(spec, shape):
    return tuple(
        parse_lit(d) if e is monomorphic_dim else e for e, d in zip(spec, shape)
    )


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def finalize_spec(spec, shape):
    return tuple(
        _parse_lit(d) if e is _monomorphic_dim else e for e, d in zip(spec, shape)
    )
# Buggy code
def parse_spec(spec=""):
    if not spec:
        return ShapeSpec(())
    if spec[0] == "(":
        if spec[-1] != ")":
            raise ShapeSyntaxError(spec)
        spec = spec[1:-1]
    dims = map(parse_dim, spec.replace(" ", "").strip(",").split(","))
    return ShapeSpec(dims)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def parse_spec(spec=""):
    if not spec:
        return ShapeSpec(())
    if spec[0] == "(":
        if spec[-1] != ")":
            raise ShapeSyntaxError(spec)
        spec = spec[1:-1]
    dims = map(_parse_dim, spec.replace(" ", "").strip(",").split(","))
    return ShapeSpec(dims)
# Buggy code
def __getitem__(self, idx):
    if type(idx) is tuple:
        return parse_spec("(" + ",".join(map(str, idx)) + ")")
    else:
        return parse_spec(str(idx))


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def __getitem__(self, idx):
    return parse_spec(
        ("(" + ",".join(map(str, idx)) + ")") if type(idx) is tuple else str(idx)
    )
# Buggy code
def is_pure(self):
    return all(ensure_poly(poly).is_constant for poly in self.shape_expr)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "/Users/necula/Source/jax/jax/interpreters/xla.py", line 230, in primitive_computation
return c.Build()
File "/Users/necula/Source/jax/jax/lib/xla_bridge.py", line 281, in Build
*args, **kwargs)
File "/Users/necula/Source/jax/build/jaxlib/xla_client.py", line 734, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 6), got 10.:

# Fixed code
def is_pure(self):
    return all(type(poly) is not Poly or poly.is_constant for poly in self.shape_expr)