{"record":{"id":"7ad58cfa315b4f7b","repo":"jax-ml/jax","slug":"reduce-sum-on-operand-operand-str-short-true-is","errorCode":null,"errorMessage":"reduce_sum on operand {operand.str_short(True)} is not allowed when jax_allow_f16_reductions=False.","messagePattern":"reduce_sum on operand (.+?) is not allowed when jax_allow_f16_reductions=False\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":8586,"sourceCode":"\ndef _reduce_op_reduced_rule(operand, out_sharding, name):\n  if out_sharding is not None and out_sharding.spec.reduced:\n    raise ValueError(\n        f'out_sharding passed to {name} cannot be reduced. Got {out_sharding=}')\n  return getr(operand)\n\ndef _reduce_sum_ur_rule(operand, *, axes, out_sharding):\n  out_unreduced, kind = _reduce_op_unreduced_rule(\n      operand, axes, out_sharding, UnreducedKind.sum, 'reduce_sum')\n  out_reduced = _reduce_op_reduced_rule(operand, out_sharding, 'reduce_sum')\n  return out_unreduced, out_reduced, kind\n\ndef _reduce_sum_dtype_rule(operand, *, axes, **_):\n  dt = _reduce_number_dtype_rule('reduce_sum', operand)\n  if (operand.dtype in [np.float16, dtypes.bfloat16] and\n      not config.allow_f16_reductions.value and\n      not all(core.definitely_equal(operand.shape[d], 1) for d in axes)):\n    raise ValueError(f\"reduce_sum on operand {operand.str_short(True)} is not \"\n                     \"allowed when jax_allow_f16_reductions=False.\")\n  return dt\n\nreduce_sum_p = standard_primitive(\n  _reduce_op_shape_rule, _reduce_sum_dtype_rule,\n  'reduce_sum', sharding_rule=_reduce_op_sharding_rule_with_out_sharding,\n  vma_rule=partial(core.standard_vma_rule, 'reduce_sum'),\n  ur_rule=_reduce_sum_ur_rule)\nad.deflinear2(reduce_sum_p, _reduce_sum_transpose_rule)\nbatching.defreducer(reduce_sum_p)\n\ndef _reduce_prod_jvp_rule(primals, tangents, *, axes):\n  reducer = lambda x, y: [mul(x, y)]\n  primals_out, tangents_out = _reduce_jvp(reducer, [_const(primals[0], 1)],\n                                          primals, tangents, axes)\n  return primals_out[0], tangents_out[0]\n\ndef _reduce_op_sharding_rule(operand, *, axes):","sourceCodeStart":8568,"sourceCodeEnd":8604,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L8568-L8604","documentation":"JAX refuses reduce_sum on float16/bfloat16 inputs when the config flag jax_allow_f16_reductions is False, unless every reduced axis has size 1. This is because f16 accumulation on many backends loses precision or is unsupported, so it is gated behind an explicit opt-in flag.","triggerScenarios":"lax.reduce_sum(x16, axes) where x16 has dtype float16 or bfloat16 and any reduced axis is larger than 1, under default config. Often reached indirectly via jnp.sum on a half-precision array during tracing.","commonSituations":"Mixed-precision training pipelines that sum bf16 losses or gradients; TPU/GPU half-precision training; upgrading JAX where the flag default or enforcement changed. The flag defaults differently per backend/version.","solutions":["Cast the operand to float32 before reducing: jnp.sum(x.astype(jnp.float32), axis).astype(x.dtype).","Opt in explicitly: set jax.config.update('jax_allow_f16_reductions', True) (with JAX_JAX_ALLOW_F16_REDUCTIONS=True env var) if your backend's f16 reductions are acceptable.","Check which layer produced the half-precision sum (e.g., a loss function) and keep reductions in fp32 there."],"exampleFix":"# before\nloss = jnp.sum(squared_err_bf16)\n# after\nloss = jnp.sum(squared_err_bf16.astype(jnp.float32))","handlingStrategy":"validation","validationCode":"def safe_sum(x, axis=None):\n    if x.dtype in (jnp.float16, jnp.bfloat16):\n        x = x.astype(jnp.float32)\n    return jnp.sum(x, axis=axis)","typeGuard":"def needs_f32_upcast(x):\n    return x.dtype in (jnp.float16, jnp.bfloat16)","tryCatchPattern":null,"preventionTips":["Keep loss/gradient reductions in fp32 in mixed-precision pipelines.","Set jax_allow_f16_reductions explicitly if you rely on half reductions."],"tags":["jax","float16","bfloat16","reduction","mixed-precision","config-flag"],"backgroundTag":"half-precision-reduction-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}