jax-ml/jax · error · TypeError

dynamic_update_slice update shape must be smaller than opera

Error message

dynamic_update_slice update shape must be smaller than operand shape, got update shape {} for operand shape {}.

What it means

Thrown by the shape rule of lax.dynamic_update_slice when the update array is larger than the operand array in at least one dimension. JAX requires update.shape[i] <= operand.shape[i] for every axis, because the operation writes update into a window of operand at the given start indices. The message prints both shapes so you can compare them directly.

Source

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

  out = mlir.dynamic_slice(ctx, aval_out, x, start_indices=start_indices)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

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"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make update no larger than operand in every dimension: pad or truncate the operand first (e.g. lax.pad / jnp.pad) or slice the update.
2. Check the printed shapes in the message and fix whichever array's shape was computed incorrectly (often a hardcoded batch/time dim).
3. If you intended a full-array write, use operand.at[...].set(update) with compatible shapes instead.

Example fix

// before
operand = jnp.zeros((3, 4))
update = jnp.ones((5, 2))
out = lax.dynamic_update_slice(operand, update, (0, 0))  # TypeError

// after
operand = jnp.zeros((5, 4))
update = jnp.ones((5, 2))
out = lax.dynamic_update_slice(operand, update, (0, 0))

def safe_dus(operand, update, starts):
    assert all(o >= u for o, u in zip(operand.shape, update.shape)), \
        f"update {update.shape} larger than operand {operand.shape}"
    return lax.dynamic_update_slice(operand, update, starts)

def is_valid_update(operand: jax.Array, update: jax.Array) -> bool:
    return operand.ndim == update.ndim and all(
        o >= u for o, u in zip(operand.shape, update.shape))

When it happens

Trigger: Calling lax.dynamic_update_slice(operand, update, start_indices) (or lax.dynamic_update_slice_p) where any dimension of update exceeds the corresponding dimension of operand, e.g. operand shape (3, 4) with update shape (5, 2).

Common situations: Padding/overwriting buffers computed with the wrong batch or sequence length; off-by-one sizing when building sliding-window updates; passing a full array instead of a slice as the update after a refactor or shape change upstream in a pipeline.

Related errors

• conv_general_dilated lhs and rhs must have the same number o
• Arguments to sort must have equal shapes, got: {shapes}
• start_indices arguments to dynamic_update_slice must be scal
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out

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