jax-ml/jax · error · ValueError

ragged_all_to_all input_offsets must be integer type.

Error message

ragged_all_to_all input_offsets must be integer type.

What it means

ragged_all_to_all describes per-device payload boundaries via offset arrays; input_offsets must be an integer-dtype array. The abstract eval rejects floating (or other) dtypes with ValueError before lowering.

Source

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

        ir.IntegerType.get_signless(64), mlir.COLLECTIVE_CHANNEL_ID
    )

  return hlo.CustomCallOp(
      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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast: input_offsets = input_offsets.astype(np.int64)
2. Compute offsets with integer arithmetic from the start
3. Validate dtypes before calling

Example fix

// before
lax.ragged_all_to_all(x, out, np.array([0.0, 5.0]), sizes, 'i')
// after
lax.ragged_all_to_all(x, out, np.array([0, 5], dtype=np.int64), sizes, 'i')

import numpy as np
assert np.issubdtype(np.asarray(input_offsets).dtype, np.integer), 'input_offsets must be int'

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

When it happens

Trigger: Passing input_offsets as float32/float64 (e.g. from np.array([...]) default or computed float math) to lax.ragged_all_to_all.

Common situations: Computing offsets with float arithmetic; loading offsets from JSON/np arrays that default to float.

Related errors

• ragged_all_to_all send_sizes must be integer type.
• ragged_all_to_all output_offsets 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/c80b9ed3222a8071. Report an issue: GitHub.