# Buggy code
def primitive_computation(prim, axis_env, backend, tuple_args, *avals, **params):
    c = xb.make_computation_builder(f"primitive_computation_{prim.name}")
    c.set_op_metadata(
        xc.OpMetadata(op_type=prim.name, op_name=str(pp_eqn_compact(prim.name, params)))
    )
    platform = xb.get_backend(backend).platform
    xla_args, _ = _xla_callable_args(c, avals, tuple_args)
    # return val always set as a side-effect on c
    if prim in backend_specific_translations[platform]:
        rule = backend_specific_translations[platform][prim]
        ans = rule(c, *xla_args, **params)
    elif prim in translations:
        rule = translations[prim]
        ans = rule(c, *xla_args, **params)
    elif prim in initial_style_translations:
        rule = initial_style_translations[prim]
        ans = rule(
            c,
            axis_env,
            extend_name_stack(prim.name),
            avals,
            backend,
            *xla_args,
            **params,
        )
    else:
        raise NotImplementedError(f"XLA translation rule for {prim} not found")
    assert isinstance(ans, xe.XlaOp)
    c.clear_op_metadata()
    try:
        return c.build()
    except RuntimeError as e:
        msg = (
            " ".join(map(str, e.args)) + "\n"
            "This is a bug in JAX's shape-checking rules; please report it!\n"
            "https://github.com/google/jax/issues\n"
        )
        raise RuntimeError(msg) from e


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-88-7216fb989347> in <module>()
4
5 with jax.disable_jit():
----> 6   jnp.repeat(jnp.zeros((1, 2)), repeats=2, axis=0, total_repeat_length=2)
7
8 def run_tests():

11 frames
google3/third_party/py/jax/_src/numpy/lax_numpy.py in repeat(a, repeats, axis, total_repeat_length)
2921       x=zeros([total_repeat_length], dtype=int32),
2922       idx=scatter_indices,
-> 2923       y=1)
2924   # Cumsum again to get scatter indices for repeat, e.g. [0,1,1,3,3,3,3,3]
2925   gather_indices = cumsum(block_split_indicators) - 1

google3/third_party/py/jax/_src/ops/scatter.py in index_add(x, idx, y, indices_are_sorted, unique_indices)
141   """
142   return _scatter_update(
--> 143       x, idx, y, lax.scatter_add, indices_are_sorted, unique_indices)
144
145

google3/third_party/py/jax/_src/ops/scatter.py in _scatter_update(x, idx, y, scatter_op, indices_are_sorted, unique_indices)
50   treedef, static_idx, dynamic_idx = jnp._split_index_for_jit(idx)
51   return _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx,
---> 52                        indices_are_sorted, unique_indices)
53
54

google3/third_party/py/jax/_src/ops/scatter.py in _scatter_impl(x, y, scatter_op, treedef, static_idx, dynamic_idx, indices_are_sorted, unique_indices)
62
63   idx = jnp._merge_static_and_dynamic_indices(treedef, static_idx, dynamic_idx)
---> 64   indexer = jnp._index_to_gather(jnp.shape(x), idx)
65
66   # Broadcast `y` to the slice output shape.

google3/third_party/py/jax/_src/numpy/lax_numpy.py in _index_to_gather(x_shape, idx)
4057     advanced_pairs = ((_normalize_index(e, x_shape[j]), i, j)
4058                       for e, i, j in advanced_pairs)
-> 4059     advanced_indexes, idx_advanced_axes, x_advanced_axes = zip(*advanced_pairs)
4060     advanced_axes_are_contiguous = np.all(np.diff(idx_advanced_axes) == 1)
4061

google3/third_party/py/jax/_src/numpy/lax_numpy.py in <genexpr>(.0)
4056       if isscalar(e) or isinstance(e, (Sequence, ndarray)))
4057     advanced_pairs = ((_normalize_index(e, x_shape[j]), i, j)
-> 4058                       for e, i, j in advanced_pairs)
4059     advanced_indexes, idx_advanced_axes, x_advanced_axes = zip(*advanced_pairs)
4060     advanced_axes_are_contiguous = np.all(np.diff(idx_advanced_axes) == 1)

google3/third_party/py/jax/_src/numpy/lax_numpy.py in _normalize_index(index, axis_size)
3785
3786   return lax.select(
-> 3787     lax.lt(index, _constant_like(index, 0)),
3788     lax.add(index, _constant_like(index, axis_size)),
3789     index)

google3/third_party/py/jax/_src/lax/lax.py in lt(x, y)
379 def lt(x: Array, y: Array) -> Array:
380   r"""Elementwise less-than: :math:`x < y`."""
--> 381   return lt_p.bind(x, y)
382
383 def convert_element_type(operand: Array, new_dtype: DType) -> Array:

google3/third_party/py/jax/core.py in bind(self, *args, **params)
261     top_trace = find_top_trace(args)
262     tracers = map(top_trace.full_raise, args)
--> 263     out = top_trace.process_primitive(self, tracers, params)
264     return map(full_lower, out) if self.multiple_results else full_lower(out)
265

google3/third_party/py/jax/core.py in process_primitive(self, primitive, tracers, params)
571
572   def process_primitive(self, primitive, tracers, params):
--> 573     return primitive.impl(*tracers, **params)
574
575   def process_call(self, primitive, f, tracers, params):

google3/third_party/py/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
232   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
233   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
--> 234   return compiled_fun(*args)
235
236

google3/third_party/py/jax/interpreters/xla.py in _execute_compiled_primitive(prim, compiled, result_handler, *args)
347   device, = compiled.local_devices()
348   input_bufs = list(it.chain.from_iterable(device_put(x, device) for x in args if x is not token))
--> 349   out_bufs = compiled.execute(input_bufs)
350   if FLAGS.jax_debug_nans: check_nans(prim, out_bufs)
351   return result_handler(*out_bufs)

RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 0: want s32[1]{0}, got pred[]

# Fixed code
def primitive_computation(prim, axis_env, backend, tuple_args, *avals, **params):
    c = xb.make_computation_builder(f"primitive_computation_{prim.name}")
    c.set_op_metadata(
        xc.OpMetadata(op_type=prim.name, op_name=str(pp_eqn_compact(prim.name, params)))
    )
    platform = xb.get_backend(backend).platform
    xla_args, _ = _xla_callable_args(c, avals, tuple_args)
    # return val always set as a side-effect on c
    if prim in backend_specific_translations[platform]:
        rule = backend_specific_translations[platform][prim]
        ans = rule(c, *xla_args, **params)
    elif prim in translations:
        rule = translations[prim]
        ans = rule(c, *xla_args, **params)
    elif prim in initial_style_translations:
        rule = initial_style_translations[prim]
        ans = rule(
            c,
            axis_env,
            extend_name_stack(prim.name),
            avals,
            backend,
            *xla_args,
            **params,
        )
    else:
        raise NotImplementedError(f"XLA translation rule for {prim} not found")
    assert isinstance(ans, xe.XlaOp)
    c.clear_op_metadata()
    try:
        return c.build(ans)
    except RuntimeError as e:
        msg = (
            " ".join(map(str, e.args)) + "\n"
            "This is a bug in JAX's shape-checking rules; please report it!\n"
            "https://github.com/google/jax/issues\n"
        )
        raise RuntimeError(msg) from e