log(a*a) → 2*log(a) returns NaN for negative inputs

[kimjune01]

Human section

I'm working with some dev tooling for generalizing tests, and I found these cases
involving LogSimplify that I believe will be a strict correctness improvement. If the
following analysis is valuable, I'd be happy to share the process I used to find the
bugs mechanically.

Machine section

Component: EnzymeHLOOptPass (default --enzyme-hlo-opt)
File: src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp, struct LogSimplify (L27376)
Class: value-soundness. Optimized program returns NaN where unoptimized returns a finite real.

Root cause: four sub-cases apply a real-analysis identity valid only for positive bases,
with no domain guard. CheckedOpRewritePattern checks disable-attr + static-shape only
(CheckedRewrite.h). The original is real for the full input domain; the rewritten form is
real only for positive operands. Domain narrowed.

#	line	rewrite	unsound when	example (f64) before to after
1	27406	log(a*a) -> 2*log(a)	any a<0	a=-2: log(4)=1.386 to 2*log(-2)=NaN
2	27391	log(pow(x,y)) -> y*log(x)	x<0 (even y), or y=0	y=0,x=-2: log(pow(-2,0))=log(1)=0 to 0*log(-2)=NaN; also y=2: log(4)=1.386 to NaN
3	27416	log(a*b) -> log(a)+log(b) (one const)	const <0	b=-3,a=-2: log(6)=1.792 to NaN
4	27469	log(a/b) -> log(a)-log(b) (one const)	const <0	b=-3,a=-2: log(2/3)=-0.405 to NaN

Note #1/#2 trigger on a runtime negative alone (no constant needed); a*a is always >= 0 so the
input log(a*a) is real for every a != 0, but 2*log(a) is NaN for all a<0. #3/#4 need the
constant operand negative (positive constant is sound; both sides NaN for x<0).

Why this is not a precision/fast-math question. It is tempting to file these under the same
opt-in float policy as log(sqrt(x)) -> 0.5*log(x), but they are a different class. That rewrite
is an exact identity in the reals and only perturbs float rounding/quantization, which is what the
opt-in flag is for. These four are false in the reals themselves: as partial functions the two
sides have different domains (log(a*a) is defined for a != 0, 2*log(a) only for a > 0), so
the divergence appears before floating point is involved and no precision flag should sanction it.
That also gives a clean guard rule: a sub-case is sound exactly when the rewrite is a real-field
identity on the operand's sign, which for #3/#4 is precisely "the constant operand is non-negative"
(positive const: both sides agree, NaN-for-NaN on negatives; negative const: the rewrite
manufactures a NaN). This is the same finite-to-NaN domain-narrowing as the CbrtOp derivative in
#2571; both are the optimized form undefined where the original is real.

Repro (case #1):

func.func @repro(%a: tensor<f64>) -> tensor<f64> {
  %sq = stablehlo.multiply %a, %a : tensor<f64>
  %r  = stablehlo.log %sq : tensor<f64>
  return %r : tensor<f64>
}
// enzymexlamlir-opt --enzyme-hlo-opt  ==>  2.0 * stablehlo.log(%a)
// %a = -2.0 :  unopt = log(4) = 1.3863   opt = 2*log(-2) = NaN

(The rewrite firing is already asserted by test/lit_tests/logsimplify.mlir @main2; that test
checks IR text only and never evaluates, which is why the value divergence is invisible to it.)

Fix: gate each sub-case on a provably-non-negative base (reuse the non-negativity analysis
already used by AbsPositiveSimplify / guaranteedNonNegativeResult). Per case:

• Use workspace clone #1 log(a*a) -> 2*log(abs(a)) is sound unconditionally (one extra abs, a*a and abs(a)
  share a domain), so it needs no guard.
• Update README.md #2 fire only when x is provably non-negative (or pow's result is).
• Patch 1 #3/pip does not automatically import jax #4 fire only when the constant operand is non-negative.

The abs form for #1 fixes the domain (finite-to-NaN) issue but is not bit-exact under the
multiply's overflow/underflow/rounding, so it still belongs under the float-rewrite policy as a
precision matter, distinct from the domain bug above.

Scope: These four are the only rewrites in this file that turn a finite real into NaN (a
domain-guard bug). Differences that are purely overflow/underflow/rounding are a separate
fast-math question and I set them aside, including the same log rewrites at extreme magnitudes and
log(a+a) -> log(2)+log(a). At the domain level log(sqrt)/log(cbrt)/log(rsqrt),
SquareAbsSimplify, SignAbsSimplify, PowSimplify, DivideSqrtToMultiplyRsqrt, multiply/power
(ChainedMultiplyToPower; pow(-2,3) returns -8, the unsoundness hypothesis did not reproduce),
and the NoNan* patterns are sound.

Acceptance: after the fix, for a<0 the optimized and unoptimized programs agree (both compute
log(a*a) correctly, i.e. a finite real).

[wsmoses]

makes sense, PR welcome!

[kimjune01]

Opened #2576 with the fix (four domain guards, reusing the existing guaranteedNonNegativeResult analysis).

Since it's from a fork, the CI workflows are waiting on a maintainer to approve the run. Could you approve it when you get a chance? I can't build Enzyme-JAX locally, so the updated CHECK blocks in logsimplify.mlir are a best-effort prediction of the post-pass IR (operand order, SSA numbering, the inserted stablehlo.abs); the source guards are the substantive change. If CI flags a CHECK mismatch I'll push a fixup.

One open question is in the PR: I rewrote log(a*a) to 2*log(abs(a)) (unconditionally domain-sound) rather than gating it like the others. Happy to gate it instead if you'd prefer to keep it under the opt-in float policy.