jax-ml/jax · error · TypeError

JAX does not support string indexing; got {idx=}

Error message

JAX does not support string indexing; got {idx=}

What it means

IndexType.from_index raises TypeError when a Python str is used as an index into a JAX array. JAX does not implement NumPy's string/field-based indexing (there are no structured dtypes with named fields). Kept as TypeError (rather than IndexError) only for backward compatibility.

Source

Thrown at jax/_src/numpy/indexing.py:91

      return cls.ELLIPSIS
    elif isinstance(idx, slice):
      return cls.SLICE
    elif isinstance(idx, indexing.Slice):
      return cls.DYNAMIC_SLICE
    elif _is_integer_index(idx):
      return cls.INTEGER
    elif _is_boolean_index(idx):
      return cls.BOOLEAN
    elif isinstance(idx, (Array, np.ndarray)):
      if dtypes.issubdtype(idx.dtype, np.integer):
        return cls.ARRAY
      else:
        raise TypeError(
          f"Indexer must have integer or boolean type, got indexer with type {idx.dtype}")
    elif isinstance(idx, str):
      # TODO(jakevdp): this TypeError is for backward compatibility.
      # We should switch to IndexError for consistency.
      raise TypeError(f"JAX does not support string indexing; got {idx=}")
    elif isinstance(idx, Sequence):
      if not idx:  # empty indices default to float, so special-case this.
        return cls.ARRAY
      idx_aval = api.eval_shape(array_constructors.asarray, idx)
      if idx_aval.dtype == bool:
        return cls.BOOLEAN
      elif dtypes.issubdtype(idx_aval.dtype, np.integer):
        return cls.ARRAY
      else:
        raise TypeError(
          f"Indexer must have integer or boolean type, got indexer with type {idx_aval.dtype}")
    elif isinstance(idx, (float, complex, np.generic)):
      raise TypeError(
        f"Indexer must have integer or boolean type, got indexer with type {np.dtype(type(idx))}")
    else:
      raise IndexError("only integers, slices (`:`), ellipsis (`...`), newaxis (`None`)"
                       f" and integer or boolean arrays are valid indices. Got {idx}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Restructure data: use a dict of separate arrays instead of a structured array
  2. If selecting columns, use an integer column index or slice instead
  3. Keep structured arrays in NumPy and convert plain numeric data to jax

Example fix

// before
name = records['name']
// after
name = records_dict['name']  # plain dict of jax arrays
Defensive patterns

Strategy: validation

Validate before calling

if isinstance(idx, str):
    raise ValueError('use a dict of arrays, not string indexing')

Type guard

def is_valid_jax_index(idx) -> bool:
    return not isinstance(idx, str)

Prevention

When it happens

Trigger: x['field_name'] or x['some_string'] on a jax Array, typically code written for NumPy structured/record arrays.

Common situations: Porting NumPy record-array code (arr['names']) to JAX, or accidentally passing a column name/key as an index to a plain jax array.

Related errors


AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/be7fb8795e61b8d5. Report an issue: GitHub.