jax-ml/jax · error · ValueError

One of src or dst must be a BufferedRef.

Error message

One of src or dst must be a BufferedRef.

What it means

sync_copy requires that at least one endpoint be a BufferedRef (that's the whole point of a pipeline-memory copy). If neither src nor dst is a BufferedRef, there is nothing buffered to sync and this ValueError is raised.

Source

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

      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)


@tree_util.register_pytree_node_class
@dataclasses.dataclass(frozen=True, eq=False)
class PipelineStep:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check that one endpoint actually carries pipeline_mode/buffering; if not, use ordinary ref.set / pallas copy instead of sync_copy
2. Verify the BufferedRef wasn't replaced by a plain ref by a trivial-windowing optimization
3. If both are plain refs, do the copy directly (dst[...] = src[...])

Example fix

# before
sync_copy(plain_src_ref, plain_dst_ref)
# after
plain_dst_ref[...] = plain_src_ref[...]

assert isinstance(src, BufferedRef) or isinstance(dst, BufferedRef), \
    'sync_copy requires a BufferedRef endpoint'

def has_bref(src, dst) -> bool:
    return isinstance(src, BufferedRef) or isinstance(dst, BufferedRef)

When it happens

Trigger: Calling sync_copy with two ordinary references/arrays, e.g. copying HBM-to-HBM or passing unwrapped Refs after a refactor removed the buffering spec.

Common situations: A spec that was expected to be pipelined silently became non-pipelined (e.g. trivial windowing made it a plain ref), then a copy helper is invoked on it; generic copy helper applied to non-pipeline tensors.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Only one of src or dst can 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/4b660869e84fbfc4. Report an issue: GitHub.