jax-ml/jax · error · TypeError

start_indices arguments to dynamic_update_slice must be scal

Error message

start_indices arguments to dynamic_update_slice must be scalars, got indices {start_indices}

What it means

Thrown when any of the start index arguments to lax.dynamic_update_slice is not a scalar (has ndim != 0). Each start index must be a 0-D array/scalar so the op knows the single offset per dimension. Passing vectors or higher-rank arrays of indices triggers this error.

Source

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

mlir.register_lowering(dynamic_slice_p, _dynamic_slice_lower)


def _dynamic_update_slice_shape_rule(operand, update, *start_indices):
  if operand.ndim != update.ndim:
    msg = ("dynamic_update_slice update must have the same rank as operand, "
           "got update shape {} for operand shape {}.")
    raise TypeError(msg.format(update.shape, operand.shape))
  if operand.ndim != len(start_indices):
    msg = ("dynamic_update_slice start_indices must have length equal to the "
           "rank of operand, got indices {} for operand shape {}.")
    raise TypeError(msg.format(start_indices, operand.shape))
  if not all(map(operator.ge, operand.shape, update.shape)):
    msg = ("dynamic_update_slice update shape must be smaller than operand "
           "shape, got update shape {} for operand shape {}.")
    raise TypeError(msg.format(update.shape, operand.shape))
  if any(idx.ndim != 0 for idx in start_indices):
    raise TypeError("start_indices arguments to dynamic_update_slice must be "
                    f"scalars, got indices {start_indices}")
  return operand.shape

def _dynamic_update_slice_sharding_rule(operand, update, *start_indices):
  if operand.sharding != update.sharding:
    raise core.ShardingTypeError(
        "dynamic_update_slice operand sharding must be equal to update"
        " sharding, got operand sharding"
        f" {operand.str_short(mesh_axis_types=True)} and update sharding"
        f" {update.str_short(mesh_axis_types=True)}.")
  return operand.sharding

def _dus_unreduced_rule(operand, update):
  if core.getu(operand) != core.getu(update):
    raise core.ShardingTypeError(
        "dynamic_update_slice operand and update must be unreduced along the"
        " same axes. Got operand sharding"
        f" {operand.str_short(mesh_axis_types=True)} and update sharding"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Convert each index to a scalar, e.g. use int(i) or i.reshape(()) / jnp.asarray(i) with shape ().
If you need many offsets at once, switch to lax.gather or operand.at[indices].get/set with an index array.
Verify with assert idx.ndim == 0 before calling.

Example fix

// before
starts = jnp.arange(3)  # shape (3,)
out = lax.dynamic_update_slice(buf, upd, starts)  # TypeError

// after
out = buf
for k in range(3):
    out = lax.dynamic_update_slice(out, upd, (k, 0))

Defensive patterns

Strategy: type-guard

Validate before calling

starts = tuple(i.reshape(()) if hasattr(i, 'reshape') else jnp.asarray(i) for i in starts)
assert all(jnp.asarray(i).ndim == 0 for i in starts)

Type guard

def are_scalar_indices(starts) -> bool:
    return all(jnp.asarray(s).ndim == 0 for s in starts)

Prevention

Normalize every start index with jnp.asarray(i).reshape(()) before the call.
Use lax.gather for vectorized multi-offset reads instead of passing index arrays to dynamic_update_slice.

When it happens

Trigger: Calling lax.dynamic_update_slice(operand, update, idx_array) where idx_array has shape (n,) instead of shape (), e.g. passing a Python list, a tuple of arrays per-axis with vectors, or broadcasting-style index arrays.

Common situations: Confusing dynamic_update_slice (single offset per dim) with lax.gather (vectorized index gathering); passing the output of jnp.arange or an index tensor computed from a loop; migrating numpy code where a length-1 array was implicitly treated as a scalar.

Related errors

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