Fix CbrtOp derivative NaN gradient on negative inputs (#2571)

[kimjune01]

Fixes #2571 (maintainer-approved: "PR also welcome here!").

Problem

The CbrtOp AD rule (HLODerivatives.td) computes the gradient as (1/3) * pow(x, -2/3). A non-integer exponent on a negative base is NaN (StableHLO power spec: power([-36.0], [1.1]) = -nan), but stablehlo.cbrt is real and smooth for x < 0 (cbrt(-8) = -2). So for finite x < 0 the primal is correct but the gradient is NaN, where the true derivative 1/(3*cbrt(x)^2) is a finite real.

Fix

Compute the derivative from the already-real primal, dropping the power op entirely:

d/dx = DiffeRet / (3 * cbrt(x)^2)

This mirrors jax.lax.cbrt's JVP (g * (1/3) * integer_pow(ans, -2)). Because the failure came from routing through power on a negative base, removing power means the fix cannot reproduce it. Verified against jax.grad(jnp.cbrt) across the sign range (float64):

x	jax.grad	old pow(x,-2/3)/3	new 1/(3*cbrt^2)
-8	0.083333	NaN	0.083333
-2	0.209987	NaN	0.209987
-0.5	0.529134	NaN	0.529134
2	0.209987	0.209987	0.209987

At x = 0 the derivative is a genuine singularity (+/-inf), the same boundary SqrtOp already special-cases.

Test

test/lit_tests/diffrules/stablehlo/cbrt.mlir's @main numerically checks the gradient under stablehlo-translate --interpret, but only sampled positive inputs. Flipped to [-8.0, -27.0]: true gradients are sign-independent so %expected is unchanged, but the old rule produced NaN there. The FORWARD/REVERSE FileCheck blocks are updated to the new (power-free) IR.

Note: I cannot build Enzyme-JAX locally, so the FORWARD/REVERSE CHECK lines are a best-effort prediction of the post-optimization IR. If CI's FileCheck flags an operand-order or constant-format mismatch, I will push a fixup to match. The source fix and the numeric @main check are the substantive parts.

Found mechanically with a value-and-gradient soundness gate over the rewrite/AD-rule library; happy to share the process if useful.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 31.01%. Comparing base (c3406f6) to head (e43d876).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2575      +/-   ##
==========================================
+ Coverage   28.62%   31.01%   +2.38%     
==========================================
  Files         225      225              
  Lines       44316    44438     +122     
==========================================
+ Hits        12686    13782    +1096     
+ Misses      31630    30656     -974     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kimjune01]

@wsmoses only red here is the A100 job, and it died with a Malt.TerminatedWorkerException in nn/convolution and nn/scatter_gather (looks like a cuDNN/OOM crash, unrelated to this cbrt gradient change). CPU and TPU are both green. mind kicking off a re-run when you get a sec?