[BUG] JVP silently returns a zero tangent when differentiating through mx.where
bug
**Describe the bug**
JVP silently returns a zero tangent when differentiating through mx.where(cond, constant, traced_value) specifically when the constant is in the TRUE branch and the traced value is in the FALSE branch. The bug fires regardless of whether cond is ever True at runtime.
**To Reproduce**
Include code snippet
```
"""Minimum reproducer: MLX `mx.jvp` silently returns 0 through
`mx.where(cond, <constant>, <traced_value>)`.
# TLDR
`mx.jvp(f, ...)` returns a *zero tangent* when `f` contains
`mx.where(cond, constant, traced_value)` — i.e. when the constant is in
the **TRUE branch** and the traced (tangent-bearing) value is in the
**FALSE branch**. `mx.vjp` and central finite differences both return
the correct gradient.
The bug fires regardless of:
- whether `cond` is ever True at runtime (it is False everywhere here),
- what the constant is (NaN, ±inf, 0.0, 999.0, mx.array(0.0), ...),
- whether the `traced_value` was produced by simple arithmetic.
Swapping the branches — putting the traced value in the TRUE branch and
the constant in the FALSE branch — gives the correct JVP. So does using
two traced values in both branches.
# Expected output (MLX 0.31.2, Apple Silicon)
TRUE=traced FALSE=const : VJP=+2.000e+00 JVP=+2.000e+00 ratio=1.00e+00
TRUE=const FALSE=traced : VJP=+2.000e+00 JVP=+0.000e+00 ratio=2.00e+30 <-- BUG
TRUE=traced FALSE=traced : VJP=+2.000e+00 JVP=+2.000e+00 ratio=1.00e+00
FD ground truth : +1.999736e+00
# Fix for now
Swap the branches: replace `mx.where(cond, const, traced)` with
`mx.where(mx.logical_not(cond), traced, const)`.
# Tested
MLX 0.31.2, Apple Silicon .
"""
import mlx.core as mx
def _bench(f, label):
theta = mx.array(0.0, dtype=mx.float32)
v = mx.array(1.0, dtype=mx.float32)
cot = mx.array(1.0, dtype=mx.float32)
primal = float(f(theta))
_, vj = mx.vjp(f, [theta], [cot]); mx.eval(vj[0]); vjp = float(vj[0])
_, jv = mx.jvp(f, [theta], [v]); mx.eval(jv[0]); jvp = float(jv[0])
ratio = abs(vjp) / max(abs(jvp), 1e-30)
flag = " <-- BUG" if ratio > 1e6 else ""
print(f" {label:25s}: primal={primal:+.4f} VJP={vjp:+.4e} "
f"JVP={jvp:+.4e} ratio={ratio:.2e}{flag}")
return ratio > 1e6
# Variant A: traced value in TRUE branch, constant in FALSE branch — works
def fA(theta):
y = (theta + 1.0) ** 2
return mx.where(y > -1.0, y, 999.0) # cond always True
# Variant B: constant in TRUE branch, traced value in FALSE branch — BUG
def fB(theta):
y = (theta + 1.0) ** 2
return mx.where(y < -1.0, 999.0, y) # cond always False
# Variant C: both branches traced — works
def fC(theta):
y = (theta + 1.0) ** 2
z = (theta + 2.0)
return mx.where(y > -1.0, y, z)
print("=" * 78)
print("MLX `mx.jvp` silent-zero bug through `mx.where(cond, const, traced)`")
print("=" * 78)
_bench(fA, "TRUE=traced FALSE=const")
buggy = _bench(fB, "TRUE=const FALSE=traced")
_bench(fC, "TRUE=traced FALSE=traced")
# Finite-difference ground truth for fB (same primal as fA/fC)
theta = mx.array(0.0, dtype=mx.float32); h = 1e-4
rp = float(fB(theta + mx.array(h, dtype=mx.float32)))
rm = float(fB(theta - mx.array(h, dtype=mx.float32)))
print(f"\n FD ground truth on fB : {(rp - rm) / (2.0 * h):+.6e}")
print()
if buggy:
print("==> BUG CONFIRMED: `mx.jvp` through `mx.where(cond, const, traced)` "
"returns a zero tangent, even when `cond` is always False at "
"runtime. `mx.vjp` and finite-differences both return the "
"correct gradient (+2.0).")
print()
print(" Workaround in user code: swap the branches —")
print(" mx.where(mx.logical_not(cond), traced, const)")
else:
print("==> Bug not triggered in this run.")
```
**Expected behavior**
I expect both the VJP and JVP answers to match. This does not happen here.
**Desktop (please complete the following information):**
- Version MLX 0.31.2
0 条评论