Repository: google/jax

Issue:
Regression bug in associative_scan
Hi,

There has been a regression in JAX code, which I believe to be related to this [commit](https://github.com/google/jax/commit/d3db7bd4be96e1fc16daa4571ade7fbfce13197f#diff-cc3f3a1265e2fee9115dc5a9f619d402e280fb48a7ed7d91261b42813d200f1b). 

I have created PR https://github.com/google/jax/pull/5165 with a reproducing failing test.

I will try and look further into what the problem could be, but if the person who made said commit could also have a look it would be great (@hawkinsp I believe this is you).

Adrien


Discussion:
Thanks for raising this!

What's the nature of the regression? An error, or a performance regression, or something else?
The error appears to be:
```
FAILED tests/lax_control_flow_test.py::LaxControlFlowTest::testAssociativeScanFailing_2 - TypeError: Slice size at index 0 in gather op is out of range, must be within [0, 1), got 1.
```

with traceback:
```
Traceback (most recent call last):
  File "/Users/phawkins/.pyenv/versions/py3.9.0/lib/python3.9/site-packages/absl/testing/parameterized.py", line 282, in bound_param_test
    return test_method(self, **testcase_params)
  File "/Users/phawkins/t/issue5164/tests/lax_control_flow_test.py", line 2488, in testAssociativeScanFailing
    _ = lax.associative_scan(fn, elems=(ms, vs))
  File "/Users/phawkins/p/jax/jax/_src/lax/control_flow.py", line 2492, in associative_scan
    scans = _scan(elems_flat)
  File "/Users/phawkins/p/jax/jax/_src/lax/control_flow.py", line 2476, in _scan
    even_elems = combine(
  File "/Users/phawkins/p/jax/jax/_src/lax/control_flow.py", line 2431, in combine
    c = fn(a, b)
  File "/Users/phawkins/p/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/api.py", line 1197, in batched_fun
    out_flat = batching.batch(flat_fun, args_flat, in_axes_flat,
  File "/Users/phawkins/p/jax/jax/interpreters/batching.py", line 35, in batch
    return batched_fun.call_wrapped(*in_vals)
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 160, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/phawkins/t/issue5164/tests/lax_control_flow_test.py", line 2486, in fn
    return m1 + m2, jsp.linalg.solve(m1, v2) + jsp.linalg.solve(m2, v1)
  File "/Users/phawkins/p/jax/jax/_src/scipy/linalg.py", line 193, in solve
    return _solve(a, b, sym_pos, lower)
  File "/Users/phawkins/p/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/api.py", line 382, in f_jitted
    return cpp_jitted_f(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/api.py", line 278, in cache_miss
    out_flat = xla.xla_call(
  File "/Users/phawkins/p/jax/jax/core.py", line 1229, in bind
    return call_bind(self, fun, *args, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1220, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1232, in process
    return trace.process_call(self, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/interpreters/batching.py", line 163, in process_call
    vals_out = call_primitive.bind(f, *vals, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1229, in bind
    return call_bind(self, fun, *args, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1220, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1232, in process
    return trace.process_call(self, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 598, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/Users/phawkins/p/jax/jax/interpreters/xla.py", line 569, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 251, in memoized_fun
    ans = call(fun, *args)
  File "/Users/phawkins/p/jax/jax/interpreters/xla.py", line 645, in _xla_callable
    jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(fun, abstract_args)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 1230, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 1211, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 160, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/phawkins/p/jax/jax/_src/scipy/linalg.py", line 168, in _solve
    return np_linalg.solve(a, b)
  File "/Users/phawkins/p/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/api.py", line 382, in f_jitted
    return cpp_jitted_f(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/api.py", line 278, in cache_miss
    out_flat = xla.xla_call(
  File "/Users/phawkins/p/jax/jax/core.py", line 1229, in bind
    return call_bind(self, fun, *args, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1220, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1232, in process
    return trace.process_call(self, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/interpreters/batching.py", line 163, in process_call
    vals_out = call_primitive.bind(f, *vals, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1229, in bind
    return call_bind(self, fun, *args, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1220, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1232, in process
    return trace.process_call(self, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 1085, in process_call
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(f, self.main, in_avals)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 1211, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 160, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/phawkins/p/jax/jax/_src/numpy/linalg.py", line 450, in solve
    return lax_linalg._solve(a, b)
  File "/Users/phawkins/p/jax/jax/_src/lax/linalg.py", line 263, in _solve
    return custom_solve(b)
  File "/Users/phawkins/p/jax/jax/_src/lax/control_flow.py", line 2211, in custom_linear_solve
    out_flat = linear_solve_p.bind(
  File "/Users/phawkins/p/jax/jax/core.py", line 271, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "/Users/phawkins/p/jax/jax/interpreters/batching.py", line 149, in process_primitive
    val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
  File "/Users/phawkins/p/jax/jax/_src/lax/control_flow.py", line 2296, in _linear_solve_batching_rule
    solve_jaxpr_batched, solve_x_bat = batching.batch_jaxpr(
  File "/Users/phawkins/p/jax/jax/interpreters/batching.py", line 411, in batch_jaxpr
    jaxpr_out, _, consts = pe.trace_to_jaxpr_dynamic(f, avals_in)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 1201, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 1211, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 160, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/phawkins/p/jax/jax/core.py", line 141, in jaxpr_as_fun
    return eval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.consts, *args)
  File "/Users/phawkins/p/jax/jax/core.py", line 352, in eval_jaxpr
    ans = eqn.primitive.bind(*(subfuns + in_vals), **bind_params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1229, in bind
    return call_bind(self, fun, *args, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1220, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1232, in process
    return trace.process_call(self, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/interpreters/batching.py", line 163, in process_call
    vals_out = call_primitive.bind(f, *vals, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1229, in bind
    return call_bind(self, fun, *args, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1220, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1232, in process
    return trace.process_call(self, fun, tracers, params)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 1085, in process_call
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(f, self.main, in_avals)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 1211, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 160, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/phawkins/p/jax/jax/core.py", line 352, in eval_jaxpr
    ans = eqn.primitive.bind(*(subfuns + in_vals), **bind_params)
  File "/Users/phawkins/p/jax/jax/core.py", line 271, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "/Users/phawkins/p/jax/jax/interpreters/batching.py", line 149, in process_primitive
    val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
  File "/Users/phawkins/p/jax/jax/_src/lax/lax.py", line 4324, in _gather_batching_rule
    return gather(operand, start_indices, dimension_numbers=dnums,
  File "/Users/phawkins/p/jax/jax/_src/lax/lax.py", line 868, in gather
    return gather_p.bind(
  File "/Users/phawkins/p/jax/jax/core.py", line 271, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 1073, in process_primitive
    out_avals = primitive.abstract_eval(*avals, **params)
  File "/Users/phawkins/p/jax/jax/_src/lax/lax.py", line 1989, in standard_abstract_eval
    shapes, dtypes = shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/_src/lax/lax.py", line 4223, in _gather_shape_rule
    raise TypeError(f"Slice size at index {i} in gather op is out of range, "
TypeError: Slice size at index 0 in gather op is out of range, must be within [0, 1), got 1.
```

I didn't yet have time to try to decode what this means.
Hi,

It's an error. I have some code which implements Kalman filtering using prefix sum operations [(see this)](https://arxiv.org/abs/1905.13002). It used to work like a charm, but I've come back to it recently to see it fail a bit randomly (it's related to batching it seems). I've narrowed it down to the combination of associative_scan and lu permutations, hence the test added in the PR.

Adrien
@hawkinsp thanks, if you disable jit in your env you'll see that the stack stops at the lu solving bit.
I'm not 100% sure if its the LU that's to blame or associative scan but looking at the file history it seemed more plausible for it to be the associative_scan

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
@@ -4352,7 +4352,7 @@ def _gather_batching_rule(batched_args, batch_dims, *, dimension_numbers,
     counts = broadcasted_iota(start_indices.dtype, tuple(count_shape), 0)
     start_indices = concatenate([counts, start_indices], len(count_shape) - 1)
 
-    slice_sizes = (1,) + slice_sizes
+    slice_sizes = (_min(operand.shape[0], 1),) + slice_sizes
     collapsed_slice_dims = (0,) + tuple(np.add(1, dimension_numbers.collapsed_slice_dims))
     offset_dims = tuple(np.add(1, dimension_numbers.offset_dims))
     start_index_map = (0,) + tuple(np.add(1, dimension_numbers.start_index_map))


Test patch:
diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py
--- a/tests/lax_vmap_test.py
+++ b/tests/lax_vmap_test.py
@@ -74,7 +74,12 @@ def _CheckBatching(self, op, bdim_size, bdims, shapes, dtypes, rng,
     args = [rng(shape, dtype) for shape, dtype in zip(batched_shapes, dtypes)]
     args_slice = args_slicer(args, bdims)
     ans = api.vmap(op, bdims)(*args)
-    expected = np.stack([op(*args_slice(i)) for i in range(bdim_size)])
+    if bdim_size == 0:
+      args = [rng(shape, dtype) for shape, dtype in zip(shapes, dtypes)]
+      out = op(*args)
+      expected = np.zeros((0,) + out.shape, out.dtype)
+    else:
+      expected = np.stack([op(*args_slice(i)) for i in range(bdim_size)])
     self.assertAllClose(ans, expected, rtol=rtol, atol=atol)
 
   @parameterized.named_parameters(itertools.chain.from_iterable(
@@ -642,6 +647,8 @@ def testFft(self, fft_ndims, shape, bdims):
       for bdims in all_bdims(shape, idxs.shape)))
   def testGather(self, shape, dtype, idxs, dnums, slice_sizes, bdims):
     fun = partial(lax.gather, dimension_numbers=dnums, slice_sizes=slice_sizes)
+    self._CheckBatching(fun, 0, bdims, [shape, idxs.shape], [dtype, idxs.dtype],
+                        jtu.rand_default(self.rng()))
     self._CheckBatching(fun, 5, bdims, [shape, idxs.shape], [dtype, idxs.dtype],
                         jtu.rand_default(self.rng()))