{"record":{"id":"a303bf640635df7f","repo":"jax-ml/jax","slug":"denominator-should-be-reduced-along-the-same-axes","errorCode":null,"errorMessage":"Denominator should be reduced along the same axes numerator is unreduced on. Got {x=}, {y=}","messagePattern":"Denominator should be reduced along the same axes numerator is unreduced on\\. Got (.+?), (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5156,"sourceCode":"\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))\n\nrem_p = standard_naryop([_int | _float, _int | _float], 'rem')\nad.defjvp(","sourceCodeStart":5138,"sourceCodeEnd":5174,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5138-L5174","documentation":"Raised by the unreduced-sharding rule for lax.div: in JAX's named-sharding 'unreduced' mechanism, the numerator of a division may be unreduced along some axes, but the denominator must be reduced along exactly those same axes. This error fires when the numerator's unreduced axes don't match the denominator's reduced axes, which would make the division semantically ill-defined (partial sums in the denominator but not the numerator).","triggerScenarios":"Calling jax.lax.div (or '/' on arrays) where operands carry NamedSharding specs with unreduced axes, e.g. after pjit/jit with a sharding where x is unreduced on axis a but y is not reduced on axis a; also passing an unreduced denominator directly (a sibling error).","commonSituations":"Migrating from pmap/pjit collectives to automatic partially-reduced sharding; specifying out_specs / in_specs with unreduced markers inconsistently between operands of a division; sum-then-divide patterns where only one side went through a reduced collective.","solutions":["Ensure the denominator is fully reduced (e.g. apply all_gather/reduce on y) before dividing","Make the numerator and denominator sharding specs consistent: reduce y along exactly the axes x is unreduced on","Do the division after an explicit lax.psum on the denominator","Check the sharding specs of both operands with jax.debug.inspect_array_sharding to spot the mismatch"],"exampleFix":"// before\nz = jnp.sum(x_unreduced, axis=0) / y  # y not reduced on same axes\n\n// after\ny_red = lax.psum(y, axis=0) if is_unreduced(y) else y\nz = jnp.sum(x_unreduced, axis=0) / y_red","handlingStrategy":"validation","validationCode":"def safe_div(x, y):\n    xu, yr = get_unreduced(x), get_reduced(y)\n    if yr:\n        raise ValueError('unreduced denominator')\n    if xu and xu != yr:\n        y = lax.psum(y, axis=tuple(sorted(xu - yr))) if (xu - yr) else y\n    return lax.div(x, y)","typeGuard":"def is_div_sharding_ok(x, y) -> bool:\n    return not getu(y) and (not getu(x) or getu(x) == getr(y))","tryCatchPattern":null,"preventionTips":["Keep denominators fully reduced before division in sharded code","Apply psum/all_gather to operands before arithmetic when specs are uncertain","Inspect operand sharding specs with jax.debug.inspect_array_sharding before dividing"],"tags":["jax","sharding","lax","division","distributed"],"backgroundTag":"sharding-spec-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}