gurudesh/pytracebugs huggingface dataset:
# Buggy code
def process_primitive(self, primitive, tracers, params):
    vals, shape_exprs = unzip2((t.val, t.shape_expr) for t in tracers)
    if primitive in shape_parameterized_primitive_rules:
        rule = shape_parameterized_primitive_rules[primitive]
        out, out_shape = rule(shape_envs, vals, shape_exprs, **params)
    else:
        out_shape = shape_rules[primitive](*(t.aval for t in tracers), **params)
        logical_shapes = map(partial(eval_shape_expr, shape_envs.logical), shape_exprs)
        out = masking_rules[primitive](vals, logical_shapes, **params)
    if not primitive.multiple_results:
        return MaskTracer(self, out, out_shape)
    else:
        return map(partial(MaskTracer, self), out, out_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 process_primitive(self, primitive, tracers, params):
    vals, shape_exprs = unzip2((t.val, t.shape_expr) for t in tracers)
    if primitive in shape_parameterized_primitive_rules:
        rule = shape_parameterized_primitive_rules[primitive]
        out, out_shape = rule(shape_envs, vals, shape_exprs, **params)
    else:
        avals = [t.aval for t in tracers]
        out = primitive.abstract_eval(*avals, **params)
        out_shape = [o.shape for o in out] if primitive.multiple_results else out.shape
        logical_shapes = map(
            partial(eval_polymorphic_shape, values_dict=shape_envs.logical), shape_exprs
        )
        out = masking_rules[primitive](vals, logical_shapes, **params)
    if not primitive.multiple_results:
        return MaskTracer(self, out, out_shape)
    else:
        return map(partial(MaskTracer, self), out, out_shape)
# Buggy code
def __init__(self, trace, shape_expr):
    self._trace = trace
    self.shape_expr = 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 __init__(self, trace, val, shape_expr):
    self._trace = trace
    self.val = val
    self.shape_expr = shape_expr
# Buggy code
def aval(self):
    return ShapedArray(self.shape_expr, None)


# 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 aval(self):
    return ShapedArray(self.shape_expr, self.val.dtype)
# Buggy code
def full_lower(self):
    return 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 full_lower(self):
    if self.is_pure():
        return core.full_lower(self.val)
    else:
        return self
# Buggy code
def pure(self, val):
    return ShapeCheckTracer(self, onp.shape(val))


# 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 pure(self, val):
    return MaskTracer(self, val, onp.shape(val))
# Buggy code
def lift(self, val):
    return ShapeCheckTracer(self, onp.shape(val))


# 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 lift(self, val):
    return MaskTracer(self, val, onp.shape(val))
# Buggy code
def sublift(self, val):
    return ShapeCheckTracer(self, val.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 sublift(self, val):
    return MaskTracer(self, val.val, val.shape_expr)
# Buggy code
def process_primitive(self, primitive, tracers, params):
    avals = [t.aval for t in tracers]
    shape_rule = shape_rules.get(primitive)
    if shape_rule is None:
        raise NotImplementedError(
            "Shape rule for {} not implemented yet.".format(primitive)
        )
    out_shape = shape_rule(*avals, **params)
    return ShapeCheckTracer(self, out_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 process_primitive(self, primitive, tracers, params):
    vals, shape_exprs = unzip2((t.val, t.shape_expr) for t in tracers)
    if primitive in shape_parameterized_primitive_rules:
        rule = shape_parameterized_primitive_rules[primitive]
        out, out_shape = rule(shape_envs, vals, shape_exprs, **params)
    else:
        avals = [t.aval for t in tracers]
        out = primitive.abstract_eval(*avals, **params)
        out_shape = [o.shape for o in out] if primitive.multiple_results else out.shape
        logical_shapes = map(
            partial(eval_polymorphic_shape, values_dict=shape_envs.logical), shape_exprs
        )
        out = masking_rules[primitive](vals, logical_shapes, **params)
    if not primitive.multiple_results:
        return MaskTracer(self, out, out_shape)
    else:
        return map(partial(MaskTracer, self), out, out_shape)
# Buggy code
def process_call(self, call_primitive, f: lu.WrappedFun, tracers, params):
    # TODO apply proper subtrace:
    return map(self.full_raise, f.call_wrapped(*tracers))


# 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 process_call(self, call_primitive, f: lu.WrappedFun, tracers, params):
    raise NotImplementedError  # TODO mask-of-jit
# Buggy code
def broadcast_shapes(*shapes):
    """Returns the shape that results from NumPy broadcasting of `shapes`."""
    if len(shapes) == 1:
        return shapes[0]
    ndim = _max(len(shape) for shape in shapes)
    shapes = onp.array([(1,) * (ndim - len(shape)) + shape for shape in shapes])
    is_zero = onp.any(shapes == 0, axis=0)
    max_shape = onp.max(shapes, axis=0)
    result_shape = onp.where(is_zero, 0, max_shape)
    if not onp.all((shapes == result_shape) | (shapes == 1)):
        raise ValueError(
            "Incompatible shapes for broadcasting: {}".format(tuple(map(tuple, shapes)))
        )
    return tuple(map(_canonicalize_dimension, result_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 broadcast_shapes(*shapes):
    """Returns the shape that results from NumPy broadcasting of `shapes`."""
    if len(shapes) == 1:
        return shapes[0]
    ndim = _max(len(shape) for shape in shapes)
    shapes = onp.array([(1,) * (ndim - len(shape)) + shape for shape in shapes])
    is_zero = onp.any(shapes == 0, axis=0)
    max_shape = onp.max(shapes, axis=0)
    result_shape = onp.where(is_zero, 0, max_shape)
    if not onp.all((shapes == result_shape) | (shapes == 1)):
        raise ValueError(
            "Incompatible shapes for broadcasting: {}".format(tuple(map(tuple, shapes)))
        )
    return canonicalize_shape(result_shape)
# Buggy code
def iota(dtype: DType, size: int) -> Array:
    """Wraps XLA's `Iota
    <https://www.tensorflow.org/xla/operation_semantics#iota>`_
    operator.
    """
    size = int(size)
    dtype = dtypes.canonicalize_dtype(dtype)
    lazy_expr = lazy.iota(dtype, size)
    aval = ShapedArray((size,), dtype)
    return xla.DeviceArray(aval, None, lazy_expr, xla.DeviceConstant())


# 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 iota(dtype: DType, size: int) -> Array:
    """Wraps XLA's `Iota
    <https://www.tensorflow.org/xla/operation_semantics#iota>`_
    operator.
    """
    size = size if type(size) is masking.Poly else int(size)
    shape = canonicalize_shape((size,))
    dtype = dtypes.canonicalize_dtype(dtype)
    lazy_expr = lazy.iota(dtype, shape[0])
    aval = ShapedArray(shape, dtype)
    return xla.DeviceArray(aval, None, lazy_expr, xla.DeviceConstant())
# Buggy code
def standard_primitive(shape_rule, dtype_rule, name, translation_rule=None):
    prim = Primitive(name)
    prim.def_impl(partial(xla.apply_primitive, prim))
    prim.def_abstract_eval(
        partial(standard_abstract_eval, prim, shape_rule, dtype_rule)
    )
    xla.translations[prim] = translation_rule or partial(standard_translate, name)
    masking.shape_rules[prim] = shape_rule
    return prim


# 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 standard_primitive(shape_rule, dtype_rule, name, translation_rule=None):
    prim = Primitive(name)
    prim.def_impl(partial(xla.apply_primitive, prim))
    prim.def_abstract_eval(
        partial(standard_abstract_eval, prim, shape_rule, dtype_rule)
    )
    xla.translations[prim] = translation_rule or partial(standard_translate, name)
    return prim
# Buggy code
def _check_shapelike(fun_name, arg_name, obj):
    """Check that `obj` is a shape-like value (e.g. tuple of nonnegative ints)."""
    if type(obj) is tuple and masking.is_polymorphic(obj):
        return obj
    if not isinstance(obj, (tuple, list, onp.ndarray)):
        msg = "{} {} must be of type tuple/list/ndarray, got {}."
        raise TypeError(msg.format(fun_name, arg_name, type(obj)))
    # bool(obj) for an ndarray raises an error, so we check len
    if not len(obj):  # pylint: disable=g-explicit-length-test
        return
    obj_arr = onp.array(obj)
    if obj_arr.ndim != 1:
        msg = "{} {} must be rank 1, got {}."
        raise TypeError(msg.format(obj_arr.ndim))
    if not dtypes.issubdtype(obj_arr.dtype, onp.integer):
        msg = "{} {} must have every element be an integer type, got {}."
        raise TypeError(msg.format(fun_name, arg_name, tuple(map(type, obj))))
    if not (obj_arr >= 0).all():
        msg = "{} {} must have every element be nonnegative, got {}."
        raise TypeError(msg.format(fun_name, arg_name, obj))


# 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 _check_shapelike(fun_name, arg_name, obj):
    """Check that `obj` is a shape-like value (e.g. tuple of nonnegative ints)."""
    if not isinstance(obj, (tuple, list, onp.ndarray)):
        msg = "{} {} must be of type tuple/list/ndarray, got {}."
        raise TypeError(msg.format(fun_name, arg_name, type(obj)))
    # bool(obj) for an ndarray raises an error, so we check len
    if not len(obj):  # pylint: disable=g-explicit-length-test
        return
    obj_arr = onp.array(obj)
    if obj_arr.ndim != 1:
        msg = "{} {} must be rank 1, got {}."
        raise TypeError(msg.format(obj_arr.ndim))
    try:
        canonicalize_shape(obj_arr)
    except TypeError:
        msg = "{} {} must have every element be an integer type, got {}."
        raise TypeError(msg.format(fun_name, arg_name, tuple(map(type, obj))))
    if not (obj_arr >= 0).all():
        msg = "{} {} must have every element be nonnegative, got {}."
        raise TypeError(msg.format(fun_name, arg_name, obj))
# Buggy code
def _scan_masking_rule(
    shape_envs,
    padded_vals,
    shape_exprs,
    forward,
    length,
    jaxpr,
    num_consts,
    num_carry,
    linear,
):
    out_shape = _scan_shape_rule(
        shape_exprs, forward, length, jaxpr, num_consts, num_carry, linear
    )
    dynamic_length = masking.eval_dim_expr(shape_envs.logical, length)
    masked_jaxpr = _masked_scan_jaxpr(jaxpr, num_consts, num_carry)
    consts, init, xs = split_list(padded_vals, [num_consts, num_carry])
    (max_length,) = {x.shape[0] for x in xs}
    const_linear, init_linear, xs_linear = split_list(linear, [num_consts, num_carry])
    out_vals = scan_p.bind(
        *itertools.chain([dynamic_length] + consts, [0], init, xs),
        forward=forward,
        length=max_length,
        jaxpr=masked_jaxpr,
        num_consts=1 + num_consts,
        num_carry=1 + num_carry,
        linear=tuple([False] + const_linear + [False] + init_linear + xs_linear),
    )
    return out_vals[1:], out_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 _scan_masking_rule(
    shape_envs,
    padded_vals,
    shape_exprs,
    forward,
    length,
    jaxpr,
    num_consts,
    num_carry,
    linear,
):
    out_shape = _scan_shape_rule(
        shape_exprs, forward, length, jaxpr, num_consts, num_carry, linear
    )
    dynamic_length = length.evaluate(shape_envs.logical)
    masked_jaxpr = _masked_scan_jaxpr(jaxpr, num_consts, num_carry)
    consts, init, xs = split_list(padded_vals, [num_consts, num_carry])
    (max_length,) = {x.shape[0] for x in xs}
    const_linear, init_linear, xs_linear = split_list(linear, [num_consts, num_carry])
    out_vals = scan_p.bind(
        *itertools.chain([dynamic_length] + consts, [0], init, xs),
        forward=forward,
        length=max_length,
        jaxpr=masked_jaxpr,
        num_consts=1 + num_consts,
        num_carry=1 + num_carry,
        linear=tuple([False] + const_linear + [False] + init_linear + xs_linear),
    )
    return out_vals[1:], out_shape
# Buggy code
def broadcast_to(arr, shape):
    """Like Numpy's broadcast_to but doesn't necessarily return views."""
    arr = arr if isinstance(arr, ndarray) else array(arr)
    shape = tuple(map(int, shape))  # check that shape is concrete
    arr_shape = _shape(arr)
    if arr_shape == shape:
        return arr
    else:
        nlead = len(shape) - len(arr_shape)
        compatible = onp.equal(arr_shape, shape[nlead:]) | onp.equal(arr_shape, 1)
        if nlead < 0 or not onp.all(compatible):
            msg = "Incompatible shapes for broadcasting: {} and requested shape {}"
            raise ValueError(msg.format(arr_shape, shape))
        (diff,) = onp.where(onp.not_equal(shape[nlead:], arr_shape))
        new_dims = tuple(range(nlead)) + tuple(nlead + diff)
        kept_dims = tuple(onp.delete(onp.arange(len(shape)), new_dims))
        return lax.broadcast_in_dim(squeeze(arr, diff), shape, kept_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 broadcast_to(arr, shape):
    """Like Numpy's broadcast_to but doesn't necessarily return views."""
    arr = arr if isinstance(arr, ndarray) else array(arr)
    shape = canonicalize_shape(shape)  # check that shape is concrete
    arr_shape = _shape(arr)
    if arr_shape == shape:
        return arr
    else:
        nlead = len(shape) - len(arr_shape)
        compatible = onp.equal(arr_shape, shape[nlead:]) | onp.equal(arr_shape, 1)
        if nlead < 0 or not onp.all(compatible):
            msg = "Incompatible shapes for broadcasting: {} and requested shape {}"
            raise ValueError(msg.format(arr_shape, shape))
        (diff,) = onp.where(onp.not_equal(shape[nlead:], arr_shape))
        new_dims = tuple(range(nlead)) + tuple(nlead + diff)
        kept_dims = tuple(onp.delete(onp.arange(len(shape)), new_dims))
        return lax.broadcast_in_dim(squeeze(arr, diff), shape, kept_dims)
# Buggy code
def _normalize_index(index, axis_size):
    """Normalizes an index value in the range [-N, N) to the range [0, N)."""
    return lax.select(
        lax.lt(index, _constant_like(index, 0)),
        lax.add(index, _constant_like(index, axis_size)),
        index,
    )


# 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 _normalize_index(index, axis_size):
    """Normalizes an index value in the range [-N, N) to the range [0, N)."""
    if type(axis_size) is Poly:
        return index + axis_size if index < 0 else index

    return lax.select(
        lax.lt(index, _constant_like(index, 0)),
        lax.add(index, _constant_like(index, axis_size)),
        index,
    )
# Buggy code
def _static_idx(idx, size):
    """Helper function to compute the static slice start/limit/stride values."""
    assert isinstance(idx, slice)
    start, stop, step = idx.indices(size)
    if (step < 0 and stop >= start) or (step > 0 and start >= stop):
        return 0, 0, 1, False  # sliced to size zero

    if step > 0:
        return start, stop, step, False
    else:
        k = (start - stop - 1) % (-step)
        return stop + k + 1, start + 1, -step, True


# 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 _static_idx(idx: slice, size: Union[int, Poly]):
    """Helper function to compute the static slice start/limit/stride values."""
    if _any(type(s) is Poly for s in (idx.start, idx.stop, idx.step, size)):
        start, stop, step = _polymorphic_slice_indices(idx, size)
    elif isinstance(size, int):
        start, stop, step = idx.indices(size)
    else:
        raise TypeError(size)

    if type(start) is not Poly and type(stop) is not Poly:
        if (step < 0 and stop >= start) or (step > 0 and start >= stop):
            return 0, 0, 1, False  # sliced to size zero

    if step > 0:
        return start, stop, step, False
    else:
        k = (start - stop - 1) % (-step)
        return stop + k + 1, start + 1, -step, True
# Buggy code
def _check_shape(name, shape, *param_shapes):
    try:
        shape = tuple(map(int, shape))
    except TypeError as err:
        msg = "{} requires a concrete tuple of integers as shape argument, got {}."
        raise ValueError(msg.format(name, shape)) from err
    if param_shapes:
        shape_ = lax.broadcast_shapes(shape, *param_shapes)
        if shape != shape_:
            msg = (
                "{} parameter shapes must be broadcast-compatible with shape "
                "argument, and the result of broadcasting the shapes must equal "
                "the shape argument, but got result {} for shape argument {}."
            )
            raise ValueError(msg.format(name, shape_, 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 _check_shape(name, shape, *param_shapes):
    shape = abstract_arrays.canonicalize_shape(shape)

    if param_shapes:
        shape_ = lax.broadcast_shapes(shape, *param_shapes)
        if shape != shape_:
            msg = (
                "{} parameter shapes must be broadcast-compatible with shape "
                "argument, and the result of broadcasting the shapes must equal "
                "the shape argument, but got result {} for shape argument {}."
            )
            raise ValueError(msg.format(name, shape_, shape))
# Buggy code
def __getitem__(self, idx):
    if self._npy_value is None and type(idx) is int:
        ids = self._ids()
        device_buffer = self.device_buffers[ids[idx]]
        aval = ShapedArray(self.aval.shape[1:], self.aval.dtype)
        handler = xla.aval_to_result_handler(aval)
        return handler(device_buffer)
    else:
        return super(ShardedDeviceArray, self).__getitem__(idx)


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
169     msg = "primitive arguments must be colocated on the same device, got {}"
170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
172   else:
173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)

# Fixed code
def __getitem__(self, idx):
    if self._npy_value is None and type(idx) is int:
        ids = self._ids()
        device_buffer = self.device_buffers[ids[idx]]
        aval = ShapedArray(self.aval.shape[1:], self.aval.dtype)
        handler = xla.aval_to_result_handler(None, aval)
        return handler(device_buffer)
    else:
        return super(ShardedDeviceArray, self).__getitem__(idx)
# Buggy code
def aval_to_result_handler(aval):
    try:
        return xla_result_handlers[type(aval)](aval)
    except KeyError:
        raise TypeError("No xla_result_handler for type: {}".format(type(aval)))


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
169     msg = "primitive arguments must be colocated on the same device, got {}"
170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
172   else:
173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)

# Fixed code
def aval_to_result_handler(device, aval):
    try:
        return xla_result_handlers[type(aval)](device, aval)
    except KeyError:
        raise TypeError("No xla_result_handler for type: {}".format(type(aval)))
# Buggy code
def array_result_handler(aval):
    return partial(DeviceArray, raise_to_shaped(aval))


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
169     msg = "primitive arguments must be colocated on the same device, got {}"
170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
172   else:
173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)

# Fixed code
def array_result_handler(device, aval):
    return partial(DeviceArray, raise_to_shaped(aval), device)
# Buggy code
def xla_primitive_callable(prim, *arg_specs, **params):
    avals, devices = unzip2(arg_specs)
    # TODO(mattjj): make Device hashable instead of handling pairs here
    try:
        (device,) = set(d for d in devices if d is not None) or (None,)
    except ValueError:
        msg = "primitive arguments must be colocated on the same device, got {}"
        names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
        raise ValueError(msg.format(", ".join(names)))
    else:
        all_devices = it.chain(xb.devices(), xb.devices("cpu"))
        device = device and next(d for d in all_devices if (type(d), d.id) == device)
    backend = xb.get_device_backend(device)
    aval_out = prim.abstract_eval(*avals, **params)
    if prim.multiple_results:
        handlers = tuple(map(aval_to_result_handler, aval_out))
        handle_result = lambda xs: tuple(
            h(x) for h, x in zip(handlers, xs.destructure())
        )
    else:
        handle_result = aval_to_result_handler(aval_out)
    tuple_args = len(avals) > 100
    built_c = primitive_computation(prim, backend, tuple_args, *avals, **params)
    options = xb.get_compile_options(device_assignment=(device.id,) if device else None)
    compiled = built_c.Compile(compile_options=options, backend=backend)
    return partial(
        _execute_compiled_primitive, prim, compiled, backend, tuple_args, handle_result
    )


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
169     msg = "primitive arguments must be colocated on the same device, got {}"
170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
172   else:
173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)

# Fixed code
def xla_primitive_callable(prim, *arg_specs, **params):
    avals, arg_devices = unzip2(arg_specs)
    device = _device_from_arg_devices(arg_devices)
    backend = xb.get_device_backend(device)
    aval_out = prim.abstract_eval(*avals, **params)
    if not prim.multiple_results:
        handle_result = aval_to_result_handler(device, aval_out)
    else:
        handlers = tuple(map(partial(aval_to_result_handler, device), aval_out))
        handle_result = lambda xs: tuple(
            h(x) for h, x in zip(handlers, xs.destructure())
        )
    tuple_args = len(avals) > 100
    built_c = primitive_computation(prim, backend, tuple_args, *avals, **params)
    options = xb.get_compile_options(device_assignment=device and (device.id,))
    compiled = built_c.Compile(compile_options=options, backend=backend)
    return partial(
        _execute_compiled_primitive, prim, compiled, backend, tuple_args, handle_result
    )
# Buggy code
def _xla_call_impl(fun, *args, **params):
    device = params["device"]
    backend = params["backend"]
    compiled_fun = _xla_callable(fun, device, backend, *map(abstractify, args))
    try:
        return compiled_fun(*args)
    except FloatingPointError:
        print(
            "Invalid value encountered in the output of a jit function. "
            "Calling the de-optimized version."
        )
        return fun.call_wrapped(*args)  # probably won't return


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
169     msg = "primitive arguments must be colocated on the same device, got {}"
170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
172   else:
173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)

# Fixed code
def _xla_call_impl(fun, *args, **params):
    device = params["device"]
    backend = params["backend"]
    compiled_fun = _xla_callable(fun, device, backend, *map(arg_spec, args))
    try:
        return compiled_fun(*args)
    except FloatingPointError:
        print(
            "Invalid value encountered in the output of a jit function. "
            "Calling the de-optimized version."
        )
        return fun.call_wrapped(*args)  # probably won't return
# Buggy code
def _pval_to_result_handler(pval):
    pv, const = pval
    if pv is None:
        return lambda _: const
    else:
        return aval_to_result_handler(pv)


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
169     msg = "primitive arguments must be colocated on the same device, got {}"
170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
172   else:
173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)

# Fixed code
def _pval_to_result_handler(device, pval):
    pv, const = pval
    if pv is None:
        return lambda _: const
    else:
        return aval_to_result_handler(device, pv)
# Buggy code
def __init__(self, aval, device_buffer):
    self.aval = aval
    self.device_buffer = device_buffer
    # TODO(mattjj): make Device hashable
    device = device_buffer.device()
    self._device = device and (type(device), device.id)

    self._npy_value = None
    if not core.skip_checks:
        assert type(aval) is ShapedArray
        npy_value = self._value
        assert npy_value.dtype == aval.dtype and npy_value.shape == aval.shape


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
169     msg = "primitive arguments must be colocated on the same device, got {}"
170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
172   else:
173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)

# Fixed code
def __init__(self, aval, device, device_buffer):
    self.aval = aval
    self.device_buffer = device_buffer
    self._device = device and (type(device), device.id)

    self._npy_value = None
    if not core.skip_checks:
        assert type(aval) is ShapedArray
        npy_value = self._value
        assert npy_value.dtype == aval.dtype and npy_value.shape == aval.shape
# Buggy code
def _device_put_impl(x, device=None):
    try:
        a = abstractify(x)
    except TypeError:
        raise TypeError(
            "Argument '{}' of type {} is not a valid JAX type".format(x, type(x))
        )
    handler = aval_to_result_handler(a)
    return handler(device_put(x, device))


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
169     msg = "primitive arguments must be colocated on the same device, got {}"
170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
172   else:
173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)

# Fixed code
def _device_put_impl(x, device=None):
    try:
        a = abstractify(x)
    except TypeError:
        raise TypeError(
            "Argument '{}' of type {} is not a valid JAX type".format(x, type(x))
        )
    handler = aval_to_result_handler(device, a)
    return handler(device_put(x, device))
# Buggy code
def _instantiate_device_constant(const, device=None, backend=None, cutoff=1e6):
    # dispatch an XLA Computation to build the constant on the device if it's
    # large, or alternatively build it on the host and transfer it if it's small
    assert isinstance(const, DeviceConstant)
    if const.size > cutoff:
        c = xb.make_computation_builder("constant_instantiating_computation")
        xla_const = const.constant_handler(c, const)
        device_assignment = (device.id,) if device else None
        opts = xb.get_compile_options(device_assignment=device_assignment)
        compiled = c.Build(xla_const).Compile((), opts, backend=xb.get_backend(backend))
        return compiled.Execute(())
    else:
        return xc.Buffer.from_pyval(
            onp.asarray(const), device, backend=xb.get_backend(backend)
        )


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-20-ca9775de72ea> in <module>()
----> 1 x / x.shape[0]

4 frames
google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
169     msg = "primitive arguments must be colocated on the same device, got {}"
170     names = ("{}({})".format(d[0].__name__, d[1]) for d in devices if d is not None)
--> 171     raise ValueError(msg.format(", ".join(names)))
172   else:
173     all_devices = it.chain(xb.devices(), xb.devices('cpu'))

ValueError: primitive arguments must be colocated on the same device, got CpuDevice(0), GpuDevice(0)

# Fixed code
def _instantiate_device_constant(const, device=None, backend=None, cutoff=1e6):
    # dispatch an XLA Computation to build the constant on the device if it's
    # large, or alternatively build it on the host and transfer it if it's small
    assert isinstance(const, DeviceConstant)
    backend = xb.get_backend(device.platform) if device else xb.get_backend(backend)
    if const.size > cutoff:
        c = xb.make_computation_builder("constant_instantiating_computation")
        xla_const = const.constant_handler(c, const)
        device_assignment = (device.id,) if device else None
        opts = xb.get_compile_options(device_assignment=device_assignment)
        compiled = c.Build(xla_const).Compile((), opts, backend=backend)
        return compiled.Execute(())
    else:
        return xc.Buffer.from_pyval(onp.asarray(const), device, backend=backend)
# Buggy code
def update(self, name, val):
    self.check_exists(name)
    if name not in self.values:
        raise Exception("Unrecognized config option: {}".format(name))
    self.values[name] = val


# Traceback (TypeError)
TypeError                                 Traceback (most recent call last)
<ipython-input-171-c4c8dbbffe17> in <module>
1 A = np.array([[1.,2.],[2.,3.]])
----> 2 V,D = np.linalg.eig(A)

~/RESEARCH/jax/jax/numpy/linalg.py in eig(a)
95 def eig(a):
96   a = _promote_arg_dtypes(np.asarray(a))
---> 97   w, vl, vr = lax_linalg.eig(a)
98   return w, vr
99

~/RESEARCH/jax/jax/lax_linalg.py in eig(x)
47
48 def eig(x):
---> 49   w, vl, vr = eig_p.bind(x)
50   return w, vl, vr
51

~/RESEARCH/jax/jax/core.py in bind(self, *args, **kwargs)
128     top_trace = find_top_trace(args)
129     if top_trace is None:
--> 130       return self.impl(*args, **kwargs)
131
132     tracers = map(top_trace.full_raise, args)

~/RESEARCH/jax/jax/lax_linalg.py in eig_impl(operand)
153
154 def eig_impl(operand):
--> 155   return xla.apply_primitive(eig_p, operand)
156
157 def eig_translation_rule(c, operand):

~/RESEARCH/jax/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
122   abstract_args = map(abstractify, args)
123   compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
--> 124   return compiled_fun(*args)
125
126 @cache()

~/RESEARCH/jax/jax/interpreters/xla.py in _execute_compiled_primitive(prim, compiled, backend, result_handler, *args)
172   input_bufs = [device_put(x, device_num, backend=backend) for x in args]
173   out_buf = compiled.Execute(input_bufs)
--> 174   if FLAGS.jax_debug_nans: check_nans(prim, out_buf)
175   return result_handler(out_buf)
176

~/RESEARCH/jax/jax/interpreters/xla.py in check_nans(prim, bufs)
177 def check_nans(prim, bufs):
178   if prim.multiple_results:
--> 179     for buf in bufs:
180       _check_nans(prim.name, buf.shape(), buf)
181   else:

TypeError: 'jaxlib.xla_extension.PyLocalBuffer' object is not iterable

# Fixed code
def update(self, name, val):
    if self.use_absl:
        setattr(self.absl_flags.FLAGS, name, val)
    else:
        self.check_exists(name)
        if name not in self.values:
            raise Exception("Unrecognized config option: {}".format(name))
        self.values[name] = val
# Buggy code
def _execute_compiled_primitive(prim, compiled, backend, result_handler, *args):
    (device_num,) = compiled.DeviceOrdinals()
    input_bufs = [device_put(x, device_num, backend=backend) for x in args]
    out_buf = compiled.Execute(input_bufs)
    if FLAGS.jax_debug_nans:
        check_nans(prim, out_buf)
    return result_handler(out_buf)


# Traceback (TypeError)
TypeError                                 Traceback (most recent call last)
<ipython-input-171-c4c8dbbffe17> in <module>
1 A = np.array([[1.,2.],[2.,3.]])
----> 2 V,D = np.linalg.eig(A)

~/RESEARCH/jax/jax/numpy/linalg.py in eig(a)
95 def eig(a):
96   a = _promote_arg_dtypes(np.asarray(a))
---> 97   w, vl, vr = lax_linalg.eig(a)
98   return w, vr
99

~/RESEARCH/jax/jax/lax_linalg.py in eig(x)
47
48 def eig(x):
---> 49   w, vl, vr = eig_p.bind(x)
50   return w, vl, vr
51

~/RESEARCH/jax/jax/core.py in bind(self, *args, **kwargs)
128     top_trace = find_top_trace(args)
129     if top_trace is None:
--> 130       return self.impl(*args, **kwargs)
131
132     tracers = map(top_trace.full_raise, args)

~/RESEARCH/jax/jax/lax_linalg.py in eig_impl(operand)
153
154 def eig_impl(operand):
--> 155   return xla.apply_primitive(eig_p, operand)
156
157 def eig_translation_rule(c, operand):

~/RESEARCH/jax/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
122   abstract_args = map(abstractify, args)
123   compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
--> 124   return compiled_fun(*args)
125
126 @cache()

~/RESEARCH/jax/jax/interpreters/xla.py in _execute_compiled_primitive(prim, compiled, backend, result_handler, *args)
172   input_bufs = [device_put(x, device_num, backend=backend) for x in args]
173   out_buf = compiled.Execute(input_bufs)
--> 174   if FLAGS.jax_debug_nans: check_nans(prim, out_buf)
175   return result_handler(out_buf)
176

~/RESEARCH/jax/jax/interpreters/xla.py in check_nans(prim, bufs)
177 def check_nans(prim, bufs):
178   if prim.multiple_results:
--> 179     for buf in bufs:
180       _check_nans(prim.name, buf.shape(), buf)
181   else:

TypeError: 'jaxlib.xla_extension.PyLocalBuffer' object is not iterable

# Fixed code
def _execute_compiled_primitive(prim, compiled, backend, result_handler, *args):
    (device_num,) = compiled.DeviceOrdinals()
    input_bufs = [device_put(x, device_num, backend=backend) for x in args]
    out_buf = compiled.Execute(input_bufs)
    if FLAGS.jax_debug_nans:
        check_nans(prim, out_buf.destructure() if prim.multiple_results else out_buf)
    return result_handler(out_buf)
# Buggy code
def cholesky_cpu_translation_rule(c, operand):
    shape = c.GetShape(operand)
    dtype = shape.element_type().type
    if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
        return c.GetTupleElement(lapack.jax_potrf(c, operand, lower=True), 0)
    else:
        # Fall back to the HLO implementation for batched Cholesky decomposition or
        # unsupported types.
        # TODO(phawkins): support LAPACK primitives in batched mode.
        return c.Cholesky(operand)


# Traceback (numpy.linalg.LinAlgError)
/home/luke/.local/lib/python3.5/site-packages/jax/scipy/linalg.py:36: UserWarning: scipy.linalg support is experimental and may cause silent failures or wrong outputs
warnings.warn(_EXPERIMENTAL_WARNING)
[[0. 1. 2.]
[0. 4. 5.]
[0. 0. 8.]]
From Scipy
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
print(linalg.cholesky(x))
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 91, in cholesky
check_finite=check_finite)
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 40, in _cholesky
"definite" % info)
numpy.linalg.LinAlgError: 1-th leading minor of the array is not positive definite

# Fixed code
def cholesky_cpu_translation_rule(c, operand):
    shape = c.GetShape(operand)
    dtype = shape.element_type().type
    if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
        potrf_output = lapack.jax_potrf(c, operand, lower=True)
        result = c.GetTupleElement(potrf_output, 0)
        info = c.GetTupleElement(potrf_output, 1)
        return c.Select(
            c.Eq(info, c.ConstantS32Scalar(0)), result, _nan_like(c, result)
        )
    else:
        # Fall back to the HLO implementation for batched Cholesky decomposition or
        # unsupported types.
        # TODO(phawkins): support LAPACK primitives in batched mode.
        return c.Cholesky(operand)