jax-ml/jax · error · TypeError

Expected strategy to be IndexingStrategy; got {strategy}

Error message

Expected strategy to be IndexingStrategy; got {strategy}

What it means

Internal guard in rewriting_take: the strategy argument must be an instance of IndexingStrategy (STATIC_SLICE, DYNAMIC_SLICE, etc.). Passing anything else is an API-contract violation, usually from internal misuse or monkeypatching.

Source

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

def rewriting_take(
    arr: Array,
    idx: Index | tuple[Index, ...], *,
    indices_are_sorted: bool = False,
    unique_indices: bool = False,
    mode: str | slicing.GatherScatterMode | None = None,
    fill_value: ArrayLike | None = None,
    normalize_indices: bool = True,
    out_sharding: NamedSharding | PartitionSpec | None = None,
    strategy: IndexingStrategy = IndexingStrategy.AUTO,
) -> Array:
  # Computes arr[idx].
  # All supported cases of indexing can be implemented as an XLA gather,
  # followed by an optional reverse and broadcast_in_dim.
  indexer = NDIndexer.from_raw_indices(idx, arr.shape)

  if not isinstance(strategy, IndexingStrategy):
    raise TypeError(f"Expected strategy to be IndexingStrategy; got {strategy}")

  if config.check_static_indices.value and (mode is None or slicing.GatherScatterMode.from_any(mode) == slicing.GatherScatterMode.PROMISE_IN_BOUNDS):
    indexer.validate_static_indices(normalize_indices=normalize_indices)

  if strategy == IndexingStrategy.STATIC_SLICE:
    static_slice_indexer = indexer.to_static_slice(
      arr_is_sharded=indexer.is_sharded(arr),
      normalize_indices=normalize_indices,
      mode=mode)
    return _static_slice(arr, static_slice_indexer)

  if strategy == IndexingStrategy.DYNAMIC_SLICE:
    dynamic_slice_indexer = indexer.to_dynamic_slice(
      arr_is_sharded=indexer.is_sharded(arr),
      normalize_indices=normalize_indices,
      mode=mode)
    return _dynamic_slice(arr, dynamic_slice_indexer)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass an IndexingStrategy member, e.g. jax._src.numpy.indexing.IndexingStrategy.STATIC_SLICE
  2. Prefer public APIs (x[idx], x.at[idx]) instead of rewriting_take
Defensive patterns

Strategy: type-guard

Validate before calling

from jax._src.numpy.indexing import IndexingStrategy
assert isinstance(strategy, IndexingStrategy)

Type guard

def is_indexing_strategy(s) -> bool:
    from jax._src.numpy.indexing import IndexingStrategy
    return isinstance(s, IndexingStrategy)

Prevention

When it happens

Trigger: Calling the internal rewriting_take(arr, idx, strategy=...) with a string, None, or custom object instead of an IndexingStrategy enum member.

Common situations: Code poking at JAX internals; stale third-party code built against an older internal API; test helpers passing raw strings.

Related errors


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