Repository: google/jax

Issue:
jax.numpy.repeat fails with JIT disabled for dimensions of size 1
`jax.numpy.repeat` throws a `RuntimeError` when JIT is disabled and the array dimension being repeated is of size 1.
The same error occurs on both CPU and GPU.

Simple repro:
```py
import jax
import jax.numpy as jnp
with jax.disable_jit():
  jnp.repeat(jnp.zeros((1, 2)), repeats=2, axis=0, total_repeat_length=2)
```

Callstack:
```
---------------------------------------------------------------------------
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[]
```


Discussion:
I'm getting the same error, and am stuck because of it. 
I have the same issue when disabling jit, but everything works fine when it is enabled.
I've been looking into this; it seems to be coming from a strange behavior of `jnp.cumsum` with jit disabled, which can be reproduced this way:
```python
from jax import lax
import jax.numpy as jnp

print(repr(lax.cumsum(jnp.ones(1))))
```
```
DeviceArray(False, dtype=float32)
```
I'm not certain how `False` has a dtype of `float32`, but I suspect whatever bug is causing this is the root of the issue here.

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -1135,7 +1135,6 @@ def reduce(operands: Array, init_values: Array, computation: Callable,
     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,
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -329,7 +329,7 @@ def primitive_computation(prim, axis_env, backend, tuple_args, *avals, **params)
   assert isinstance(ans, xe.XlaOp)
   c.clear_op_metadata()
   try:
-    return c.build()
+    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"


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2392,5 +2392,11 @@ def testUnaryWeakTypes(self, op_name, rec_dtypes):
     self.assertTrue(py_op.aval.weak_type)
     self.assertFalse(lax_op.aval.weak_type)
 
+  def testCumsumLengthOne(self):
+    # regression test for issue 4672
+    x = lax.full((1,), 1)
+    out = lax.cumsum(x)
+    self.assertArraysEqual(out, x)
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())