pymc-devs / pymc-devs/pytensor

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

Open
#2,187 10 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

gradients numba request discussion
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.