{"record":{"id":"bf3fa13444b8ab66","repo":"jax-ml/jax","slug":"masked-swap-with-strided-store","errorCode":null,"errorMessage":"masked swap with strided store","messagePattern":"masked swap with strided store","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2576,"sourceCode":"    result_vec_type = ir.VectorType.get(\n        ctx.lowering_context.dynamic_shape_replacement_fn(aval_out.shape),\n      _dtype_to_ir_type(aval_out.dtype, is_kernel_boundary=True))\n    result = vector.shape_cast(result_vec_type, result)\n    val_vec_type = ir.VectorType.get(\n        ctx.lowering_context.dynamic_shape_replacement_fn(mem_aval.shape),\n      _dtype_to_ir_type(mem_aval.dtype, is_kernel_boundary=True))\n    val = vector.shape_cast(val_vec_type, val)\n    if mask is not None:\n      mask_vec_type = ir.VectorType.get(\n          ctx.lowering_context.dynamic_shape_replacement_fn(mem_aval.shape),\n          _dtype_to_ir_type(mask_aval.dtype)\n      )\n      mask = vector.shape_cast(mask_vec_type, mask)\n  result = _maybe_cast_load_to_bool(ctx, val_aval, result)\n\n  if need_stride:\n    if mask is not None:\n      raise NotImplementedError(\"masked swap with strided store\")\n    tpu.strided_store(val, ref, starts, strides)\n  else:\n    tpu.vector_store(val, ref, starts, strides=[], mask=mask)\n  return result\n\n\n@register_lowering_rule(\n    primitives.multiple_of_p, kernel_types=[*tpu_core.CoreType]\n)\ndef _multiple_of_lowering_rule(ctx: LoweringRuleContext, val, *, values):\n  del ctx\n  for multiple in values:\n    val = tpu.assume_multiple(val, multiple)\n  return val\n\n\ndef reduce_lowering_rule(reduce_fn, type_to_kind, type_to_identity):\n  def _lowering_rule(ctx: LoweringRuleContext, x, *, axes, **kwargs):","sourceCodeStart":2558,"sourceCodeEnd":2594,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2558-L2594","documentation":"Raised when a masked swap/store uses strided access (non-unit strides). TPU strided_store has no mask parameter, so combining a mask with a strided (e.g. transposed or dilated) store is unimplemented.","triggerScenarios":"pl.swap/pl.store with a mask where the indexer produced strides != 1, e.g. storing into a transposed view or a ref sliced with steps.","commonSituations":"Storing transposed outputs with predication; applying column-masked writes to a row-major block accessed with stride.","solutions":["Remove the stride: materialize the transposed/dilated block in a VMEM scratch and store contiguously","Remove the mask (compute full block with where and store unmasked)","Transpose the value before storing so strides become unit"],"exampleFix":"# before\npl.store(ref.T_view, val, mask=m)  # strided + masked\n# after\nval = jnp.where(m, val, pl.load(ref.T_view))\npl.store(ref.T_view, val)  # unmasked strided store is fine","handlingStrategy":"fallback","validationCode":"def safe_strided_store(ref, val, mask):\n    if mask is not None:\n        val = jnp.where(mask, val, pl.load(ref))\n        pl.store(ref, val)  # unmasked strided store\n    else:\n        pl.store(ref, val)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't combine masks with transposed/strided stores; materialize or use where","Transpose values in registers before storing"],"tags":["jax","pallas","tpu","strided-store","masked-store"],"backgroundTag":"strided-masked-store-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}