{"record":{"id":"6032ad78badb9a68","repo":"jax-ml/jax","slug":"the-denominator-cannot-be-unreduced-passed-to-div","errorCode":null,"errorMessage":"The denominator cannot be unreduced passed to `div`. Got {y=}","messagePattern":"The denominator cannot be unreduced passed to `div`\\. Got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5153,"sourceCode":"dispatch.simple_impl(mulhi_p)\nad.defjvp_zero(mulhi_p)\nmlir.register_lowering(mulhi_p, partial(_nary_lower_hlo, chlo.mulhi))\n\n\ndef _div_transpose_rule(cotangent, x, y):\n  assert ad.is_undefined_primal(x)\n  if ad.is_undefined_primal(y):\n    raise RuntimeError(\"nonlinear div can't be transposed\")\n  if type(cotangent) is ad_util.Zero:\n    return [ad_util.Zero(x.aval), None]\n  else:\n    return [_unbroadcast(x.aval, div(cotangent, y)), None]\n\ndef _div_ur_rule(x, y):\n  out_reduced = default_nary_reduced_rule(x, y)\n  x_ur, y_ur = getu(x), getu(y)\n  if y_ur:\n    raise ValueError(\n        f'The denominator cannot be unreduced passed to `div`. Got {y=}')\n  if x_ur and x_ur != getr(y):\n    raise ValueError(\n        'Denominator should be reduced along the same axes numerator is'\n        f' unreduced on. Got {x=}, {y=}')\n  out_unreduced = x_ur\n  if out_unreduced:\n    assert out_reduced == out_unreduced, (out_reduced, out_unreduced)\n    out_reduced = frozenset()  # if both are equal, set difference is empty.\n  kind = UnreducedKind.sum if out_unreduced else None\n  return out_unreduced, out_reduced, kind\n\ndiv_p = standard_naryop([_num, _num], 'div', ur_rule=_div_ur_rule)\nad.defjvp(div_p,\n          lambda g, x, y: div(g, y),\n          lambda g, x, y: mul(mul(neg(g), x), integer_pow(y, -2)))\nad.primitive_transposes[div_p] = _div_transpose_rule\nmlir.register_lowering(div_p, partial(_nary_lower_hlo, hlo.divide))","sourceCodeStart":5135,"sourceCodeEnd":5171,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5135-L5171","documentation":"In sharding propagation, the denominator of lax.div must be fully reduced: an unreduced y has no well-defined per-element division under partial shardings, so _div_ur_rule raises ValueError showing y.","triggerScenarios":"Sharded division where the denominator operand is unreduced along some axes (partial/reduced sharding mismatch) under NamedSharding/GSPMD jit.","commonSituations":"Dividing partially-sharded activations by partially-sharded normalizers (RMSNorm/LayerNorm-style) in sharded training; mixing reduced sums with unreduced tensors in elementwise division inside sharded jits.","solutions":["Constrain the denominator to a fully reduced sharding: with_sharding_constraint(y, P()) before dividing","Reshard/reduce y appropriately (all-reduce the denominator) so it is reduced on the numerator's unreduced axes","Reorder ops: normalize the denominator outside the sharded region or use lax.div after matching shardings"],"exampleFix":"// before\nout = x / y  # y unreduced under sharding propagation\n// after\ny_red = jax.lax.with_sharding_constraint(y, P())\nout = x / y_red","handlingStrategy":"validation","validationCode":"y = jax.lax.with_sharding_constraint(y, P()) if propagating_shardings else y\nout = lax.div(x, y)","typeGuard":null,"tryCatchPattern":"try:\n    out = lax.div(x, y)\nexcept ValueError as e:\n    if 'denominator cannot be unreduced' in str(e):\n        out = lax.div(x, jax.lax.with_sharding_constraint(y, P()))\n    else:\n        raise","preventionTips":["Ensure denominators are reduced/replicated in sharded norms","All-reduce normalizers before elementwise division in sharded jits"],"tags":["jax","sharding","div","unreduced","value-error"],"backgroundTag":"sharding-rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}