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

  1. Don't call low-level BufferedRef methods manually; drive prefetching through the public pipeline API (emit_pipeline / fetch_with_lookahead)
  2. Ensure pipeline_mode prefetch configuration is consistent (prefetched_count > 0 with a bound window_ref)
  3. If reproduced with supported APIs, file a jax issue with a minimal kernel
Defensive patterns

Strategy: validation

Prevention

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


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