{"record":{"id":"de7a6e3659047dca","repo":"jax-ml/jax","slug":"swap-only-supports-slices-with-stride-1-got-stri","errorCode":null,"errorMessage":"Swap only supports slices with stride 1, got {strides}","messagePattern":"Swap only supports slices with stride 1, got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_lowering.py","lineNumber":219,"sourceCode":"  ref_block_shape, *_ = ctx.block_shapes\n  ref, ref_block_shape = _transform_ref(\n      ref, ref_aval, ref_block_shape, prev_transforms\n  )\n  starts, sizes, strides, squeeze_dims, _ = tc_lowering._indexer_to_start_size_stride(\n      indexer, ref_block_shape, cast_to_index=True\n  )\n  for first_nontrivial_dim, s in enumerate(sizes):\n    if s != 1:\n      break\n  else:\n    first_nontrivial_dim = len(sizes)\n  if any(squeeze_dims[first_nontrivial_dim:]):\n    raise NotImplementedError(\n        \"Integer indexing of refs that follows a non-trivial slice is not\"\n        \" supported on SC\"\n    )\n  if not all(s == 1 for s in strides):\n    raise NotImplementedError(\n        \"Swap only supports slices with stride 1, got {strides}\"\n    )\n\n  if (out_aval.ndim == 0) != (ref_memory_space is MemorySpace.SMEM):\n    message = \"Swap only supports scalars in SMEM.\"\n    if ref_memory_space is MemorySpace.SMEM:\n      message += \" Trying to swap an array of shape {out_aval.shape}.\"\n    else:\n      message += f\" Trying to swap a scalar in {ref_memory_space!r}.\"\n    raise NotImplementedError(message)\n\n  if out_aval.ndim == 0:\n    if mask is not None:\n      raise NotImplementedError(\"Swap does not support masked scalar stores\")\n    if add:\n      # TODO(slebedev): We can use memref.atomic_rmw here, but the SC compiler\n      # doesn't support it yet.\n      raise NotImplementedError(\"Swap does not support atomic scalar adds\")","sourceCodeStart":201,"sourceCodeEnd":237,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_lowering.py#L201-L237","documentation":"SC swap/store only supports 0-d (scalar) values in SMEM. Trying to swap an array in SMEM, or a scalar in a non-SMEM space, raises NotImplementedError with this message.","triggerScenarios":"(out_aval.ndim == 0) mismatched with memory space: swapping an array-shaped value in SMEM, or a scalar in VMEM/HBM.","commonSituations":"Using SMEM buffers sized >1 element with swap semantics; scalar swaps attempted on VMEM refs.","solutions":["For scalar swaps, move the scalar into SMEM first","For array swaps in SMEM, use element-wise stores instead of swap","Load/swap arrays in VMEM and copy scalars through SMEM"],"exampleFix":"# before\nold = smem_buf[:]  # swap of array in SMEM\n# after\nold = smem_buf[0]  # scalar swap in SMEM only","handlingStrategy":"validation","validationCode":"assert (out.ndim == 0) == (space is MemorySpace.SMEM)","typeGuard":"null","tryCatchPattern":"null","preventionTips":["Only swap scalars, and only in SMEM"],"tags":["jax","pallas","tpu","sparsecore","swap","memory-space"],"backgroundTag":"unsupported-scalar-memory-access","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}