jax-ml/jax · error · NotImplementedError

Reductions with constants not supported.

Error message

Reductions with constants not supported.

What it means

When lowering a Pallas reduction, the combine function is traced to a jaxpr; if that jaxpr closes over constants (e.g. captures a Python/JAX constant from the enclosing scope), the lowering cannot build the Triton ReduceOp region and raises NotImplementedError.

Source

Thrown at jax/_src/pallas/triton/lowering.py:2491

def _reduction_lowering(body, ctx: LoweringRuleContext, a, axes):
  flat_args = tree_util.tree_leaves(a)
  (axis,) = axes

  a_structure = tree_util.tree_structure(a)
  avals_tree = tree_util.tree_unflatten(a_structure, ctx.avals_in)
  mapped_avals_tree = tree_util.tree_map(
      lambda aval: jax_core.ShapedArray((), aval.dtype), avals_tree
  )
  in_avals_ft = ft.flatten(((mapped_avals_tree, mapped_avals_tree), {}))

  debug_info = api_util.debug_info("pallas triton reduction", body, (a, a), {})
  combine_jaxpr, _ = pe.trace_to_jaxpr(
      body, in_avals_ft, debug_info=debug_info
  )

  if combine_jaxpr.consts:
    raise NotImplementedError("Reductions with constants not supported.")
  element_types = [_element_type(arg.type) for arg in flat_args]
  reduce_op = tt_dialect.ReduceOp(flat_args, axis)
  param_types = element_types * 2
  entry = reduce_op.regions[0].blocks.append(*param_types)
  with ir.InsertionPoint.at_block_begin(entry):
    results = lower_jaxpr_to_triton_ir(
        ctx.context, combine_jaxpr, None, *entry.arguments
    )
    tt_dialect.reduce_return(results)
  reduce_op.verify()
  return list(reduce_op.result)


def _reduce_lowering(body, ctx: LoweringRuleContext, a, *, axes, **kwargs):
  assert isinstance(axes, tuple)
  if not axes:
    return a
  while len(axes) > 1:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the combine function depend only on its two arguments; move constants into the array being reduced or into kernel inputs
2. Use built-in reductions (jnp.max/min/add) which have dedicated lowerings
3. Pass the constant as an extra operand so it appears in avals rather than consts

Example fix

# before
scale = 0.5
reduce_fn = lambda x, y: jnp.maximum(x, y * scale)  # captures constant

# after
reduce_fn = lambda x, y: jnp.maximum(x, y)
# apply scaling before/after the reduction instead

# ensure combine fn uses only its two params; quick check:
import jax
jaxpr = jax.make_jaxpr(lambda x, y: combine_fn(x, y))(a, b)
assert not jaxpr.consts, 'combine fn captures constants'

def pure_combine(fn, x, y) -> bool:
    import jax
    return not jax.make_jaxpr(fn)(x, y).consts

When it happens

Trigger: Writing a pallas reduction whose binary operator references a captured constant, e.g. lambda x, y: jnp.maximum(x, y * 0.5) or a min with a non-input scalar, so trace_to_jaxpr yields non-empty consts.

Common situations: Defining custom reduce ops with scale factors, thresholds, or initialization constants captured by closure; porting lambdas that worked with lax.reduce.

Related errors

• `pallas` reduce operations only support one reduce axis.
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Only single axis reduction supported
• reductions require axes to be (0,) on SparseCore, but got {a

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