IndexError: vector meaning?

[DanPuzzuoli]

I'm trying to run a jit compiled gradient and I'm getting the following error:

  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/pjit.py", line 256, in cache_miss
    outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
                                                 ^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/pjit.py", line 167, in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **params)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/core.py", line 2657, in bind
    return self.bind_with_trace(top_trace, args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/core.py", line 389, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/core.py", line 869, in process_primitive
    return primitive.impl(*tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/pjit.py", line 1212, in _pjit_call_impl
    return xc._xla.pjit(name, f, call_impl_cache_miss, [], [], donated_argnums,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/pjit.py", line 1196, in call_impl_cache_miss
    out_flat, compiled = _pjit_call_impl_python(
                         ^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/pjit.py", line 1132, in _pjit_call_impl_python
    lowering_parameters=mlir.LoweringParameters()).compile()
                                                   ^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2276, in compile
    executable = UnloadedMeshExecutable.from_hlo(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2624, in from_hlo
    xla_executable, compile_options = _cached_compilation(
                                      ^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 2531, in _cached_compilation
    xla_executable = compiler.compile_or_get_cached(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/compiler.py", line 294, in compile_or_get_cached
    return backend_compile(backend, computation, compile_options,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/devEnv311/lib/python3.11/site-packages/jax/_src/compiler.py", line 256, in backend_compile
    return backend.compile(built_c, compile_options=options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
IndexError: vector

This error does not get raised if I don't try to jit the gradient function, which makes it difficult to track down what's causing the error. I'm still trying to find a minimal example, but wanted to ask here in case anyone has any insight.

[hawkinsp]

That's a complete mystery to me. Can you share a repro?

[DanPuzzuoli]

Yeah good timing, just boiled it down to a small example:

import jax.numpy as jnp
from jax.lax import cond
from jax import jit, grad

t_span = jnp.array([0.0, 2.0])
t_eval = jnp.array([1.0, 1.5, 1.7, 2.0])

def f(a):
    t = jnp.append(jnp.append(t_span[0], t_eval), t_span[1])
    t = cond(t[-1] == t[-2], lambda x: x.at[-2].set(5.), lambda x: x, t)
    
    y = a * t
    
    return y.real.sum()

jit(grad(f))(1.)

[hawkinsp]

That appears to be a translation of a C++ failure. Filed #18106

[DanPuzzuoli]

Notes:

• The appending seems to be necessary; defining the initial value of t outside of f causes it to work.
• If you remove the jit it works.
• If you change the last entry of t_eval to 2.1 (or anything not equal to 2.0 which triggers the conditional) it works.

[hawkinsp]

Fixed by openxla/xla#6344, which will be in the next jaxlib release (or nightly).

[DanPuzzuoli]

thank you!