# Buggy code
def reduce(
    operands: Array,
    init_values: Array,
    computation: Callable,
    dimensions: Sequence[int],
) -> Array:
    """Wraps XLA's `Reduce
    <https://www.tensorflow.org/xla/operation_semantics#reduce>`_
    operator.
    """
    flat_operands, operand_tree = tree_util.tree_flatten(operands)
    flat_init_values, init_value_tree = tree_util.tree_flatten(init_values)
    if operand_tree != init_value_tree:
        raise ValueError(
            "Operands must have the same tree structure as init_values:"
            f" {operand_tree} vs. {init_value_tree}"
        )
    if len(flat_operands) != len(flat_init_values):
        raise ValueError(
            "Must have same total number of operands as init_values: "
            f" {len(flat_operands)} vs. {len(flat_init_values)}"
        )
    monoid_reducer = _get_monoid_reducer(computation, flat_init_values)
    if monoid_reducer:
        # monoid reducers bypass the weak_type_rule, so we set it explicitly.
        weak_type = dtypes.is_weakly_typed(*flat_operands) and dtypes.is_weakly_typed(
            *flat_init_values
        )
        return convert_element_type(
            monoid_reducer(*flat_operands, dimensions), weak_type=weak_type
        )
    else:
        flat_init_avals = safe_map(_abstractify, flat_init_values)
        # breakpoint()
        jaxpr, consts, out_tree = _variadic_reduction_jaxpr(
            computation, tuple(flat_init_avals), init_value_tree
        )
        out = reduce_p.bind(
            *(flat_operands + flat_init_values),
            computation=computation,
            jaxpr=jaxpr,
            consts=consts,
            dimensions=tuple(dimensions),
        )
        return tree_util.tree_unflatten(out_tree, out)


# 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 reduce(
    operands: Array,
    init_values: Array,
    computation: Callable,
    dimensions: Sequence[int],
) -> Array:
    """Wraps XLA's `Reduce
    <https://www.tensorflow.org/xla/operation_semantics#reduce>`_
    operator.
    """
    flat_operands, operand_tree = tree_util.tree_flatten(operands)
    flat_init_values, init_value_tree = tree_util.tree_flatten(init_values)
    if operand_tree != init_value_tree:
        raise ValueError(
            "Operands must have the same tree structure as init_values:"
            f" {operand_tree} vs. {init_value_tree}"
        )
    if len(flat_operands) != len(flat_init_values):
        raise ValueError(
            "Must have same total number of operands as init_values: "
            f" {len(flat_operands)} vs. {len(flat_init_values)}"
        )
    monoid_reducer = _get_monoid_reducer(computation, flat_init_values)
    if monoid_reducer:
        # monoid reducers bypass the weak_type_rule, so we set it explicitly.
        weak_type = dtypes.is_weakly_typed(*flat_operands) and dtypes.is_weakly_typed(
            *flat_init_values
        )
        return convert_element_type(
            monoid_reducer(*flat_operands, dimensions), weak_type=weak_type
        )
    else:
        flat_init_avals = safe_map(_abstractify, flat_init_values)
        jaxpr, consts, out_tree = _variadic_reduction_jaxpr(
            computation, tuple(flat_init_avals), init_value_tree
        )
        out = reduce_p.bind(
            *(flat_operands + flat_init_values),
            computation=computation,
            jaxpr=jaxpr,
            consts=consts,
            dimensions=tuple(dimensions),
        )
        return tree_util.tree_unflatten(out_tree, out)