{"record":{"id":"c7e2466d8966180f","repo":"jax-ml/jax","slug":"mosaic-kernels-cannot-be-automatically-partitioned","errorCode":null,"errorMessage":"Mosaic kernels cannot be automatically partitioned. Please wrap the call in a shard_map.","messagePattern":"Mosaic kernels cannot be automatically partitioned\\. Please wrap the call in a shard_map\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu_custom_call.py","lineNumber":425,"sourceCode":"\n\ndef _tpu_custom_call_lowering(\n    ctx: mlir.LoweringRuleContext,\n    *in_nodes,\n    config: CustomCallBackendConfig,\n    has_side_effects: TpuSideEffectType,\n    kernel_name: str | None,\n    out_avals: Any,\n    input_output_aliases: tuple[tuple[int, int], ...],\n    metadata: Any | None,\n) -> ir.OpResultList:\n  result_types, _ = mlir.ir_tree_registry.flatten([mlir.aval_to_ir_types(ctx.module_context, aval) for aval in out_avals])\n  axis_context = ctx.module_context.axis_context\n  if isinstance(axis_context, sharding_impls.SPMDAxisContext):\n    manual_axes = axis_context.manual_axes | set(axis_context.mesh.manual_axes)\n    if (axis_context.manual_axes and\n        manual_axes != frozenset(axis_context.mesh.axis_names)):\n      raise NotImplementedError(\n          \"Mosaic kernels cannot be automatically partitioned. Please wrap the\"\n          \" call in a shard_map.\"\n      )\n  elif isinstance(axis_context, sharding_impls.ShardingContext):\n    if axis_context.num_devices != 1:\n      raise NotImplementedError(\n          \"Mosaic kernels cannot be automatically partitioned. Please wrap the\"\n          \" call in a shard_map.\"\n      )\n  elif config.has_communication:\n    raise NotImplementedError(\n        \"Replica lowering for Mosaic kernels not implemented.\"\n    )\n  if all(core.is_constant_shape(aval_out.shape) for aval_out in ctx.avals_out):\n    result_shapes = None\n  else:\n    result_shapes, _ = mlir.ir_tree_registry.flatten([\n        mlir.shape_tensor(ctx.module_context, mlir.eval_dynamic_shape(ctx, aval_out.shape))","sourceCodeStart":407,"sourceCodeEnd":443,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu_custom_call.py#L407-L443","documentation":"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.","triggerScenarios":"jit with a sharded mesh where manual_axes is a strict subset of mesh axis names and the computation contains a Pallas TPU kernel.","commonSituations":"Mixing sharded jitted computations with Pallas kernels under multi-host or multi-slice meshes.","solutions":["Wrap the Pallas call in jax.experimental.shard_map.shard_map","Make all mesh axes manual (e.g. run the kernel region inside a manual computation) or run unsharded"],"exampleFix":"// before\nout = jax.jit(my_pallas_fn)(x)\n// after\nfrom jax.experimental.shard_map import shard_map\nout = jax.jit(shard_map(my_pallas_fn, mesh, in_specs=P('x'), out_specs=P('x')))(x)","handlingStrategy":"type-guard","validationCode":"mesh = jax.make_mesh(...)\nassert manual_axes == set(mesh.axis_names) before calling a Pallas kernel","typeGuard":null,"tryCatchPattern":"try:\n    jax.jit(f)(x)\nexcept NotImplementedError as e:\n    if 'shard_map' in str(e):\n        f_s = shard_map(f, mesh, ...)\n    else:\n        raise","preventionTips":["Always wrap Pallas TPU kernels in shard_map when using meshes","Keep Pallas kernels out of partially-manual SPMD regions"],"tags":["tpu","pallas","sharding","shard-map","jax"],"backgroundTag":"sharding-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}