jax-ml/jax · error · RuntimeError

`broadcast_to` is a Triton-specific primitive. Please consid

Error message

`broadcast_to` is a Triton-specific primitive. Please consider using `jnp.broadcast_to` instead.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:2705

@register_lowering_rule(lax.reduce_or_p, kernel_types=[*tpu_core.CoreType])
def _reduce_or_lowering_rule(ctx: LoweringRuleContext, x, *, axes):
  def _proxy_reduce(arg, *, axes):
    # Mosaic currently only supports float reductions, so we cast the boolean
    # arg to a float and use reduce_max to implement reduce_or.
    # TODO(b/351017807): Implement this logic in Mosaic MultiDimReductionOp
    # instead.
    float_arg = jnp.where(arg, 1.0, 0.0)
    return jnp.max(float_arg, axis=axes) > 0.0
  proxy_lowering = lower_fun(_proxy_reduce)
  return proxy_lowering(ctx, x, axes=axes)


@register_lowering_rule(state_primitives.broadcast_to_p)
def _broadcast_to_lowering_rule(
    ctx: LoweringRuleContext, x, shape: Sequence[int]
):
  raise RuntimeError(
      "`broadcast_to` is a Triton-specific primitive. Please consider using"
      " `jnp.broadcast_to` instead."
  )


@register_lowering_rule(
    lax.broadcast_in_dim_p, kernel_types=[*tpu_core.CoreType]
)
def _broadcast_in_dim_lowering_rule(
    ctx: LoweringRuleContext, val, *, shape, broadcast_dimensions, sharding
):
  del sharding
  (aval_in,) = ctx.avals_in
  (aval_out,) = ctx.avals_out
  if aval_in.shape == shape:
    return val

  if broadcast_dimensions:

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
import jax.experimental.pallas as pallas
x_b = pl.broadcast_to(x, out_shape)
// after
import jax.numpy as jnp
x_b = jnp.broadcast_to(x, out_shape)

Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp
# in kernel body always prefer:
y = jnp.broadcast_to(x, shape)

Prevention

Don't import Triton-only primitives into kernels targeting TPU
Write kernels against jnp.* ops which all backends lower

When it happens

Trigger: 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.

Common situations: Porting a Triton-targeted Pallas kernel to TPU; copy-pasting kernel code written for the GPU backend.

Related errors

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