{"record":{"id":"aa0d683467dc00a2","repo":"jax-ml/jax","slug":"dynamic-update-slice-update-shape-must-be-smaller","errorCode":null,"errorMessage":"dynamic_update_slice update shape must be smaller than operand shape, got update shape {} for operand shape {}.","messagePattern":"dynamic_update_slice update shape must be smaller than operand shape, got update shape (.+?) for operand shape (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/slicing.py","lineNumber":1653,"sourceCode":"  out = mlir.dynamic_slice(ctx, aval_out, x, start_indices=start_indices)\n  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]\n\nmlir.register_lowering(dynamic_slice_p, _dynamic_slice_lower)\n\n\ndef _dynamic_update_slice_shape_rule(operand, update, *start_indices):\n  if operand.ndim != update.ndim:\n    msg = (\"dynamic_update_slice update must have the same rank as operand, \"\n           \"got update shape {} for operand shape {}.\")\n    raise TypeError(msg.format(update.shape, operand.shape))\n  if operand.ndim != len(start_indices):\n    msg = (\"dynamic_update_slice start_indices must have length equal to the \"\n           \"rank of operand, got indices {} for operand shape {}.\")\n    raise TypeError(msg.format(start_indices, operand.shape))\n  if not all(map(operator.ge, operand.shape, update.shape)):\n    msg = (\"dynamic_update_slice update shape must be smaller than operand \"\n           \"shape, got update shape {} for operand shape {}.\")\n    raise TypeError(msg.format(update.shape, operand.shape))\n  if any(idx.ndim != 0 for idx in start_indices):\n    raise TypeError(\"start_indices arguments to dynamic_update_slice must be \"\n                    f\"scalars, got indices {start_indices}\")\n  return operand.shape\n\ndef _dynamic_update_slice_sharding_rule(operand, update, *start_indices):\n  if operand.sharding != update.sharding:\n    raise core.ShardingTypeError(\n        \"dynamic_update_slice operand sharding must be equal to update\"\n        \" sharding, got operand sharding\"\n        f\" {operand.str_short(mesh_axis_types=True)} and update sharding\"\n        f\" {update.str_short(mesh_axis_types=True)}.\")\n  return operand.sharding\n\ndef _dus_unreduced_rule(operand, update):\n  if core.getu(operand) != core.getu(update):\n    raise core.ShardingTypeError(\n        \"dynamic_update_slice operand and update must be unreduced along the\"","sourceCodeStart":1635,"sourceCodeEnd":1671,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/slicing.py#L1635-L1671","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["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.","Check the printed shapes in the message and fix whichever array's shape was computed incorrectly (often a hardcoded batch/time dim).","If you intended a full-array write, use operand.at[...].set(update) with compatible shapes instead."],"exampleFix":"// before\noperand = jnp.zeros((3, 4))\nupdate = jnp.ones((5, 2))\nout = lax.dynamic_update_slice(operand, update, (0, 0))  # TypeError\n\n// after\noperand = jnp.zeros((5, 4))\nupdate = jnp.ones((5, 2))\nout = lax.dynamic_update_slice(operand, update, (0, 0))","handlingStrategy":"validation","validationCode":"def safe_dus(operand, update, starts):\n    assert all(o >= u for o, u in zip(operand.shape, update.shape)), \\\n        f\"update {update.shape} larger than operand {operand.shape}\"\n    return lax.dynamic_update_slice(operand, update, starts)","typeGuard":"def is_valid_update(operand: jax.Array, update: jax.Array) -> bool:\n    return operand.ndim == update.ndim and all(\n        o >= u for o, u in zip(operand.shape, update.shape))","tryCatchPattern":null,"preventionTips":["Compare operand.shape and update.shape element-wise before every dynamic_update_slice call.","Pad the operand to the maximum needed size upfront so updates always fit.","Write a unit test asserting update fits whenever buffer sizes are config-driven."],"tags":["jax","lax","dynamic-update-slice","shape-mismatch"],"backgroundTag":"array-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}