Indexing a Jax Tracer?

[syedzayyan]

Hi All!

I am trying to index a jax tracer. Initially I looked into converting a tracer to a normal jax array and then I stumbled upon this: https://docs.jax.dev/en/latest/faq.html#how-can-i-convert-a-jax-tracer-to-a-numpy-array . The main problem I am running into is I need this to happen:
jnp.take(cov, X[:, 0], axis=0)
but this results in this error:
IndexError: Too many indices: 1-dimensional array indexed with 2 regular indices.

And when I checked with print(type(X)), I got a <class 'jax._src.interpreters.batching.BatchTracer'>.

The function works without any qualms when not in grad, but has this tracer thing in grad mode which is understandable.

Way more context, I am trying to implement a graph edge kernel in GPJax and I am running into indexing issues.

Any help is much appreciated!

[jakevdp]

From the error message, it sounds like X is one-dimensional, so X[:, 0] is an error: this would be true regardless of whether X is a tracer, a jax array, or a NumPy array.

The fix here would be to not attempt to index a 1D array or tracer with two indices.

The function works without any qualms when not in grad, but has this tracer thing in grad mode which is understandable.

This probably indicates that you're doing some sort of batching transformation along with the gradient. We'll need more info in order to help out – could you include a minimal reproducible example of code that shows this behavior?

[syedzayyan]

Hi! Thanks for the reply! This is a minimal reproducible example:

import jax

def f(x):
    print("inside trace, type:", type(x))
    first_col = x[:, 0]
    print("indexed first column (still a tracer):", first_col)
    return first_col

# Input array (batch dimension along axis 0)
x = jax.numpy.array([
    [0, 2],
    [0, 3],
    [0, 5],
    [0, 6]
])

# Apply vmap so a batching tracer appears during tracing
batched_f = jax.vmap(f, in_axes=0)
batched_f(x)

Here's the full error trace:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[25], line 19
     17 # Apply vmap so a batching tracer appears during tracing
     18 batched_f = jax.vmap(f, in_axes=0)
---> 19 batched_f(x)

    [... skipping hidden 7 frame]

Cell In[25], line 5, in f(x)
      3 def f(x):
      4     print("inside trace, type:", type(x))
----> 5     first_col = x[:, 0]
      6     print("indexed first column (still a tracer):", first_col)
      7     return first_col

File ~/Desktop/GPJax/.venv/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:1083, in _forward_operator_to_aval.<locals>.op(self, *args)
   1082 def op(self, *args):
-> 1083   return getattr(self.aval, f"_{name}")(self, *args)

File ~/Desktop/GPJax/.venv/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:657, in _getitem(self, item)
    656 def _getitem(self, item):
--> 657   return indexing.rewriting_take(self, item)

File ~/Desktop/GPJax/.venv/lib/python3.10/site-packages/jax/_src/numpy/indexing.py:653, in rewriting_take(arr, idx, indices_are_sorted, unique_indices, mode, fill_value, out_sharding)
    650 if out_sharding is not None:
    651   return auto_axes(internal_gather, out_sharding=out_sharding
    652                    )(arr, dynamic_idx)
--> 653 return internal_gather(arr, dynamic_idx)

File ~/Desktop/GPJax/.venv/lib/python3.10/site-packages/jax/_src/numpy/indexing.py:662, in _gather(arr, dynamic_idx, treedef, static_idx, indices_are_sorted, unique_indices, mode, fill_value)
    659 def _gather(arr, dynamic_idx, *, treedef, static_idx, indices_are_sorted,
    660             unique_indices, mode, fill_value):
    661   idx = merge_static_and_dynamic_indices(treedef, static_idx, dynamic_idx)
--> 662   indexer = index_to_gather(np.shape(arr), idx)  # shared with _scatter_update
    663   jnp_error._check_precondition_oob_gather(arr.shape, indexer.gather_indices)
    664   y = arr

File ~/Desktop/GPJax/.venv/lib/python3.10/site-packages/jax/_src/numpy/indexing.py:799, in index_to_gather(x_shape, idx, normalize_indices)
    796 advanced_axes_are_contiguous = np.all(np.diff(is_advanced) == 1)
    798 # Remove ellipses and add trailing slice(None)s.
--> 799 idx = _canonicalize_tuple_index(len(x_shape), idx)
    801 # Check for scalar boolean indexing: this requires inserting extra dimensions
    802 # before performing the rest of the logic.
    803 scalar_bool_dims: Sequence[int] = [n for n, i in enumerate(idx) if isinstance(i, bool)]

File ~/Desktop/GPJax/.venv/lib/python3.10/site-packages/jax/_src/numpy/indexing.py:1118, in _canonicalize_tuple_index(arr_ndim, idx)
   1116 if num_dimensions_consumed > arr_ndim:
   1117   index_or_indices = "index" if num_dimensions_consumed == 1 else "indices"
-> 1118   raise IndexError(
   1119       f"Too many indices: {arr_ndim}-dimensional array indexed "
   1120       f"with {num_dimensions_consumed} regular {index_or_indices}.")
   1121 ellipses = (i for i, elt in enumerate(idx) if elt is Ellipsis)
   1122 ellipsis_index = next(ellipses, None)

IndexError: Too many indices: 1-dimensional array indexed with 2 regular indices.

[syedzayyan]

This works though:

import jax

def f(x):
    print("inside trace, type:", type(x))
    first_col = x[:, 0]
    print("indexed first column (still a tracer):", first_col)
    return first_col

# Input array (batch dimension along axis 0)
x = jax.numpy.array([
    [0, 2],
    [0, 3],
    [0, 5],
    [0, 6]
])

# Apply vmap so a batching tracer appears during tracing
f(x)

inside trace, type: <class 'jaxlib._jax.ArrayImpl'>
indexed first column (still a tracer): [0 0 0 0]
Array([0, 0, 0, 0], dtype=int32)

[jakevdp]

Thanks for the repro! This is behaving as expected.

When you wrap a function in vmap, it means that the function effectively operates over a single batch of the input. Your input x is two-dimensional, and so each batch is logically a one-dimensional vector. When you write x[:, 0] within your function, you are attempting a two-dimensional indexing operation on a one-dimensional array, which leads to the error you're seeing.

To fix this, you'll either need to (1) not wrap your function in vmap, or (2) rewrite your function so that it accepts one-dimensional batches.

You can read more about vmap and automatic vectorization at https://docs.jax.dev/en/latest/automatic-vectorization.html.

[syedzayyan]

Thank you for the explanation. This is making tonnes of sense now.