futhark 0.25.28 → 0.25.29

raw patch · 15 files changed

+251/−96 lines, 15 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• CHANGELOG.md +23/−4
• futhark.cabal +1/−1
• src/Futhark/AD/Derivatives.hs +1/−1
• src/Futhark/AD/Rev.hs +53/−8
• src/Futhark/IR/Prop/Rearrange.hs +6/−0
• src/Futhark/Internalise/Defunctionalise.hs +7/−3
• src/Futhark/Internalise/Entry.hs +16/−8
• src/Futhark/Internalise/Monomorphise.hs +8/−13
• src/Futhark/Optimise/Fusion.hs +4/−3
• src/Futhark/Optimise/Fusion/RulesWithAccs.hs +4/−5
• src/Futhark/Optimise/Simplify/Rules/Index.hs +67/−6
• src/Language/Futhark/Interpreter.hs +8/−6
• src/Language/Futhark/Interpreter/AD.hs +48/−34
• src/Language/Futhark/Interpreter/Values.hs +2/−1
• src/Language/Futhark/TypeChecker/Modules.hs +3/−3

@@ -5,13 +5,32 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). -## [0.25.28]+## [0.25.29] -### Added+### Fixed -### Removed+* AD would in some cases produce code that would cause a compiler+  crash (#2228). -### Changed+* Slight error in the definition of the partial derivatives for the+  `**` operator could cause NaNs in the interpreter when using+  forward-mode AD (#2229).++* The magical machinery for inferring external API types did not+  handle arrays with uniqueness annotations consistently, resulting in+  incompatible entry point types being generated, leading to a+  compiler crash. (#2231)++* A simplification rule for array slices would in some cases produce+  type-incorrect code. (#2232)++* A bug in the defunctionaliser could cause a compiler crash in code+  that used both higher order functions and size expressions in clever+  ways (#2234).++* Fusion could crash after AD in some circumstances (#2236).++## [0.25.28]  ### Fixed 

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.28+version:        0.25.29 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to

@@ -136,7 +136,7 @@   where     derivs x y =       ( y * (x ** (y - 1)),-        condExp (x .<=. 0) 0 ((x ** y) * log x)+        (x ** y) * condExp (x .<=. 0) 0 (log x)       ) pdBinOp (FMax ft) a b =   floatBinOp derivs derivs derivs ft a b

@@ -28,6 +28,15 @@ patName (Pat [pe]) = pure $ patElemName pe patName pat = error $ "Expected single-element pattern: " ++ prettyString pat +copyIfArray :: VName -> ADM VName+copyIfArray v = do+  v_t <- lookupType v+  case v_t of+    Array {} ->+      letExp (baseString v <> "_copy") . BasicOp $+        Replicate mempty (Var v)+    _ -> pure v+ -- The vast majority of BasicOps require no special treatment in the -- forward pass and produce one value (and hence one adjoint).  We -- deal with that case here.@@ -202,13 +211,7 @@       (_pat_v, pat_adj) <- commonBasicOp pat aux e m       returnSweepCode $ do         v_adj <- letExp "update_val_adj" $ BasicOp $ Index pat_adj slice-        t <- lookupType v_adj-        v_adj_copy <--          case t of-            Array {} ->-              letExp "update_val_adj_copy" . BasicOp $-                Replicate mempty (Var v_adj)-            _ -> pure v_adj+        v_adj_copy <- copyIfArray v_adj         updateSubExpAdj v v_adj_copy         zeroes <- letSubExp "update_zero" . zeroExp =<< subExpType v         void $@@ -278,7 +281,9 @@           ses           (map (fmap $ diffBody adjs branches_free) cases)           (diffBody adjs branches_free defbody)-    zipWithM_ insAdj branches_free branches_free_adj+    -- See Note [Array Adjoints of Match]+    forM_ (zip branches_free branches_free_adj) $ \(v, v_adj) ->+      insAdj v =<< copyIfArray v_adj diffStm (Let pat aux (Op soac)) m =   vjpSOAC vjpOps pat aux soac m diffStm (Let pat aux loop@Loop {}) m =@@ -439,3 +444,43 @@ -- our current translation rules, they will be dead code.  As long as -- we are careful to run dead code elimination after revVJP, we should -- be good.++-- Note [Array Adjoints of Match]+--+-- Some unusual, but sadly not completely contrived, contain Match+-- expressions that return multiple arrays, and there the arrays+-- returned by one branch have overlapping aliases with another+-- branch, although in different places. As an example consider this:+--+--   let (X,Y) = if c+--               then (A, B)+--               else (B, A)+--+-- Because our aliasing representation cannot express mutually+-- exclusive aliases, we will consider X and Y to be aliased to each+-- other. In practice, this means it is unlikely for X or Y to be+-- consumed, because it would also consume the other (although it's+-- possible for carefully written code).+--+-- When producing adjoints for this, it will be something like+--+--   let (X_adj,Y_adj) = if c+--                       then (A_adj, B_adj)+--                       else (B_adj, A_adj)+--+-- which completely reflects the primal code. However, while it is+-- unlikely that any consumption takes place for the original primal+-- variables, it is almost guaranteed that X_adj and Y_adj will be+-- consumed (that is the main way we use adjoints after all), and due+-- to the conservative aliasing, when one is consumed, so is the+-- other! To avoid this tragic fate, we are forced to copy any+-- array-typed adjoints returned by a Match. This can be quite costly.+-- However:+--+-- 1) Futhark has pretty OK copy removal, so maybe it can get rid of+--    these by using information not available to the AD pass.+--+-- 2) In many cases, arrays will have accumulator adjoints, which are+--    not subject to this problem.+--+-- Issue #2228 was caused by neglecting to do this.

@@ -6,6 +6,7 @@     rearrangeReach,     rearrangeCompose,     isPermutationOf,+    isIdentityPerm,     transposeIndex,     isMapTranspose,   )@@ -61,6 +62,11 @@       | otherwise = do           (xs', v) <- pick (i + 1) xs y           pure (x : xs', v)++-- | Is this an identify permutation? An identity permutation is of+-- the form @[0, 1, ..., k]@.+isIdentityPerm :: [Int] -> Bool+isIdentityPerm perm = perm == [0 .. length perm - 1]  -- | If @l@ is an index into the array @a@, then @transposeIndex k n -- l@ is an index to the same element in the array @transposeArray k n

@@ -310,8 +310,11 @@     f _ = Nothing  sizesToRename :: StaticVal -> S.Set VName-sizesToRename (DynamicFun (_, sv1) sv2) =-  sizesToRename sv1 <> sizesToRename sv2+sizesToRename (DynamicFun (_, sv1) _sv2) =+  -- It is intentional that we do not look at sv2 here, as some names+  -- that are free in sv2 are actually bound by the parameters in sv1.+  -- See #2234.+  sizesToRename sv1 sizesToRename IntrinsicSV =   mempty sizesToRename HoleSV {} =@@ -882,8 +885,9 @@           globals <- asks fst           let bound_sizes = S.fromList dims' <> globals           pats' <- instAnySizes pats+          let dims'' = dims' ++ unboundSizes bound_sizes pats' -          liftValDec fname (RetType [] rettype') (dims' ++ unboundSizes bound_sizes pats') pats' e0+          liftValDec fname (RetType [] rettype') dims'' pats' e0           pure             ( Var                 (qualName fname)

@@ -185,6 +185,12 @@ entryPointTypeName (I.TypeOpaque v) = v entryPointTypeName (I.TypeTransparent {}) = error "entryPointTypeName: TypeTransparent" +elemTypeExp :: E.TypeExp E.Exp VName -> Maybe (E.TypeExp E.Exp VName)+elemTypeExp (E.TEArray _ te _) = Just te+elemTypeExp (E.TEUnique te _) = elemTypeExp te+elemTypeExp (E.TEParens te _) = elemTypeExp te+elemTypeExp _ = Nothing+ entryPointType ::   VisibleTypes ->   E.EntryType ->@@ -221,17 +227,13 @@                   rank = E.shapeRank shape                   ts' = map (strip rank) ts                   record_t = E.Scalar (E.Record fs)-                  record_te = case E.entryAscribed t of-                    Just (E.TEArray _ te _) -> Just te-                    _ -> Nothing+                  record_te = elemTypeExp =<< E.entryAscribed t               ept <- snd <$> entryPointType types (E.EntryType record_t record_te) ts'               addType desc . I.OpaqueRecordArray rank (entryPointTypeName ept)                 =<< opaqueRecordArray types rank fs' ts         E.Array _ shape et -> do           let ts' = map (strip (E.shapeRank shape)) ts-              elem_te = case E.entryAscribed t of-                Just (E.TEArray _ te _) -> Just te-                _ -> Nothing+              elem_te = elemTypeExp =<< E.entryAscribed t           ept <- snd <$> entryPointType types (E.EntryType (E.Scalar et) elem_te) ts'           addType desc . I.OpaqueArray (E.shapeRank shape) (entryPointTypeName ept) $             map valueType ts@@ -265,9 +267,15 @@                          E.entryAscribed eret                        ) of                 (Just ts, Just (E.TETuple e_ts _)) ->-                  zipWithM (entryPointType types) (zipWith E.EntryType ts (map Just e_ts)) crets+                  zipWithM+                    (entryPointType types)+                    (zipWith E.EntryType ts (map Just e_ts))+                    crets                 (Just ts, Nothing) ->-                  zipWithM (entryPointType types) (map (`E.EntryType` Nothing) ts) crets+                  zipWithM+                    (entryPointType types)+                    (map (`E.EntryType` Nothing) ts)+                    crets                 _ ->                   pure <$> entryPointType types eret (concat crets)           )

@@ -841,12 +841,12 @@     -- the thing that should be fixed, but it requires fiddling with     -- the defunctorisation of size-lifted types. -    onExps bound (Var v _ _) e-      | Just rexp <- lookup (qualLeaf v) named1 =-          onExps mempty (unReplaced rexp) e-      | otherwise =-          unless (any (`elem` bound) $ freeVarsInExp e) $-            modify (M.insert (qualLeaf v) e)+    onExps bound (Var v _ _) e = do+      unless (any (`elem` bound) $ freeVarsInExp e) $+        modify (M.insert (qualLeaf v) e)+      case lookup (qualLeaf v) named1 of+        Just rexp -> onExps mempty (unReplaced rexp) e+        Nothing -> pure ()     onExps _bound e (Var v _ _)       | Just rexp <- lookup (qualLeaf v) named2 =           onExps mempty e (unReplaced rexp)@@ -997,7 +997,7 @@         substTypesAny (fmap (fmap (second (const mempty))) . (`M.lookup` substs'))       params' = map (substPat substStructType) params   params'' <- withArgs shape_names $ mapM transformPat params'-  exp_naming <- paramGetClean+  exp_naming <- get <* put mempty    let args = S.fromList $ foldMap patNames params       arg_params = map snd exp_naming@@ -1006,7 +1006,7 @@     withParams exp_naming $       withArgs (args <> shape_names) $         hardTransformRetType (applySubst (`M.lookup` substs') rettype)-  extNaming <- paramGetClean+  extNaming <- get <* put mempty   scope <- S.union shape_names <$> askScope'   let (rettype'', new_params) = arrowArg scope args arg_params rettype'       bind_t' = substTypesAny (`M.lookup` substs') bind_t@@ -1062,11 +1062,6 @@     shape_params = filter (not . isTypeParam) tparams      updateExpTypes substs = astMap (mapper substs)--    paramGetClean = do-      ret <- get-      put mempty-      pure ret      hardTransformRetType (RetType dims ty) = do       ty' <- transformType ty

@@ -387,18 +387,19 @@     | wacc2_cons_nms <- namesFromList $ concatMap (\(_, nms, _) -> nms) w2_inps,       wacc1_indep_nms <- map getName is1,       all (`notNameIn` wacc2_cons_nms) wacc1_indep_nms = do-        -- \^ the other safety checks are done inside `tryFuseWithAccs`+        -- the other safety checks are done inside `tryFuseWithAccs`         lam1' <- fst <$> doFusionInLambda lam1         lam2' <- fst <$> doFusionInLambda lam2         let stm1 = Let pat1 aux1 (WithAcc w1_inps lam1')             stm2 = Let pat2 aux2 (WithAcc w2_inps lam2')-        mstm <- SF.tryFuseWithAccs infusible stm1 stm2+        mstm <- sequence $ SF.tryFuseWithAccs infusible stm1 stm2         case mstm of+          Nothing -> pure Nothing           Just (Let pat aux (WithAcc w_inps wlam)) -> do             (wlam', success) <- doFusionInLambda wlam             let new_stm = Let pat aux (WithAcc w_inps wlam')             if success then fusedSomething (StmNode new_stm) else pure Nothing-          _ -> error "Illegal result of tryFuseWithAccs called from vFuseNodeT."+          Just _ -> error "Illegal result of tryFuseWithAccs called from vFuseNodeT." -- vFuseNodeT _ _ _ _ = pure Nothing 

@@ -400,7 +400,7 @@   [VName] ->   Stm SOACS ->   Stm SOACS ->-  m (Maybe (Stm SOACS))+  Maybe (m (Stm SOACS)) tryFuseWithAccs   infusible   (Let pat1 aux1 (WithAcc w_inps1 lam1))@@ -424,7 +424,7 @@       all (`notElem` infusible) bs,       -- safety 3:       cs <- namesFromList $ concatMap ((\(_, xs, _) -> xs) . accTup2) acc_tup2,-      all ((`notNameIn` cs) . patElemName . fst) other_pr1 = do+      all ((`notNameIn` cs) . patElemName . fst) other_pr1 = Just $ do         let getCertPairs (t1, t2) = (paramName (accTup3 t2), paramName (accTup3 t1))             tab_certs = M.fromList $ map getCertPairs tup_common             lam2_bdy' = substituteNames tab_certs (lambdaBody lam2)@@ -472,8 +472,7 @@                 ++ map fst (other_pr1 ++ other_pr2)             res_w_inps = map (accTup2 . fst) tup_common ++ map accTup2 (acc_tup1' ++ acc_tup2')         res_w_inps' <- mapM renameLamInWAccInp res_w_inps-        let stm_res = Let (Pat res_pat) (aux1 <> aux2) $ WithAcc res_w_inps' res_lam'-        pure $ Just stm_res+        pure $ Let (Pat res_pat) (aux1 <> aux2) $ WithAcc res_w_inps' res_lam'     where       -- local helpers: @@ -535,7 +534,7 @@       renameLamInWAccInp winp = pure winp -- tryFuseWithAccs _ _ _ =-  pure Nothing+  Nothing  ------------------------------- --- simple helper functions ---

@@ -8,6 +8,8 @@ where  import Control.Monad (guard)+import Data.Bifunctor (first)+import Data.List qualified as L import Data.List.NonEmpty (NonEmpty (..)) import Data.Maybe import Futhark.Analysis.PrimExp.Convert@@ -33,6 +35,8 @@   | SubExpResult Certs SubExp  -- Fake expressions that we can recognise.+--+-- See Note [Simplifying a Slice]. fakeIndices :: [TPrimExp Int64 VName] fakeIndices = map f [0 :: Int ..]   where@@ -61,6 +65,7 @@         worthInlining e,         all (`ST.elem` vtable) (unCerts cs) ->           Just $ SubExpResult cs <$> toSubExp "index_primexp" e+      -- For the two cases below, see Note [Simplifying a Slice].       | Just inds' <- sliceIndices (Slice inds),         Just (ST.IndexedArray cs arr inds'') <- ST.index idd inds' vtable,         all (worthInlining . untyped) inds'',@@ -76,17 +81,28 @@         all (`ST.elem` vtable) (unCerts cs),         not consuming,         not $ consumed arr,-        Just inds''' <- mapM okIdx inds'' -> do-          Just $ IndexResult cs arr . Slice <$> sequence inds'''+        Just (ordering, inds''') <- first concat . unzip <$> mapM okIdx inds'',+        Just perm <- L.sort ordering `isPermutationOf` ordering ->+          if isIdentityPerm perm+            then Just $ IndexResult cs arr . Slice <$> sequence inds'''+            else Just $ do+              arr_sliced <-+                certifying cs $+                  letExp (baseString arr <> "_sliced") . BasicOp . Index arr . Slice+                    =<< sequence inds'''+              arr_sliced_tr <-+                letSubExp (baseString arr_sliced <> "_tr") $+                  BasicOp (Rearrange perm arr_sliced)+              pure $ SubExpResult mempty arr_sliced_tr       where-        matches = zip fakeIndices $ sliceDims $ Slice inds+        matches = zip fakeIndices $ zip [0 :: Int ..] $ sliceDims $ Slice inds         okIdx i =           case lookup i matches of-            Just w ->-              Just $ pure $ DimSlice (constant (0 :: Int64)) w (constant (1 :: Int64))+            Just (j, w) ->+              Just ([j], pure $ DimSlice (constant (0 :: Int64)) w (constant (1 :: Int64)))             Nothing -> do               guard $ not $ any ((`namesIntersect` freeIn i) . freeIn . fst) matches-              Just $ DimFix <$> toSubExp "index_primexp" i+              Just ([], DimFix <$> toSubExp "index_primexp" i)     Nothing -> Nothing     Just (SubExp (Var v), cs) ->       Just $ pure $ IndexResult cs v $ Slice inds@@ -258,3 +274,48 @@           True       | otherwise =           False++-- Note [Simplifying a Slice]+--+-- The 'indexOp' simplification machinery permits simplification of+-- full indexing (i.e., where every component of the Slice is a+-- DimFix). We use this in a creative way to also simplify slices.+-- For example, for a slice+--+--   A[i:j:+n]+--+-- we synthesize a uniquely recognizable index expression "ie" (see+-- 'fakeIndices'), we use `indexOp` to simplify the full indexing+--+--   A[ie]+--+-- and if that produces a simplification result+--+--   B[ie]+--+-- then we can replace the "ie" DimFix with our original slice and+-- produce+--+--   B[i:j:+n]+--+-- While the above case is trivial, this is useful for cases that+-- intermix indexing and slicing. We must be careful, however: If we+-- have an original expression+--+--   A[i0:j0:+n0,i1:j1:+n1]+--+-- for which we then synthesize the expression+--+--   A[ie0, ie1]+--+-- then if we receive back the result+--+--   B[ie1, ie0]+--+-- we cannot just replace the indexes with the original slices, as+-- that would change the shape (and semantics) of the result:+--+--   B[i1:j1:+n1,i0:j0:+n0]+--+-- In such cases we must actually insert a Rearrange operation to move+-- the dimensions of the result appropriately.

@@ -35,6 +35,7 @@ import Data.Array import Data.Bifunctor import Data.Bitraversable+import Data.Either (fromRight) import Data.List   ( find,     foldl',@@ -1351,8 +1352,9 @@      adToPrim v = putV $ AD.primitive v -    adBinOp op x y = AD.doOp op [x, y]-    adUnOp op x = AD.doOp op [x]+    adBinOp op x y =+      either (const Nothing) Just $ AD.doOp op [x, y]+    adUnOp op x = either (const Nothing) Just $ AD.doOp op [x]      fun1 f =       TermValue Nothing $ ValueFun $ \x -> f x@@ -1421,7 +1423,7 @@         _           | Just x' <- getAD x,             Just y' <- getAD y,-            Just z <- msum $ map (`bopDefAD'` (x', y')) fs -> do+            Just z <- msum $ map (`bopDefAD` (x', y')) fs -> do               breakOnNaN [adToPrim x', adToPrim y'] $ adToPrim z               pure $ putAD z         _ ->@@ -1436,7 +1438,7 @@           x' <- valf x           y' <- valf y           retf =<< op x' y'-        bopDefAD' (_, _, _, dop) (x, y) = dop x y+        bopDefAD (_, _, _, dop) (x, y) = dop x y      unopDef fs = fun1 $ \x ->       case x of@@ -1471,7 +1473,7 @@         Just [x, y]           | Just x' <- getAD x,             Just y' <- getAD y,-            Just z <- AD.doOp op [x', y'] -> do+            Right z <- AD.doOp op [x', y'] -> do               breakOnNaN [adToPrim x', adToPrim y'] $ adToPrim z               pure $ putAD z         _ ->@@ -2029,7 +2031,7 @@         -- TODO: Perhaps this could be fully abstracted by AD?         -- Making addFor private would be nice..         add x y =-          fromMaybe (error "jvp: illtyped add") $+          fromRight (error "jvp: illtyped add") $             AD.doOp (AD.OpBin $ AD.addFor $ P.primValueType $ AD.primitive x) [x, y]     def "jvp2" = Just $       -- TODO: This could be much better. Currently, it is very inefficient

@@ -15,7 +15,7 @@ where  import Control.Monad (foldM, zipWithM)-import Data.Either (isRight)+import Data.Either (fromRight, isRight) import Data.List (find, foldl') import Data.Map qualified as M import Data.Maybe (fromMaybe)@@ -70,12 +70,14 @@ mulFor Bool = LogAnd mulFor t = error $ "mulFor: " ++ show t +type Depth = Int+ -- Types and utility functions-- -- When taking the partial derivative of a function, we -- must differentiate between the values which are kept -- constant, and those which are not data ADValue-  = Variable Int ADVariable+  = Variable Depth ADVariable   | Constant PrimValue   deriving (Show) @@ -88,7 +90,7 @@   | JVP JVPValue   deriving (Show) -depth :: ADValue -> Int+depth :: ADValue -> Depth depth (Variable d _) = d depth (Constant _) = 0 @@ -97,6 +99,12 @@ primal (Variable _ (JVP (JVPValue v _))) = primal v primal (Constant v) = Constant v +primalFor :: Depth -> ADValue -> ADValue+primalFor cur v@(Variable tag _) | cur /= tag = v+primalFor _ (Variable _ (VJP (VJPValue t))) = tapePrimal t+primalFor cur (Variable _ (JVP (JVPValue v _))) = primalFor cur v+primalFor _ (Constant v) = Constant v+ primitive :: ADValue -> PrimValue primitive (Variable _ v) = varPrimal v primitive (Constant v) = v@@ -106,9 +114,11 @@ varPrimal (JVP (JVPValue v _)) = primitive $ primal v  -- Evaluates a PrimExp using doOp-evalPrimExp :: M.Map VName ADValue -> PrimExp VName -> Maybe ADValue-evalPrimExp m (LeafExp n _) = M.lookup n m-evalPrimExp _ (ValueExp pv) = Just $ Constant pv+evalPrimExp :: M.Map VName ADValue -> PrimExp VName -> Either String ADValue+evalPrimExp m (LeafExp n _) =+  maybe (Left $ "Unknown variable " <> show n) Right $ M.lookup n m+evalPrimExp _ (ValueExp pv) =+  Right $ Constant pv evalPrimExp m (BinOpExp op x y) = do   x' <- evalPrimExp m x   y' <- evalPrimExp m y@@ -141,23 +151,25 @@ -- This function performs a mathematical operation on a -- list of operands, performing automatic differentiation -- if one or more operands is a Variable (of depth > 0)-doOp :: Op -> [ADValue] -> Maybe ADValue+doOp :: Op -> [ADValue] -> Either String ADValue doOp op o   | not $ opTypeMatch op (map primValueType pv) =       -- This function may be called with arguments of invalid types,       -- because it is used as part of an overloaded operator.-      Nothing+      Left $ unwords ["invalid types for op", show op, "and operands", show o]   | otherwise = do       let dep = case op of             OpCmp _ -> 0 -- AD is not well-defined for comparason operations             -- There are no derivatives for those written in             -- PrimExp (check lookupPDs)             _ -> maximum (map depth o)-      if dep == 0 then constCase else nonconstCase dep+      if dep == 0+        then maybe (Left "failed to evaluate const") Right constCase+        else nonconstCase dep   where     pv = map primitive o -    divideDepths :: Int -> ADValue -> Either ADValue ADVariable+    divideDepths :: Depth -> ADValue -> Either ADValue ADVariable     divideDepths _ v@(Constant {}) = Left v     divideDepths d v@(Variable d' v') = if d' < d then Left v else Right v' @@ -198,7 +210,7 @@       -- have to perform the necessary steps for AD        -- First, we calculate the value for the previous depth-      let oprev = map primal o+      let oprev = map (primalFor dep) o       vprev <- doOp op oprev        -- Then we separate the values of the maximum depth from@@ -209,28 +221,31 @@         -- Finally, we perform the necessary steps for the given         -- type of AD         Just (Right (VJP {})) ->-          Just . Variable dep . VJP . VJPValue $ vjpHandleOp op (map extractVJP o') vprev+          Right . Variable dep . VJP . VJPValue $+            vjpHandleOp op (map extractVJP o') vprev         Just (Right (JVP {})) ->-          Variable dep . JVP . JVPValue vprev <$> jvpHandleFn op (map extractJVP o')+          Variable dep . JVP . JVPValue vprev+            <$> jvpHandleOp op (map extractJVP o')         _ ->           -- Since the maximum depth is non-zero, there must be at           -- least one variable of depth > 0           error "find isRight"  calculatePDs :: Op -> [ADValue] -> [ADValue]-calculatePDs op p =+calculatePDs op args =   -- Create a unique VName for each operand-  let n = map (\i -> VName (nameFromString $ "x" ++ show i) i) [1 .. length p]+  let n = map (\i -> VName (nameFromString $ "x" ++ show i) i) [1 .. length args]       -- Put the operands in the environment-      m = M.fromList $ zip n p+      m = M.fromList $ zip n args        -- Look up, and calculate the partial derivative       -- of the operation with respect to each operand       pde =         fromMaybe (error "lookupPDs failed") $           lookupPDs op $-            map (`LeafExp` opReturnType op) n-   in map (fromMaybe (error "evalPrimExp failed") . evalPrimExp m) pde+            zipWith (\v val -> LeafExp v $ primValueType $ primitive val) n args+      res = map (either (error . ("evalPrimExp failed: " <>)) id . evalPrimExp m) pde+   in res  -- VJP / Reverse mode automatic differentiation-- -- In reverse mode AD, the entire computation@@ -244,7 +259,7 @@ data Tape   = -- | This represents a variable. Each variable is given a unique ID,     -- and has an initial value-    TapeID Int ADValue+    TapeID Depth ADValue   | -- | This represents a constant.     TapeConst ADValue   | -- | This represents the application of a mathematical operation.@@ -279,7 +294,7 @@   let s'' = case op of         OpConv op' ->           -- In case of type conversion, simply convert the sensitivity-          [ fromMaybe (error "deriveTape: doOp failed") $+          [ fromRight (error "deriveTape: doOp failed") $               doOp (OpConv $ flipConvOp op') [s]           ]         _ ->@@ -291,10 +306,10 @@       foldl' (M.unionWith add) M.empty pd   where     add x y =-      fromMaybe (error "deriveTape: add failed") $+      fromRight (error "deriveTape: add failed") $         doOp (OpBin $ addFor $ opReturnType op) [x, y]     mul x y =-      fromMaybe (error "deriveTape: mul failed") $+      fromRight (error "deriveTape: mul failed") $         doOp (OpBin $ mulFor $ opReturnType op) [x, y]  -- JVP / Forward mode automatic differentiation--@@ -304,27 +319,26 @@ data JVPValue = JVPValue ADValue ADValue   deriving (Show) --- | This calculates the derivative part of the JVPValue resulting+-- | This calculates the tangent part of the JVPValue resulting -- from the application of a mathematical operation on one or more -- JVPValues.-jvpHandleFn :: Op -> [Either ADValue JVPValue] -> Maybe ADValue-jvpHandleFn op p = do+jvpHandleOp :: Op -> [Either ADValue JVPValue] -> Either String ADValue+jvpHandleOp op p = do   case op of     OpConv _ ->       -- In case of type conversion, simply convert-      -- the old derivative-      doOp op [derivative $ head p]+      -- the old tangent+      doOp op [tangent $ head p]     _ -> do-      -- Calculate the new derivative using the chain-      -- rule+      -- Calculate the new tangent using the chain rule       let pds = calculatePDs op $ map primal' p-      vs <- zipWithM mul pds $ map derivative p-      foldM add (Constant $ blankPrimValue $ opReturnType op) vs+      vs <- zipWithM mul pds $ map tangent p+      foldM add (Constant $ blankPrimValue op_t) vs   where+    op_t = opReturnType op     primal' (Left v) = v     primal' (Right (JVPValue v _)) = v-    derivative (Left v) = Constant $ blankPrimValue $ primValueType $ primitive v-    derivative (Right (JVPValue _ d)) = d-+    tangent (Left _) = Constant $ blankPrimValue $ opReturnType op+    tangent (Right (JVPValue _ d)) = d     add x y = doOp (OpBin $ addFor $ opReturnType op) [x, y]     mul x y = doOp (OpBin $ mulFor $ opReturnType op) [x, y]

@@ -200,7 +200,8 @@ valueAccum :: (a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m) valueAccum f i v@(ValuePrim {}) = f i v valueAccum f i v@(ValueAD {}) = f i v-valueAccum f i (ValueRecord m) = second ValueRecord $ M.mapAccum (valueAccum f) i m+valueAccum f i (ValueRecord m) =+  second ValueRecord $ M.mapAccum (valueAccum f) i m valueAccum f i (ValueArray s a) = do   -- TODO: This could probably be better   -- Transform into a map

@@ -343,12 +343,12 @@ missingType :: (Pretty a) => Loc -> a -> Either TypeError b missingType loc name =   Left . TypeError loc mempty $-    "Module does not define a type named" <+> pretty name <> "."+    "Module does not define a type named" <+> dquotes (pretty name) <> "."  missingVal :: (Pretty a) => Loc -> a -> Either TypeError b missingVal loc name =   Left . TypeError loc mempty $-    "Module does not define a value named" <+> pretty name <> "."+    "Module does not define a value named" <+> dquotes (pretty name) <> "."  topLevelSize :: Loc -> VName -> Either TypeError b topLevelSize loc name =@@ -358,7 +358,7 @@ missingMod :: (Pretty a) => Loc -> a -> Either TypeError b missingMod loc name =   Left . TypeError loc mempty $-    "Module does not define a module named" <+> pretty name <> "."+    "Module does not define a module named" <+> dquotes (pretty name) <> "."  mismatchedType ::   Loc ->