{"record":{"id":"9ca06bc0f7729410","repo":"jax-ml/jax","slug":"expected-value-and-mask-to-have-the-same-shape-bu","errorCode":null,"errorMessage":"Expected value and mask to have the same shape, but got value shape {val_aval.shape} vs. mask shape {mask_aval.shape}.","messagePattern":"Expected value and mask to have the same shape, but got value shape (.+?) vs\\. mask shape (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2482,"sourceCode":"\n\n@register_lowering_rule(primitives.swap_p, ensure_mlir_values=False)\ndef _masked_swap_lowering_rule(\n    ctx: LoweringRuleContext, *args_flat, args_tree, **_\n):\n  ref, transforms, val, mask = args_tree.unflatten(args_flat)\n  ref_aval, transforms_avals, val_aval, mask_aval = args_tree.unflatten(\n      ctx.avals_in\n  )\n  prev_transforms, idx = _canonicalize_transforms_to_indexer(\n      ref_aval, transforms, transforms_avals\n  )\n\n  if mask is not None:\n    if  val_aval.dtype.itemsize != 4:\n      raise NotImplementedError(\"masked swap with non-32-bit data\")\n    if val_aval.shape != mask_aval.shape:\n      raise ValueError(\n          \"Expected value and mask to have the same shape, but got\"\n          f\" value shape {val_aval.shape} vs. mask shape {mask_aval.shape}.\"\n      )\n\n  ref_block_shape, *_ = ctx.block_shapes\n  ref, ref_block_shape = _transform_ref(\n      ref, ref_aval, ref_block_shape, prev_transforms\n  )\n\n  ref_type = ir.MemRefType(ref.type)\n  memory_space = str(ref_type.memory_space)\n  is_smem_store = memory_space == \"#tpu.memory_space<smem>\"\n  is_vmem_store = memory_space == \"#tpu.memory_space<vmem>\"\n  (aval_out,) = ctx.avals_out\n  if not isinstance(val, ir.Value):\n    val = ir_constant(val, mlir_type=_dtype_to_ir_type(val_aval.dtype))\n  if not is_smem_store and not ref_block_shape:\n    raise NotImplementedError(","sourceCodeStart":2464,"sourceCodeEnd":2500,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2464-L2500","documentation":"Raised when a masked store's value shape differs from its mask shape. The TPU masked-swap lowering applies the mask lane-for-lane, so the mask must broadcast-exactly match the stored value's shape.","triggerScenarios":"pl.store(ref, val, mask=m) with val.shape != m.shape, e.g. mask computed from a different block shape or with extra/missing dims.","commonSituations":"Reusing a mask computed for a different block shape; forgetting to broadcast/reshape the mask after squeezing a dim; off-by-one in BlockSpec block shapes.","solutions":["Make the mask the same shape as the value: mask = jnp.broadcast_to(mask, val.shape)","Check that the BlockSpec block shapes for value and mask refs agree","Recompute the mask from the current block indices (pl.program_id / block_start)"],"exampleFix":"# before\npl.store(out_ref, val, mask=mask)  # mask shape (8, 128), val shape (128, 8)\n# after\npl.store(out_ref, val, mask=jnp.broadcast_to(mask.T, val.shape))","handlingStrategy":"validation","validationCode":"val, mask = jnp.broadcast_arrays(val, mask)\nassert val.shape == mask.shape\npl.store(ref, val, mask=mask)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always broadcast mask to value shape right before storing","Derive masks from the same block-shape math as the value"],"tags":["jax","pallas","tpu","shape-mismatch","masked-store"],"backgroundTag":"mask-value-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}