{"record":{"id":"b044b5d384cfcc94","repo":"jax-ml/jax","slug":"broadcast-to-is-a-triton-specific-primitive-ple","errorCode":null,"errorMessage":"`broadcast_to` is a Triton-specific primitive. Please consider using `jnp.broadcast_to` instead.","messagePattern":"`broadcast_to` is a Triton-specific primitive\\. Please consider using `jnp\\.broadcast_to` instead\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2705,"sourceCode":"\n@register_lowering_rule(lax.reduce_or_p, kernel_types=[*tpu_core.CoreType])\ndef _reduce_or_lowering_rule(ctx: LoweringRuleContext, x, *, axes):\n  def _proxy_reduce(arg, *, axes):\n    # Mosaic currently only supports float reductions, so we cast the boolean\n    # arg to a float and use reduce_max to implement reduce_or.\n    # TODO(b/351017807): Implement this logic in Mosaic MultiDimReductionOp\n    # instead.\n    float_arg = jnp.where(arg, 1.0, 0.0)\n    return jnp.max(float_arg, axis=axes) > 0.0\n  proxy_lowering = lower_fun(_proxy_reduce)\n  return proxy_lowering(ctx, x, axes=axes)\n\n\n@register_lowering_rule(state_primitives.broadcast_to_p)\ndef _broadcast_to_lowering_rule(\n    ctx: LoweringRuleContext, x, shape: Sequence[int]\n):\n  raise RuntimeError(\n      \"`broadcast_to` is a Triton-specific primitive. Please consider using\"\n      \" `jnp.broadcast_to` instead.\"\n  )\n\n\n@register_lowering_rule(\n    lax.broadcast_in_dim_p, kernel_types=[*tpu_core.CoreType]\n)\ndef _broadcast_in_dim_lowering_rule(\n    ctx: LoweringRuleContext, val, *, shape, broadcast_dimensions, sharding\n):\n  del sharding\n  (aval_in,) = ctx.avals_in\n  (aval_out,) = ctx.avals_out\n  if aval_in.shape == shape:\n    return val\n\n  if broadcast_dimensions:","sourceCodeStart":2687,"sourceCodeEnd":2723,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2687-L2723","documentation":"Pallas Mosaic (TPU backend) has no lowering rule for the Triton-specific `broadcast_to` primitive, so the registered rule deliberately raises to tell you this operation only exists in the Triton path. Use `jnp.broadcast_to` inside your kernel body instead, which Mosaic knows how to lower.","triggerScenarios":"Calling `state_primitives.broadcast_to` (or importing broadcast_to from jax.experimental.pallas and using it) inside a Pallas kernel that is compiled for TPU/Mosaic instead of Triton.","commonSituations":"Porting a Triton-targeted Pallas kernel to TPU; copy-pasting kernel code written for the GPU backend.","solutions":["Replace broadcast_to(...) with jnp.broadcast_to(...) in the kernel body","Avoid state-primitive broadcasts on TPU; restructure the kernel to compute the broadcast via normal jnp ops"],"exampleFix":"// before\nimport jax.experimental.pallas as pallas\nx_b = pl.broadcast_to(x, out_shape)\n// after\nimport jax.numpy as jnp\nx_b = jnp.broadcast_to(x, out_shape)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\n# in kernel body always prefer:\ny = jnp.broadcast_to(x, shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't import Triton-only primitives into kernels targeting TPU","Write kernels against jnp.* ops which all backends lower"],"tags":["jax","pallas","tpu","mosaic","broadcast"],"backgroundTag":"unsupported-operation-on-backend","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}