jax-ml/jax · error · TypeError

index arguments to dynamic_update_slice must be integers of

Error message

index arguments to dynamic_update_slice must be integers of the same type, got {}

What it means

The dtype rule of dynamic_update_slice requires every start-index argument to be an integer dtype and all of them to share the same dtype (e.g. all int32). Passing mixed dtypes (int32 and int64) or float indices raises this TypeError. This mirrors XLA's requirement that all index arguments have a uniform integer type.

Source

Thrown at jax/_src/lax/slicing.py:1697

  if core.getr(operand) != core.getr(update):
    raise core.ShardingTypeError(
        "dynamic_update_slice operand and update must be reduced along the"
        " same axes. Got operand sharding"
        f" {operand.str_short(mesh_axis_types=True)} and update sharding"
        f" {update.str_short(mesh_axis_types=True)}.")
  return core.getr(operand)

def _dynamic_update_slice_ur_rule(operand, update, *start_indices):
  out_u, kind = _dus_unreduced_rule(operand, update)
  return out_u, _dus_reduced_rule(operand, update), kind

def _dynamic_update_slice_dtype_rule(operand, update, *start_indices):
  lax.check_same_dtypes("dynamic_update_slice", operand, update)
  if any(i.dtype != start_indices[0].dtype or
         not dtypes.issubdtype(i.dtype, np.integer) for i in start_indices):
    msg = ("index arguments to dynamic_update_slice must be integers of the "
           "same type, got {}")
    raise TypeError(msg.format(", ".join(i.dtype.name for i in start_indices)))
  return operand.dtype

def _dynamic_update_slice_jvp(primals, tangents):
  operand, update = primals[:2]
  start_indices = primals[2:]
  g_operand, g_update = tangents[:2]
  val_out = dynamic_update_slice_p.bind(operand, update, *start_indices)
  if type(g_operand) is ad_util.Zero and type(g_update) is ad_util.Zero:
    tangent_out = ad_util.p2tz(val_out)
  else:
    g_operand = ad.instantiate_zeros(g_operand)
    g_update = ad.instantiate_zeros(g_update)
    tangent_out = dynamic_update_slice_p.bind(g_operand, g_update, *start_indices)
  return val_out, tangent_out

def _dynamic_update_slice_transpose_rule(t, operand, update, *start_indices):
  assert all(not ad.is_undefined_primal(x) for x in start_indices)
  if type(t) is ad_util.Zero:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast all indices to one integer type before the call: start = tuple(jnp.asarray(i, dtype=jnp.int32) for i in starts).
If indices come from float math, round and cast explicitly: jnp.floor(pos).astype(jnp.int32).
Standardize on one index dtype project-wide (usually int32) to avoid mixed-width args.

Example fix

# before
starts = (np.int64(2), np.int32(0))
out = lax.dynamic_update_slice(buf, upd, *starts)  # TypeError

# after
starts = tuple(int(i) for i in starts)
out = lax.dynamic_update_slice(buf, upd, *starts)

Defensive patterns

Strategy: validation

Validate before calling

idx_dtype = jnp.int32
starts = tuple(jnp.asarray(i, dtype=idx_dtype) for i in starts)
out = lax.dynamic_update_slice(operand, update, *starts)

Try / catch

try:
    out = lax.dynamic_update_slice(operand, update, *starts)
except TypeError as e:
    if 'same type' in str(e):
        starts = tuple(jnp.asarray(i, dtype=jnp.int32) for i in starts)
        out = lax.dynamic_update_slice(operand, update, *starts)
    else:
        raise

Prevention

Pick one index dtype (int32) and cast all indices to it at the boundary of your module.
Avoid mixing numpy int64 counters with jnp int32 indices; convert loop variables explicitly.

When it happens

Trigger: Calling lax.dynamic_update_slice with start indices of mixed integer widths (one np.int32, one np.int64), or with float values like 0.5 or jnp.float32 positions.

Common situations: On platforms (Windows/older numpy) where the default int is int32 while loop counters or jnp.arange defaults produce int64; combining Python ints with numpy index arrays; passing positions computed in float (e.g. from an interpolation step) without casting.

Related errors

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