pymc-devs / pymc-devs/pytensor
Gradient of pt.minimum is wrong in numba mode due to fast_math assoc flag
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
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 :(
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by running the supplied MWE and compare the cvm and NUMBA results for the gradient of pt.minimum under fast_run. Then inspect the NUMBA mode and fast_math handling around the forward and backward computations. Done means the NUMBA result agrees with the cvm result for this case while retaining fast-math performance, with a regression test covering the mismatch.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100