{"record":{"id":"aa0a18787f892b5a","repo":"jax-ml/jax","slug":"reductions-with-constants-not-supported","errorCode":null,"errorMessage":"Reductions with constants not supported.","messagePattern":"Reductions with constants not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":2491,"sourceCode":"\ndef _reduction_lowering(body, ctx: LoweringRuleContext, a, axes):\n  flat_args = tree_util.tree_leaves(a)\n  (axis,) = axes\n\n  a_structure = tree_util.tree_structure(a)\n  avals_tree = tree_util.tree_unflatten(a_structure, ctx.avals_in)\n  mapped_avals_tree = tree_util.tree_map(\n      lambda aval: jax_core.ShapedArray((), aval.dtype), avals_tree\n  )\n  in_avals_ft = ft.flatten(((mapped_avals_tree, mapped_avals_tree), {}))\n\n  debug_info = api_util.debug_info(\"pallas triton reduction\", body, (a, a), {})\n  combine_jaxpr, _ = pe.trace_to_jaxpr(\n      body, in_avals_ft, debug_info=debug_info\n  )\n\n  if combine_jaxpr.consts:\n    raise NotImplementedError(\"Reductions with constants not supported.\")\n  element_types = [_element_type(arg.type) for arg in flat_args]\n  reduce_op = tt_dialect.ReduceOp(flat_args, axis)\n  param_types = element_types * 2\n  entry = reduce_op.regions[0].blocks.append(*param_types)\n  with ir.InsertionPoint.at_block_begin(entry):\n    results = lower_jaxpr_to_triton_ir(\n        ctx.context, combine_jaxpr, None, *entry.arguments\n    )\n    tt_dialect.reduce_return(results)\n  reduce_op.verify()\n  return list(reduce_op.result)\n\n\ndef _reduce_lowering(body, ctx: LoweringRuleContext, a, *, axes, **kwargs):\n  assert isinstance(axes, tuple)\n  if not axes:\n    return a\n  while len(axes) > 1:","sourceCodeStart":2473,"sourceCodeEnd":2509,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L2473-L2509","documentation":"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.","triggerScenarios":"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.","commonSituations":"Defining custom reduce ops with scale factors, thresholds, or initialization constants captured by closure; porting lambdas that worked with lax.reduce.","solutions":["Make the combine function depend only on its two arguments; move constants into the array being reduced or into kernel inputs","Use built-in reductions (jnp.max/min/add) which have dedicated lowerings","Pass the constant as an extra operand so it appears in avals rather than consts"],"exampleFix":"# before\nscale = 0.5\nreduce_fn = lambda x, y: jnp.maximum(x, y * scale)  # captures constant\n\n# after\nreduce_fn = lambda x, y: jnp.maximum(x, y)\n# apply scaling before/after the reduction instead","handlingStrategy":"validation","validationCode":"# ensure combine fn uses only its two params; quick check:\nimport jax\njaxpr = jax.make_jaxpr(lambda x, y: combine_fn(x, y))(a, b)\nassert not jaxpr.consts, 'combine fn captures constants'","typeGuard":"def pure_combine(fn, x, y) -> bool:\n    import jax\n    return not jax.make_jaxpr(fn)(x, y).consts","tryCatchPattern":null,"preventionTips":["Write combine functions as pure two-argument lambdas","Pass external scalars as kernel operands, not closures"],"tags":["jax","pallas","triton","reduction","closures"],"backgroundTag":"closure-captured-constant-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}