jax-ml/jax · error · ValueError

Only one of src or dst can be a BufferedRef.

Error message

Only one of src or dst can be a BufferedRef.

What it means

sync_copy copies data between a pipelined BufferedRef and a regular HBM reference. Exactly one side must be the buffered ref; if both src and dst are BufferedRefs, there is no plain memory to copy through, and ValueError 'Only one of src or dst can be a BufferedRef.' is raised.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:1660

      core_id * num_iters,
      core_id * base_num_iters + rem,
  )
  offsets = jax_util.tuple_update(
      (0,) * len(grid),
      partition_dimension,
      grid_offset,
  )
  return new_grid, offsets


def sync_copy(src: REF | BufferedRef, dst: REF | BufferedRef, indices):
  """Perform a synchronous copy from src to dst."""
  bref: BufferedRef
  hbm_ref: REF
  if isinstance(src, BufferedRef):
    bref = src
    if isinstance(dst, BufferedRef):
      raise ValueError("Only one of src or dst can be a BufferedRef.")
    hbm_ref = dst
    copy_in = False
  else:
    if not isinstance(dst, BufferedRef):
      raise ValueError("One of src or dst must be a BufferedRef.")
    bref = dst
    hbm_ref = src
    copy_in = True
  window_ref = bref.current_ref
  if not bref.is_trivial_windowing:
    hbm_slice = bref.get_dma_slice(_ref_to_value_aval(hbm_ref), indices)
    bref_slice = bref._to_window_slice(hbm_slice)
    hbm_ref = hbm_ref.at[hbm_slice]
    window_ref = window_ref.at[bref_slice]
  if copy_in:
    tpu_helpers.sync_copy(hbm_ref, window_ref)
  else:
    tpu_helpers.sync_copy(window_ref, hbm_ref)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Copy through an intermediate non-buffered Ref (or VMEM scratch) between the two buffered refs
2. Perform the transfer inside the kernel body by reading from one buffer and writing the other
3. Restructure so only one spec is pipelined

Example fix

# before
sync_copy(in_bref, out_bref)
# after
scratch = alloc_buffer(...)  # plain ref
sync_copy(in_bref, scratch)
sync_copy(scratch, out_bref)

assert not (isinstance(src, BufferedRef) and isinstance(dst, BufferedRef)), \
    'sync_copy needs exactly one BufferedRef'

def exactly_one_bref(src, dst) -> bool:
    return isinstance(src, BufferedRef) != isinstance(dst, BufferedRef)

When it happens

Trigger: Calling sync_copy (or a helper like commit / prefetch hooks that route through it) with BufferedRef arguments on both sides, e.g. copying between two pipelined buffers.

Common situations: Chaining two pipelined stages and trying to copy buffer-to-buffer directly instead of going through kernel body reads/writes or an intermediate plain Ref.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• One of src or dst must be a BufferedRef.
• Element with {padding=} is not supported.
• Grid must be specified when using lookahead.
• Output buffering does not support lookahead.
• Specified {device=} which requires a copy since the source d

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