jax-ml/jax · error · NotImplementedError

Mosaic kernels cannot be automatically partitioned. Please w

Error message

Mosaic kernels cannot be automatically partitioned. Please wrap the call in a shard_map.

What it means

A Mosaic (Pallas TPU) kernel is being lowered under an SPMD axis context where only some axes are manual; automatic partitioning of Mosaic kernels is not supported, so all mesh axes must be manual or the call must be wrapped in shard_map.

Source

Thrown at jax/_src/tpu_custom_call.py:425

def _tpu_custom_call_lowering(
    ctx: mlir.LoweringRuleContext,
    *in_nodes,
    config: CustomCallBackendConfig,
    has_side_effects: TpuSideEffectType,
    kernel_name: str | None,
    out_avals: Any,
    input_output_aliases: tuple[tuple[int, int], ...],
    metadata: Any | None,
) -> 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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the Pallas call in jax.experimental.shard_map.shard_map
2. Make all mesh axes manual (e.g. run the kernel region inside a manual computation) or run unsharded

Example fix

// before
out = jax.jit(my_pallas_fn)(x)
// after
from jax.experimental.shard_map import shard_map
out = jax.jit(shard_map(my_pallas_fn, mesh, in_specs=P('x'), out_specs=P('x')))(x)

mesh = jax.make_mesh(...)
assert manual_axes == set(mesh.axis_names) before calling a Pallas kernel

try:
    jax.jit(f)(x)
except NotImplementedError as e:
    if 'shard_map' in str(e):
        f_s = shard_map(f, mesh, ...)
    else:
        raise

When it happens

Trigger: jit with a sharded mesh where manual_axes is a strict subset of mesh axis names and the computation contains a Pallas TPU kernel.

Common situations: Mixing sharded jitted computations with Pallas kernels under multi-host or multi-slice meshes.

Related errors

• Partitioned callback not implemented on {platform} backend.
• Mesh must be provided for shard_map with checkify.
• Custom Partitioning rules must return Sharding.
• Explicit sharding is not currently supported in the pallas-t
• 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/c7e2466d8966180f. Report an issue: GitHub.