Gradient of pt.minimum is wrong in numba mode due to fast_math assoc flag

[jessegrabowski]

MWE:

  import numpy as np
  import pytensor
  import pytensor.tensor as pt
  from pytensor.compile.mode import Mode

  x, y, z, w, b = pt.vectors("x", "y", "z", "w", "b")
  m = pt.minimum(x * y * z * w, b).sum()
  gx = pytensor.grad(m, x)

  xv = np.array([0.9887127230177662])
  yv = np.array([1.000617328480952])
  zv = np.array([0.9995082121912299])
  wv = np.array([404.146799095179])
  bv = np.array([((xv[0] * yv[0]) * zv[0]) * wv[0]])   # exactly the left-assoc bits of x*y*z*w

  cvm = pytensor.function([x, y, z, w, b], gx, mode=Mode(linker="cvm", optimizer="fast_run"))
  nb  = pytensor.function([x, y, z, w, b], gx, mode="NUMBA")
  print(f"cvm: {cvm(xv, yv, zv, wv, bv)}")   # [404.19741326]
  print(f"nb:  {nb(xv, yv, zv, wv, bv)}")    # [0.]

The robot claims that we should be storing the floating-point result of the forward for re-use in the backward to avoid this. Obviously I can just disable fastmath, but I want my math to be fast :(

[ricardoV94]

I wouldn't say it's wrong. If we consider this a bug we would have to go back and really reduce the scope of a lot of fusion/algebra/reassociative rewrites we do. The gradient is discontinuous and numerical precision can creep up at the discontinuity.

[jessegrabowski]

it's right in every backend except numba

[ricardoV94]

Well here is a similar example with cvm:

import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import Mode

a, c, b = pt.vectors("a", "c", "b")
m = pt.minimum(pt.exp(a) * pt.exp(c), b).sum()
ga = pytensor.grad(m, a)

av = np.array([-2.486104997138254])
cv = np.array([-1.5791369604234018])
bv = np.exp(av) * np.exp(cv)

cvm      = pytensor.function([a, c, b], ga, mode=Mode(linker="cvm", optimizer="fast_run"))
cvm_excl = pytensor.function([a, c, b], ga, mode=Mode(linker="cvm", optimizer="fast_run").excluding("local_mul_exp_to_exp_add"))
print(cvm(av, cv, bv))       # [0.]          <- exp(a)*exp(c) rewritten to exp(a+c)
print(cvm_excl(av, cv, bv))  # [0.01715884]  <- product left intact

(Even excluding the rewrite it fails in numba and jax backends, but hopefully the point come across)

[ricardoV94]

The robot claims that we should be storing the floating-point result of the forward for re-use in the backward to avoid this. Obviously I can just disable fastmath, but I want my math to be fast :(

We aren't computing it twice it's a low level LLVM thing, the graph is just:

Composite{((1.0 - eq(minimum(i3, (i4 * i2 * i1 * i0)), i3)) * i0 * i1 * i2)} [id A] 0
 ├─ w [id B]
 ├─ z [id C]
 ├─ y [id D]
 ├─ b [id E]
 └─ x [id F]

Inner graphs:

Composite{((1.0 - eq(minimum(i3, (i4 * i2 * i1 * i0)), i3)) * i0 * i1 * i2)} [id A]
 ← mul [id G]
    ├─ sub [id H]
    │  ├─ 1.0 [id I]
    │  └─ EQ [id J]
    │     ├─ Minimum [id K]
    │     │  ├─ i3 [id L]
    │     │  └─ mul [id M]
    │     │     ├─ i4 [id N]
    │     │     ├─ i2 [id O]
    │     │     ├─ i1 [id P]
    │     │     └─ i0 [id Q]
    │     └─ i3 [id L]
    ├─ i0 [id Q]
    ├─ i1 [id P]
    └─ i2 [id O]

[jessegrabowski]

That one is a bit more artificial though, because it depends on inducing a different computation entirely. exp is a pretty complex function

[ricardoV94]

Sorry:

Composite{(eq(minimum((i4 * i2 * i1 * i0), i3), (i4 * i2 * i1 * i0)) * i0 * i1 * i2)} [id A] 0
 ├─ w [id B]
 ├─ z [id C]
 ├─ y [id D]
 ├─ b [id E]
 └─ x [id F]

Inner graphs:

Composite{(eq(minimum((i4 * i2 * i1 * i0), i3), (i4 * i2 * i1 * i0)) * i0 * i1 * i2)} [id A]
 ← mul [id G]
    ├─ EQ [id H]
    │  ├─ Minimum [id I]
    │  │  ├─ mul [id J]
    │  │  │  ├─ i4 [id K]
    │  │  │  ├─ i2 [id L]
    │  │  │  ├─ i1 [id M]
    │  │  │  └─ i0 [id N]
    │  │  └─ i3 [id O]
    │  └─ mul [id J]
    │     └─ ···
    ├─ i0 [id N]
    ├─ i1 [id M]
    └─ i2 [id L]

I was testing min(b, ...) when I plotted the output above... which again is just the sort of stuff we want to do in #2140 and would "break" this gradient in every Op

[ricardoV94]

That one is a bit more artificial though, because it depends on inducing a different computation entirely. exp is a pretty complex function

My point is that we do many graph operations that are invalid under float point precision rules, and you will have a hard time convincing me this is special to do something about (and what would that be? pay the no fastmath cost on every CAReduce?)

[jessegrabowski]

It seems like it's just EQ though

[ricardoV94]

@jessegrabowski bot summary, which makes sense to me. Numba is not "losing precision" or comparing two different values... at least in a naive sense. what happened is that it reassociated the graph, so your manually built b_value is different than the one numba computes on that branch:

1. numba__fastmath = False fixes it

  fastmath default (True): [0.]
  fastmath = False:        [404.19741326]   # matches cvm

  2. The "just EQ" graph is not computed twice — LLVM confirms it

  The composite lowers to a single product SSA value %.11 that feeds both the minimum and the eq:
  %.11 = fmul fast double %.10, %arg.w     ; the ONE product value
  %.13 = select i1 %.12, %arg.b, %.11      ; minimum(b, prod)
  %.14 = fcmp fast oeq double %.13, %.11    ; eq(minimum, prod)  <- reuses %.11
  So the robot's "store the forward result for the backward" suggestion is moot — it already is one value. And the eq is bit-exact (fcmp oeq) with or without fast-math.
  The EQ is not what differs.
  
  3. What actually differs: LLVM reassociates the product fmul chain

  With fast flags LLVM's Reassociate pass rebuilds x*y*z*w into a different association than NumPy used to build b (which Jesse constructed as the exact left-assoc
  product). That shifts the product by 1 ULP — 399.6351251002443 vs 399.63512510024435 — and at the exact tie that flips which minimum branch is taken, zeroing the
  gradient. Compare the two builds:
  ; fastmath=True            ; fastmath=False
  fmul fast z, y             fmul z, w
  fmul fast %, x             fmul y, %
  fmul fast %, w  → ((z·y)·x)·w     fmul x, %  → x·(y·(z·w))

  4. Minimal trigger: {reassoc, nsz} together — no single flag

  Tested on the real composite by injecting flag subsets into PyTensor's opinionated set {arcp, contract, afn, reassoc, nsz}:

  ┌──────────────────────────────────────┬─────────┐
  │                flags                 │  grad   │
  ├──────────────────────────────────────┼─────────┤
  │ False                                │ 404.197 │
  ├──────────────────────────────────────┼─────────┤
  │ any single flag                      │ 404.197 │
  ├──────────────────────────────────────┼─────────┤
  │ full set minus reassoc               │ 404.197 │
  ├──────────────────────────────────────┼─────────┤
  │ full set minus nsz                   │ 404.197 │
  ├──────────────────────────────────────┼─────────┤
  │ full set minus arcp / contract / afn │ 0.      │
  ├──────────────────────────────────────┼─────────┤
  │ {reassoc, nsz}                       │ 0.      │
  └──────────────────────────────────────┴─────────┘

  So it takes exactly reassoc + nsz acting together: reassoc lets LLVM reorder the multiply, and nsz removes the signed-zero constraint that otherwise blocks the
  Reassociate pass. arcp, contract, and afn are irrelevant here.

  Bottom line for the thread: it's not a recompute and it's not the EQ — it's LLVM legally reassociating the product feeding the EQ, which is exactly the class of
  float-reassociating transform PyTensor itself does (your local_mul_exp_to_exp_add example). Confirms it's not a numba-specific bug; the only numba-specific knob is
  that {reassoc, nsz} ship on by default.

[ricardoV94]

It seems like it's just EQ though

Was that because I brought CAReduce? I read minimum -> min. wrong. In this case we would need to disable it in Composite / Elemwise then?