ITADN

[BUG] JVPs of power, divmod, and slice_update mishandle partly traced inputs

#3634Openqflen 创建于 2026-06-06
Q
qflencommented
**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 条评论