jax-ml/jax · error · NotImplementedError

masked swap with strided store

Error message

masked swap with strided store

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:2576

    result_vec_type = ir.VectorType.get(
        ctx.lowering_context.dynamic_shape_replacement_fn(aval_out.shape),
      _dtype_to_ir_type(aval_out.dtype, is_kernel_boundary=True))
    result = vector.shape_cast(result_vec_type, result)
    val_vec_type = ir.VectorType.get(
        ctx.lowering_context.dynamic_shape_replacement_fn(mem_aval.shape),
      _dtype_to_ir_type(mem_aval.dtype, is_kernel_boundary=True))
    val = vector.shape_cast(val_vec_type, val)
    if mask is not None:
      mask_vec_type = ir.VectorType.get(
          ctx.lowering_context.dynamic_shape_replacement_fn(mem_aval.shape),
          _dtype_to_ir_type(mask_aval.dtype)
      )
      mask = vector.shape_cast(mask_vec_type, mask)
  result = _maybe_cast_load_to_bool(ctx, val_aval, result)

  if need_stride:
    if mask is not None:
      raise NotImplementedError("masked swap with strided store")
    tpu.strided_store(val, ref, starts, strides)
  else:
    tpu.vector_store(val, ref, starts, strides=[], mask=mask)
  return result


@register_lowering_rule(
    primitives.multiple_of_p, kernel_types=[*tpu_core.CoreType]
)
def _multiple_of_lowering_rule(ctx: LoweringRuleContext, val, *, values):
  del ctx
  for multiple in values:
    val = tpu.assume_multiple(val, multiple)
  return val


def reduce_lowering_rule(reduce_fn, type_to_kind, type_to_identity):
  def _lowering_rule(ctx: LoweringRuleContext, x, *, axes, **kwargs):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the stride: materialize the transposed/dilated block in a VMEM scratch and store contiguously
2. Remove the mask (compute full block with where and store unmasked)
3. Transpose the value before storing so strides become unit

Example fix

# before
pl.store(ref.T_view, val, mask=m)  # strided + masked
# after
val = jnp.where(m, val, pl.load(ref.T_view))
pl.store(ref.T_view, val)  # unmasked strided store is fine

def safe_strided_store(ref, val, mask):
    if mask is not None:
        val = jnp.where(mask, val, pl.load(ref))
        pl.store(ref, val)  # unmasked strided store
    else:
        pl.store(ref, val)

When it happens

Trigger: 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.

Common situations: Storing transposed outputs with predication; applying column-masked writes to a row-major block accessed with stride.

Related errors

• masked swap with non-32-bit data
• Expected value and mask to have the same shape, but got valu
• SMEM store does not support masks
• Swap only supports scalars in SMEM.
• Compiler params for platform {platform} cannot be used for {

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/bf3fa13444b8ab66. Report an issue: GitHub.