{"record":{"id":"18aeafe481aa18fc","repo":"jax-ml/jax","slug":"replica-lowering-for-mosaic-kernels-not-implemente","errorCode":null,"errorMessage":"Replica lowering for Mosaic kernels not implemented.","messagePattern":"Replica lowering for Mosaic kernels not implemented\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu_custom_call.py","lineNumber":436,"sourceCode":") -> 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))\n        for aval_out in ctx.avals_out\n    ])\n  extra_attributes: dict[str, ir.Attribute] | None = None\n  # Add kernel_name and kernel_metadata as attributes to the custom call op.\n  # This is because we do not want to pollute the backend_config with this\n  # information.\n  if kernel_name is not None:\n    extra_attributes = dict(kernel_name=ir.StringAttr.get(kernel_name))\n  # If the IR version we originally generated the ASM string with is not the\n  # same as the one we should have used, we need to downgrade the ASM string.\n  ir_version = get_ir_version(ctx)","sourceCodeStart":418,"sourceCodeEnd":454,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu_custom_call.py#L418-L454","documentation":"A Mosaic kernel that performs inter-device communication is being lowered in replicated mode, which has no implementation.","triggerScenarios":"A Pallas kernel using communication primitives (e.g. collective barriers/DMAs between devices) lowered with has_communication=True outside an SPMD/manual context.","commonSituations":"Porting cross-core communication kernels from Mosaic to Pallas without shard_map or a proper mesh.","solutions":["Wrap the kernel in shard_map so communication is explicitly managed","Remove/disable communication ops if replication was intended","Upgrade JAX where replica lowering support may exist"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    out = kernel(x)\nexcept NotImplementedError as e:\n    if 'Replica lowering' in str(e):\n        out = shard_map(kernel, mesh, ...)(x)\n    else:\n        raise","preventionTips":["Avoid cross-device communication ops in replicated Pallas kernels","Document which kernels require shard_map"],"tags":["tpu","pallas","communication","replication","jax"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}