jax-ml/jax · error · ValueError
next_fetch is None
Error message
next_fetch is None
What it means
advance_next_fetch advances the prefetch pointer by one grid step. It is a programming invariant that next_fetch was set before the call; if it is None, the internal state is uninitialized or misused and the guard raises immediately with 'next_fetch is None'.
Source
Thrown at jax/_src/pallas/mosaic/pipeline.py:986
def wait_out(self, dst_ref, grid_indices):
"""Waits for output copy to finish."""
assert self.is_output
if not self.is_buffered: return
assert self.sem_sends is not None
wait_slot = self.current_wait_out_slot
dst_slice = self.get_dma_slice(_ref_to_value_aval(dst_ref), grid_indices)
src_slice = self._to_window_slice(dst_slice)
# Single-buffered outputs are synchronously copied.
if self.buffer_count > 1:
tpu_primitives.make_async_copy(
self._window_ref_at(wait_slot, src_slice), # nb: doesn't matter
dst_ref.at[dst_slice], # only dst shape is important
self.sem_sends.at[wait_slot],
).wait()
def advance_next_fetch(self, grid):
if self.next_fetch is None:
raise ValueError("next_fetch is None")
return self.with_next_fetch(_next_index(tuple(self.next_fetch), grid))
def fetch_with_lookahead(buffered_ref, src_ref,
grid,
grid_offsets,
predicate: jax.Array | bool = True,
max_num_fetches: int | None = None,
update_slots: bool = True):
"""Fetch future blocks using unbounded lookahead.
Args:
buffered_ref: the BufferedRef to fetch for.
src_ref: the source Ref.
grid: the grid bounds.
grid_offsets: the grid offsets (used for megacore).
predicate: a boolean predicate for whether to perform the fetch.
max_num_fetches: the maximum number of fetches to perform. If None,View on GitHub (pinned to 1e1c6a8fc0)
Solutions
- Don't call low-level BufferedRef methods manually; drive prefetching through the public pipeline API (emit_pipeline / fetch_with_lookahead)
- Ensure pipeline_mode prefetch configuration is consistent (prefetched_count > 0 with a bound window_ref)
- If reproduced with supported APIs, file a jax issue with a minimal kernel
Defensive patterns
Strategy: validation
Prevention
- Use only the public pipeline API (emit_pipeline/fetch_with_lookahead); never manipulate BufferedRef state directly
- Report reproductions of this invariant error to the JAX project
When it happens
Trigger: Calling advance_next_fetch on a BufferedRef whose next_fetch state was never initialized (e.g. constructing a fetch sequence manually or a state-replacement path dropping next_fetch) — normally reached via initialize_step on prefetched inputs.
Common situations: Almost exclusively a JAX-internal invariant violation or library misuse via lower-level APIs; end users rarely trigger it directly.
Related errors
- Unexpected unhandled transforms: {remaining_ref_transforms}
- `JaxprInputEffect` {eff} does not have a corresponding jaxpr
- This gmm kernel only supports either (m, k) x (g, k, n) -> (
- Group sizes {group_sizes.shape=} must match first dimension
- Explicit sharding is not currently supported in the pallas-t
AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27).
Data as JSON: /api/errors/df3f29a7c813d4d1.
Report an issue: GitHub.