diff --git a/src/enzyme_ad/jax/Utils.cpp b/src/enzyme_ad/jax/Utils.cpp index 402c8fa9db..ef630bb906 100644 --- a/src/enzyme_ad/jax/Utils.cpp +++ b/src/enzyme_ad/jax/Utils.cpp @@ -987,10 +987,10 @@ NonNegativeResultAnalysis::State NonNegativeResultAnalysis::localGuaranteed( // (mul a a) is always non-negative if (auto mulOp = dyn_cast(op)) { - auto lhsOp = mulOp.getLhs().getDefiningOp(); - auto rhsOp = mulOp.getRhs().getDefiningOp(); - - if (lhsOp == rhsOp) { + // Compare the operand Values, not their defining ops: two distinct block + // arguments both have a null defining op, so a defining-op comparison + // would wrongly prove mul(%arg0, %arg1) non-negative. + if (mulOp.getLhs() == mulOp.getRhs()) { return State::GUARANTEED; } } diff --git a/test/lit_tests/abspositive.mlir b/test/lit_tests/abspositive.mlir index 079597880d..f2afaad02f 100644 --- a/test/lit_tests/abspositive.mlir +++ b/test/lit_tests/abspositive.mlir @@ -95,3 +95,12 @@ func.func @test8(%arg0: tensor<4xf64>) -> tensor<4xf64> { // CHECK-NEXT: %0 = stablehlo.logistic %arg0 {enzymexla.non_negative = [#enzymexla]} : tensor<4xf64> // CHECK-NEXT: return %0 : tensor<4xf64> // CHECK-NEXT: } + +// mul of two distinct block args must NOT be proven non-negative: the product +// is negative whenever the args have opposite signs, so the abs has to stay. +func.func @test9(%arg0: tensor<12xf64>, %arg1: tensor<12xf64>) -> tensor<12xf64> { + %0 = stablehlo.multiply %arg0, %arg1 : tensor<12xf64> + // CHECK: stablehlo.abs + %1 = stablehlo.abs %0 : tensor<12xf64> + return %1 : tensor<12xf64> +}