Fix #2199.

CHANGELOG.md

@@ -29,6 +29,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 * Inconsistent handling of types in lambda lifting (#2197).
 
+* Invalid primal results from `vjp2` in interpreter (#2199).
+
 ## [0.25.24]
 
 ### Added

src/Language/Futhark/Interpreter.hs

@@ -2010,7 +2010,7 @@ initialCtx =
         let drvs = M.map (Just . putAD) $ M.unionsWith add $ map snd m
 
         -- Extract the output values, and the partial derivatives
-        let ov = modifyValue (\i _ -> fst $ m !! i) o
+        let ov = modifyValue (\i _ -> fst $ m !! (length m - 1 - i)) o
         let od =
               fromMaybe (error "vjp: differentiation failed") $
                 modifyValueM (\i vo -> M.findWithDefault (ValuePrim . putV . P.blankPrimValue . P.primValueType . AD.primitive <$> getAD vo) i drvs) v

tests/ad/issue2199.fut

@@ -0,0 +1,14 @@
+-- ==
+-- entry: test_primal test_rev
+-- input { [1.0,2.0] [3.0,4.0] [5.0, 6.0] }
+-- output { 3.0 7.0 11.0 }
+
+def op (x0, y0, z0) (x1, y1, z1) : (f64, f64, f64) = (x0 + x1, y0 + y1, z0 + z1)
+def ne = (0f64, 0f64, 0f64)
+
+def primal xs = reduce_comm op ne xs
+
+entry test_primal as bs cs = primal (zip3 as bs cs)
+
+entry test_rev as bs cs =
+  (vjp2 (\(as, bs, cs) -> test_primal as bs cs) (as, bs, cs) (1, 1, 1)).0