jax-ml/jax · error · ValueError

ragged_all_to_all output_offsets must be integer type.

Error message

ragged_all_to_all output_offsets must be integer type.

What it means

output_offsets passed to ragged_all_to_all must be an integer-dtype array; the abstract eval rejects non-integer dtypes before compilation.

Source

Thrown at jax/_src/lax/parallel.py:1649

      result=[output.type],
      inputs=[operand, output, input_offsets, send_sizes, output_offsets,
              recv_sizes],
      call_target_name=ir.StringAttr.get('ragged_all_to_all'),
      backend_config=ir.DictAttr.get(ragged_all_to_all_attrs),
      api_version=ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 4),
  ).results

def _ragged_all_to_all_effectful_abstract_eval(
    operand, output, input_offsets, send_sizes, output_offsets, recv_sizes,
    axis_name, axis_index_groups
):
  del operand, axis_index_groups
  if not dtypes.issubdtype(input_offsets.dtype, np.integer):
    raise ValueError("ragged_all_to_all input_offsets must be integer type.")
  if not dtypes.issubdtype(send_sizes.dtype, np.integer):
    raise ValueError("ragged_all_to_all send_sizes must be integer type.")
  if not dtypes.issubdtype(output_offsets.dtype, np.integer):
    raise ValueError("ragged_all_to_all output_offsets must be integer type.")
  if not dtypes.issubdtype(recv_sizes.dtype, np.integer):
    raise ValueError("ragged_all_to_all recv_sizes must be integer type.")
  if len(input_offsets.shape) != 1 or input_offsets.shape[0] < 1:
    raise ValueError(
        "ragged_all_to_all input_offsets must be rank 1 with positive dimension"
        " size, but got shape {}".format(input_offsets.shape)
    )
  if len(send_sizes.shape) != 1 or send_sizes.shape[0] < 1:
    raise ValueError(
        "ragged_all_to_all send_sizes must be rank 1 with positive dimension"
        " size, but got shape {}".format(send_sizes.shape)
    )
  if len(output_offsets.shape) != 1 or output_offsets.shape[0] < 1:
    raise ValueError(
        "ragged_all_to_all output_offsets must be rank 1 with positive"
        " dimension size, but got shape {}".format(output_offsets.shape)
    )
  if len(recv_sizes.shape) != 1 or recv_sizes.shape[0] < 1:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast output_offsets to an integer dtype
2. Keep a single validated int offsets helper used for all sides

Example fix

// before
lax.ragged_all_to_all(x, out, offs, sizes, 'i', output_offsets=np.array([0., 4.]))
// after
lax.ragged_all_to_all(x, out, offs, sizes, 'i', output_offsets=np.array([0, 4], dtype=np.int64))

assert np.issubdtype(np.asarray(output_offsets).dtype, np.integer), 'output_offsets must be int'

def is_int_array(a): return np.issubdtype(np.asarray(a).dtype, np.integer)

When it happens

Trigger: Passing a float-typed output_offsets array to lax.ragged_all_to_all.

Common situations: Reusing a single float offsets buffer for both input and output offsets; serialization round-trips producing floats.

Related errors

• ragged_all_to_all input_offsets must be integer type.
• ragged_all_to_all send_sizes must be integer type.
• ragged_all_to_all recv_sizes must be integer type.
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass

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