jax-ml/jax · error · NotImplementedError

Replica lowering for Mosaic kernels not implemented.

Error message

Replica lowering for Mosaic kernels not implemented.

What it means

A Mosaic kernel that performs inter-device communication is being lowered in replicated mode, which has no implementation.

Source

Thrown at jax/_src/tpu_custom_call.py:436

) -> ir.OpResultList:
  result_types, _ = mlir.ir_tree_registry.flatten([mlir.aval_to_ir_types(ctx.module_context, aval) for aval in out_avals])
  axis_context = ctx.module_context.axis_context
  if isinstance(axis_context, sharding_impls.SPMDAxisContext):
    manual_axes = axis_context.manual_axes | set(axis_context.mesh.manual_axes)
    if (axis_context.manual_axes and
        manual_axes != frozenset(axis_context.mesh.axis_names)):
      raise NotImplementedError(
          "Mosaic kernels cannot be automatically partitioned. Please wrap the"
          " call in a shard_map."
      )
  elif isinstance(axis_context, sharding_impls.ShardingContext):
    if axis_context.num_devices != 1:
      raise NotImplementedError(
          "Mosaic kernels cannot be automatically partitioned. Please wrap the"
          " call in a shard_map."
      )
  elif config.has_communication:
    raise NotImplementedError(
        "Replica lowering for Mosaic kernels not implemented."
    )
  if all(core.is_constant_shape(aval_out.shape) for aval_out in ctx.avals_out):
    result_shapes = None
  else:
    result_shapes, _ = mlir.ir_tree_registry.flatten([
        mlir.shape_tensor(ctx.module_context, mlir.eval_dynamic_shape(ctx, aval_out.shape))
        for aval_out in ctx.avals_out
    ])
  extra_attributes: dict[str, ir.Attribute] | None = None
  # Add kernel_name and kernel_metadata as attributes to the custom call op.
  # This is because we do not want to pollute the backend_config with this
  # information.
  if kernel_name is not None:
    extra_attributes = dict(kernel_name=ir.StringAttr.get(kernel_name))
  # If the IR version we originally generated the ASM string with is not the
  # same as the one we should have used, we need to downgrade the ASM string.
  ir_version = get_ir_version(ctx)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the kernel in shard_map so communication is explicitly managed
2. Remove/disable communication ops if replication was intended
3. Upgrade JAX where replica lowering support may exist

try:
    out = kernel(x)
except NotImplementedError as e:
    if 'Replica lowering' in str(e):
        out = shard_map(kernel, mesh, ...)(x)
    else:
        raise

When it happens

Trigger: A Pallas kernel using communication primitives (e.g. collective barriers/DMAs between devices) lowered with has_communication=True outside an SPMD/manual context.

Common situations: Porting cross-core communication kernels from Mosaic to Pallas without shard_map or a proper mesh.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}
• Accumulators are not available on TPU {info.chip_version}

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