jax-ml/jax · error · ValueError

ragged_all_to_all send_sizes must be integer type.

Error message

ragged_all_to_all send_sizes must be integer type.

What it means

send_sizes for ragged_all_to_all must have an integer dtype; the abstract eval checks dtypes.issubdtype(send_sizes.dtype, np.integer) and raises ValueError otherwise.

Source

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

  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(
        "ragged_all_to_all output_offsets must be rank 1 with positive"
        " dimension size, but got shape {}".format(output_offsets.shape)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast send_sizes to np.int64/int32
2. Use integer literals when constructing the array
3. Add a dtype assertion helper for all four offset/size arrays

Example fix

// before
lax.ragged_all_to_all(x, out, offs, np.array([4.0, 8.0]), 'i')
// after
lax.ragged_all_to_all(x, out, offs, np.array([4, 8], dtype=np.int32), 'i')

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

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

When it happens

Trigger: Passing float send_sizes to lax.ragged_all_to_all.

Common situations: Deriving sizes from shape arithmetic in float; default numpy float arrays from literals like np.array([4, 8]) is fine but np.array([4., 8.]) is not.

Related errors

• ragged_all_to_all input_offsets 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/4d6492d44f513996. Report an issue: GitHub.