[BUG] JVPs of power, divmod, and slice_update mishandle partly traced inputs
**Describe the bug**
Three more `Primitive::jvp` implementations mishandle partly traced inputs — the same class as #3627/#3629 (fixed by #3633), found while auditing the remaining jvp implementations against the packed-tangents convention (`tangents[i]` is the tangent of input `argnums[i]`; a jvp returns one tangent per output):
1. `Power::jvp` delegates to `vjp`, which scales every partial by `cotangents[0]`. With both inputs traced, both partials are multiplied by the first input's tangent and the second tangent is ignored - silently wrong tangents.
2. `DivMod::jvp` returns one tangent for a two-output primitive, so the second output never enters the tangent map and its consumer is invoked with empty `argnums`/`tangents` - segfault in release builds.
3. `DynamicSliceUpdate::jvp` ignores `argnums` and reads `tangents[0]`/`tangents[1]` unconditionally - segfault when only one input is traced, and mispaired tangents for `argnums = {0, 2}`.
**To Reproduce**
```python
import mlx.core as mx
# 1) power: silently wrong with both inputs traced
_, (j,) = mx.jvp(lambda a, b: a**b, [mx.array(2.0), mx.array(3.0)],
[mx.array(0.0), mx.array(1.0)])
print(j) # 0.0 — expected d/db 2^b = ln(2)*8 ~ 5.5452
_, (j,) = mx.jvp(lambda a, b: a**b, [mx.array(2.0), mx.array(3.0)],
[mx.array(1.0), mx.array(0.0)])
print(j) # 17.5452 — expected d/da a^3 = 12.0 (got 12 + 5.5452)
# 2) divmod: segfault when the second output is consumed
mx.jvp(lambda x: -(mx.divmod(x, mx.array(3.0))[1]), [mx.array(7.0)], [mx.array(1.0)])
# 3) slice_update with array start indices, only the update traced: segfault
x = mx.zeros((4,))
mx.jvp(lambda u: mx.slice_update(x, u, mx.array([1]), axes=[0]),
[mx.ones(2)], [mx.ones(2)])
```
**Expected behavior**
1. `power` uses each input's own tangent: 5.5452 and 12.0 respectively.
2. `divmod` returns a tangent per output (zeros, matching its vjp semantics); no crash.
3. `slice_update` treats untraced inputs as zero tangents; no crash.
In a debug build, 2) aborts at `Negative::jvp`'s `assert(argnums.size() == 1)`, confirming the downstream primitive is invoked with no tangents. `mx.vjp`/`mx.grad` are unaffected for all three.
I can have a fix for these too soon.
**Desktop**
- macOS, Apple Silicon
- main @ 8f0e8b14
0 条评论