{"record":{"id":"752e029301212f49","repo":"jax-ml/jax","slug":"start-indices-arguments-to-dynamic-update-slice-mu","errorCode":null,"errorMessage":"start_indices arguments to dynamic_update_slice must be scalars, got indices {start_indices}","messagePattern":"start_indices arguments to dynamic_update_slice must be scalars, got indices (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/slicing.py","lineNumber":1655,"sourceCode":"\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\"\n        \" same axes. Got operand sharding\"\n        f\" {operand.str_short(mesh_axis_types=True)} and update sharding\"","sourceCodeStart":1637,"sourceCodeEnd":1673,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/slicing.py#L1637-L1673","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"// before\nstarts = jnp.arange(3)  # shape (3,)\nout = lax.dynamic_update_slice(buf, upd, starts)  # TypeError\n\n// after\nout = buf\nfor k in range(3):\n    out = lax.dynamic_update_slice(out, upd, (k, 0))","handlingStrategy":"type-guard","validationCode":"starts = tuple(i.reshape(()) if hasattr(i, 'reshape') else jnp.asarray(i) for i in starts)\nassert all(jnp.asarray(i).ndim == 0 for i in starts)","typeGuard":"def are_scalar_indices(starts) -> bool:\n    return all(jnp.asarray(s).ndim == 0 for s in starts)","tryCatchPattern":null,"preventionTips":["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."],"tags":["jax","lax","dynamic-update-slice","scalar-indices"],"backgroundTag":"index-not-scalar","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}