packages feed

horde-ad (empty) → 0.1.0.0

raw patch · 73 files changed

+38382/−0 lines, 73 filesdep +Booleandep +assert-failuredep +atomic-counter

Dependencies added: Boolean, assert-failure, atomic-counter, base, bytestring, containers, criterion, data-default, deepseq, dependent-enummap, dependent-sum, enummapset, ghc-typelits-knownnat, ghc-typelits-natnormalise, horde-ad, ilist, inspection-testing, mnist-idx, orthotope, ox-arrays, random, some, tasty, tasty-hunit, tasty-quickcheck, vector, zlib

Files

+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Revision history for horde-ad++## [v0.1.0.0](https://github.com/Mikolaj/horde-ad/compare/v0.0.0.0...v0.1.0.0)++- First version. Released on an unsuspecting world.
+ CREDITS.md view
@@ -0,0 +1,25 @@+Credits go to++Mikolaj Konarski, the original author of the library++Conal Elliott, talking beautifully about the Essence of Automatic Differentiation++Simon Peyton Jones, who sketched in "Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation" talk at https://www.youtube.com/watch?v=qGkgLtpseMs algorithms that this Haskell codebase implements++Faustyna Krawiec, Neel Krishnaswami, Tom Ellis, Andrew Fitzgibbon, Richard Eisenberg, the remaining authors of the paper "Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation" from POPL 2022 that describes and proves correct a comprehensive AD formalism, including some ingenious algorithmic ideas this Haskell codebase implements++Oleg Grenrus, whose https://hackage.haskell.org/package/overloaded contains inspiring examples related to AD++Edward Kmett, author of the archetype Haskell https://hackage.haskell.org/package/ad library++Justin Le, author of the https://hackage.haskell.org/package/backprop library and the amazing articles at https://blog.jle.im++Lennart Augustsson, author of the incredibly clever shapely typed Haskell tensor library https://hackage.haskell.org/package/orthotope++Tom Smeding, who influenced the major implementation decisions of horde-ad and solved the hardest technical problems++and other benefactors and contributors, in chronological order:+Stanisław Findeisen, Cale Gibbard, Well-Typed LLP+++Authors of files `t10k-images-idx3-ubyte.gz`, `t10k-labels-idx1-ubyte.gz`, `train-images-idx3-ubyte.gz`, `train-labels-idx1-ubyte.gz` from http://yann.lecun.com/exdb/mnist, copied into this repo into directory `samplesData/`, copyright (according to https://keras.io/api/datasets/mnist) Yann LeCun and Corinna Cortes for the MNIST dataset, which is a derivative work from original NIST datasets, copyright by their authors and made available under the terms of the Creative Commons Attribution-Share Alike 3.0 license.
+ LICENSE view
@@ -0,0 +1,26 @@+BSD 3-Clause License++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions+are met:+1. Redistributions of source code must retain the above copyright notice, this+   list of conditions and the following disclaimer.+2. Redistributions in binary form must reproduce the above copyright notice,+   this list of conditions and the following disclaimer in the documentation+   and/or other materials provided with the distribution.+3. Neither the name of the copyright holder nor the names of its+   contributors may be used to endorse or promote products derived from+   this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,275 @@+# HORDE-AD: Higher Order Reverse Derivatives Efficiently++[![Hackage](https://img.shields.io/hackage/v/horde-ad.svg)](https://hackage.haskell.org/package/horde-ad)++Welcome to the Automatic Differentiation library originally inspired by the paper [_"Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation"_](https://dl.acm.org/doi/10.1145/3498710). Compared to the paper and to classic taping AD Haskell packages, the library additionally efficiently supports array operations and generation of symbolic derivative programs, though the efficiency is confined to a narrowly typed class of source programs with limited higher-orderness. A detailed account of the extension is in the paper [_"Dual-Numbers Reverse AD for Functional Array Languages"_](http://arxiv.org/abs/2507.12640) by Tom Smeding, Mikolaj Konarski, Simon Peyton Jones and Andrew Fitzgibbon.++This is an early prototype, both in terms of the engine performance, the API and the preliminary tools and examples built with it. At this development stage, it's not coded defensively but exactly the opposite: it will fail on cases not found in current tests so that new code and tests have to be added and old code optimized for the new specimens reported in the wild. The user should also be ready to add missing primitives and any obvious tools that should be predefined but aren't, such as weight normalization (https://github.com/Mikolaj/horde-ad/issues/42). It's already possible to differentiate basic neural network architectures, such as fully connected, recurrent, convolutional and residual. The library should also be suitable for defining exotic machine learning architectures and non-machine learning systems, given that no notion of a neural network nor of a computation graph are hardwired into the formalism, but instead they are compositionally and type-safely built up from general automatic differentiation building blocks.++Mature Haskell libraries with similar capabilities, but varying efficiency, are https://hackage.haskell.org/package/ad and https://hackage.haskell.org/package/backprop. See also https://github.com/Mikolaj/horde-ad/blob/master/CREDITS.md. Benchmarks suggest that horde-ad has competitive performance on CPU.+<!--+The benchmarks at SOMEWHERE show that this library has performance highly competitive with (i.e. faster than) those and PyTorch on CPU.+-->+It is hoped that the (well-typed) separation of AD logic and tensor manipulation backend will enable similar speedups on numerical accelerators, when their support is implemented. Contributions to this and other tasks are very welcome. The newcomer-friendly tickets are listed at https://github.com/Mikolaj/horde-ad/issues?q=is%3Aissue%20state%3Aopen%20label%3A%22good%20first%20issue%22. Please don't hesitate to ask questions on github, on Matrix, via email.+++## Computing the derivative of a simple function++Here is an example of a Haskell function to be differentiated:+```hs+-- A function that goes from R^3 to R.+foo :: RealFloat a => (a, a, a) -> a+foo (x, y, z) =+  let w = x * sin y+  in atan2 z w + z * w  -- note that w appears twice+```++The gradient of `foo` instantiated to `Double` can be expressed in Haskell with horde-ad as:+```hs+gradFooDouble :: (Double, Double, Double) -> (Double, Double, Double)+gradFooDouble = fromDValue . cgrad foo . fromValue+```++which can be verified by computing the gradient at `(1.1, 2.2, 3.3)`:+```hs+>>> gradFooDouble (1.1, 2.2, 3.3)+(2.4396285219055063, -1.953374825727421, 0.9654825811012627)+```++We can instantiate `foo` to matrices; the operations within (`sin`, `+`, `*`, etc.) applying elementwise:+```hs+type Matrix2x2 f r = f (TKS '[2, 2] r)  -- TKS means shapely-typed tensor kind+type ThreeMatrices r = (Matrix2x2 Concrete r, Matrix2x2 Concrete r, Matrix2x2 Concrete r)+threeSimpleMatrices :: ThreeMatrices Double+threeSimpleMatrices = (srepl 1.1, srepl 2.2, srepl 3.3)  -- srepl replicates its argument to fill the whole matrix+fooMatrixValue :: Matrix2x2 Concrete Double+fooMatrixValue = foo threeSimpleMatrices+>>> fooMatrixValue+sfromListLinear [2,2] [4.242393641025528,4.242393641025528,4.242393641025528,4.242393641025528])+```++Instantiated to matrices, `foo` now returns a matrix, not a scalar &mdash; but a gradient can only be computed of a function that returns a scalar.+To remediate this, let's sum the whole output of `foo` and only then compute its gradient: (note that `ssum0` returns a zero-dimensional array; `kfromS` extracts the (single) scalar from that)+```hs+gradSumFooMatrix :: ThreeMatrices Double -> ThreeMatrices Double+gradSumFooMatrix = cgrad (kfromS . ssum0 . foo)+```++This works as well as before:+```hs+>>> gradSumFooMatrix threeSimpleMatrices+(sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])+```++### Efficiency: preserving sharing++We noted above that `w` appears twice in `foo`.  A property of tracing-based AD systems is that such re-use may not be captured, with explosive results.+In `cgrad`, such sharing is preserved, so `w` is processed only once during gradient computation and this property is guaranteed for the `cgrad` tool universally, without any action required from the user.+`horde-ad` also allows computing _symbolic_ derivative programs: using this API, a program is differentiated only once, after which it can be run on many different input values.+In this case, however, sharing is _not_ automatically preserved, so shared variables have to be explicitly marked using `tlet`, as shown below in `fooLet`.+This also makes the type of the function more specific: it now does not work on an arbitrary `Num` any more, but instead on an arbitrary `horde-ad` tensor that implements the standard arithmetic operations, some of which (e.g., `atan2H`) are implemented in custom numeric classes.+```hs+fooLet :: (RealFloatH (f r), ADReady f)  -- ADReady means the container type f supports AD+       => (f r, f r, f r) -> f r+fooLet (x, y, z) =+  tlet (x * sin y) $ \w ->+    atan2H z w + z * w  -- note that w still appears twice+```++### Vector-Jacobian product (VJP) and symbolic derivatives++The most general symbolic reverse derivative program for this function can be obtained using the `vjpArtifact` tool. Because the `vjp` family of tools permits non-scalar codomains (but expects an incoming cotangent argument to compensate, visible in the code below as `dret`), we illustrate it using the original `fooLet` from the previous section, without the need to add `ssum0`.+```hs+artifact :: AstArtifactRev (X (ThreeConcreteMatrices Double)) (TKS '[2, 2] Double)+artifact = vjpArtifact fooLet threeSimpleMatrices+```++The vector-Jacobian product program (presented below with additional formatting) looks like ordinary functional code with nested pairs and projections representing tuples. A quick inspection of the code reveals that computations are not repeated, which is thanks to the `tlet` used above.++```hs+>>> printArtifactPretty artifact+"\dret m1 ->+   let m3 = sin (tproject2 (tproject1 m1))+       m4 = tproject1 (tproject1 m1) * m3+       m5 = recip (tproject2 m1 * tproject2 m1 + m4 * m4)+       m7 = (negate (tproject2 m1) * m5) * dret + tproject2 m1 * dret+    in tpair+         (tpair (m3 * m7)+                (cos (tproject2 (tproject1 m1)) * (tproject1 (tproject1 m1) * m7)))+         ((m4 * m5) * dret + m4 * dret)"+```++A concrete value of this symbolic reverse derivative at the same input as before can be obtained by interpreting its program in the context of the operations supplied by the horde-ad library. (Note that the output happens to be the same as `gradSumFooMatrix threeSimpleMatrices` above, which used `cgrad` on `kfromS . ssum0 . foo`; the reason is that `srepl 1.0` happens to be the reverse derivative of `kfromS . ssum0`.)+```hs+>>> vjpInterpretArtifact artifact (toTarget threeSimpleMatrices) (srepl 1.0)+((sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627]) :: ThreeConcreteMatrices Double)+```++Note that, as evidenced by the `printArtifactPretty` call above, `artifact` contains the complete and simplified code of the VJP of `fooLet`, so repeated calls of `vjpInterpretArtifact artifact` won't ever repeat differentiation nor simplification and will only incur the cost of straightforward interpretation. The repeated call would fail with an error if the provided argument had a different shape than `threeSimpleMatrices`. However, for the examples we show here, such a scenario is ruled out by the types, because all tensors we present are shaped, meaning their full shape is stated in their type and so can't differ for two (tuples of) tensors of the same type. More loosely-typed variants of all the tensor operations, where runtime checks can really fail, are available in the horde-ad API and can be mixed and matched freely.++A shorthand that creates a symbolic gradient program, simplifies it and interprets it with a given input on the default CPU backend is called `grad` and is used exactly the same as (but with often much better performance on the same program than) `cgrad`:+```hs+>>> grad (kfromS . ssum0 . fooLet) threeSimpleMatrices+(sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])+```+++## For all shapes and sizes++An important feature of this library is a type system for tensor shape+arithmetic. The following code is part of a convolutional neural network+definition, for which horde-ad computes the shape of the gradient from+the shape of the input data and the initial parameters.+The Haskell compiler is able to infer many tensor shapes, deriving them both+from dynamic dimension arguments (the first two lines of parameters+to the function below) and from static type-level hints.++Let's look at the body of the `convMnistTwoS` function before we look at its signature. It is common to see neural network code like that, with shape annotations in comments, hidden from the compiler:+```hs+convMnistTwoS+  kh@SNat kw@SNat h@SNat w@SNat+  c_out@SNat _n_hidden@SNat batch_size@SNat+    -- integer parameters denoting basic dimensions+  input              -- input images, shape [batch_size, 1, h, w]+  (ker1, bias1)      -- layer1 kernel, shape [c_out, 1, kh+1, kw+1]; and bias, shape [c_out]+  (ker2, bias2)      -- layer2 kernel, shape [c_out, c_out, kh+1, kw+1]; and bias, shape [c_out]+  ( weightsDense     -- dense layer weights, shape [n_hidden, c_out * (h/4) * (w/4)]+  , biasesDense )    -- dense layer biases, shape [n_hidden]+  ( weightsReadout   -- readout layer weights, shape [10, n_hidden]+  , biasesReadout )  -- readout layer biases, shape [10]+  =                  -- -> output classification, shape [10, batch_size]+  assumeEquality @(Div (Div w 2) 2) @(Div w 4) $+  assumeEquality @(Div (Div h 2) 2) @(Div h 4) $+  let t1 = convMnistLayerS kh kw h w+                           (SNat @1) c_out batch_size+                           ker1 (sfromPrimal input) bias1+      t2 = convMnistLayerS kh kw (SNat @(h `Div` 2)) (SNat @(w `Div` 2))+                           c_out c_out batch_size+                           ker2 t1 bias2+      m1 = sreshape t2+      denseLayer = weightsDense `smatmul2` str m1+                   + str (sreplicate biasesDense)+  in weightsReadout `smatmul2` reluS denseLayer+     + str (sreplicate biasesReadout)+```+But we don't just want the shapes in comments and in runtime expressions; we want them as a compiler-verified documentation in the form of the type signature of the function:+```hs+convMnistTwoS+  :: forall kh kw h w c_out n_hidden batch_size target r.+     ( 1 <= kh  -- kernel height is large enough+     , 1 <= kw  -- kernel width is large enough+     , ADReady target, GoodScalar r, Differentiable r )+         -- GoodScalar means r is supported as array elements by horde-ad+  => SNat kh -> SNat kw -> SNat h -> SNat w+  -> SNat c_out -> SNat n_hidden -> SNat batch_size+       -- ^ these boilerplate lines tie type parameters to the corresponding+       -- SNat value parameters denoting basic dimensions+  -> PrimalOf target (TKS '[batch_size, 1, h, w] r)  -- `input` shape [batch_size, 1, h, w]+  -> ( ( target (TKS '[c_out, 1, kh + 1, kw + 1] r)  -- `ker1` shape [c_out, 1, kh+1, kw+1]+       , target (TKS '[c_out] r) )                   -- `bias2` shape [c_out]+     , ( target (TKS '[c_out, c_out, kh + 1, kw + 1] r)  -- `ker2` shape [c_out, c_out, kh+1, kw+1]+       , target (TKS '[c_out] r) )                       -- `bias2` shape [c_out]+     , ( target (TKS '[n_hidden, c_out * (h `Div` 4) * (w `Div` 4) ] r)+           -- `weightsDense` shape [n_hidden, c_out * (h/4) * (w/4)]+       , target (TKS '[n_hidden] r) )    -- `biasesDense` shape [n_hidden]+     , ( target (TKS '[10, n_hidden] r)  -- `weightsReadout` shape [10, n_hidden]+       , target (TKS '[10] r) ) )        -- `biasesReadout` shape [10]+  -> target (TKS '[SizeMnistLabel, batch_size] r)  -- -> `output` shape [10, batch_size]+convMnistTwoS+  kh@SNat kw@SNat h@SNat w@SNat+  c_out@SNat _n_hidden@SNat batch_size@SNat+  input  -- input images+  (ker1, bias1)  -- layer1 kernel+  (ker2, bias2)  -- layer2 kernel+  ( weightsDense  -- dense layer weights+  , biasesDense )  -- dense layer biases+  ( weightsReadout  -- readout layer weights+  , biasesReadout )  -- readout layer biases+  =+...+```++This style gets verbose and the Haskell compiler needs some convincing to accept such programs, but type-safety is the reward. In practice, at least the parameters of the objective function are best expressed with shaped tensors, while the implementation can (zero-cost) convert the tensors to loosely typed variants as needed.++The full neural network definition from which this function is taken can be found at++https://github.com/Mikolaj/horde-ad/tree/master/example++in file `MnistCnnShaped2.hs` and the directory contains several other sample neural networks for MNIST digit classification. Among them are recurrent, convolutional and fully connected networks based on fully typed tensors (sizes of all dimensions are tracked in the types, as above) as well as on the weakly typed ranked tensors, where only tensor ranks are tracked. It's possible to mix the two typing styles within one function signature and even within one shape description.+++Compilation from source+-----------------------++Horde-ad uses [ox-arrays](https://hackage.haskell.org/package/ox-arrays)+as its CPU backend library, which in turn is inspired by and depends+on [orthotope](https://hackage.haskell.org/package/orthotope).+Neither of these packages require any special installation.+Some of the other+[Haskell packages we depend on](https://github.com/Mikolaj/horde-ad/blob/master/horde-ad.cabal)+need their usual C library dependencies to be installed manually,+e.g., package zlib needs the C library zlib1g-dev or an equivalent.+At this time, we don't depend on any GPU hardware nor bindings.++For development, copying the included `cabal.project.local.development`+to `cabal.project.local` provides a sensible default to run `cabal build` with+and get compilation results relatively fast. For extensive testing,+on the other hand, a command like++    cabal test minimalTest --enable-optimization++ensures that the code is compiled with optimization and consequently+executes the rather computation-intensive testsuites in reasonable time.+++Running tests+-------------++The `parallelTest` test suite consists of large tests and runs in parallel++    cabal test parallelTest --enable-optimization++which is likely to cause the extra printf messages coming from within+the tests to be out of order. To keep your screen tidy, simply redirect+`stderr`, e.g.,++    cabal test parallelTest --enable-optimization 2>/dev/null++The remainder of the test suite is set up to run sequentially, in part+to simplify automatic checking of results that may vary slightly+depending on execution order++    cabal test CAFlessTest --enable-optimization++The tests in this suite also don't contribute big Haskell CAFs,+which makes them more reliable micro-benchmarks. Ordinary+benchmarks that use package criterion are even more trustworthy,+but they don't have a comparable coverage at this time.+++Coding style+------------++Stylish Haskell is used for slight auto-formatting at buffer save; see+[.stylish-haskell.yaml](https://github.com/Mikolaj/horde-ad/blob/master/.stylish-haskell.yaml).+As defined in the file, indentation is 2 spaces wide and screen is+80-columns wide. Spaces are used, not tabs. Spurious whitespace avoided.+Spaces around arithmetic operators encouraged.+Generally, relax and try to stick to the style apparent in a file+you are editing. Put big formatting changes in separate commits.++Haddocks should be provided for all module headers and for the main+functions and types from the most important modules.+Apart of that, only particularly significant functions and types+are distinguished by having a haddock. If minor ones have comments,+they should not be haddocks and they are permitted to describe+implementation details and be out of date. Prefer assertions in place+of comments, unless too verbose.+++Copyright+---------++Copyright 2023--2025 Mikolaj Konarski, Well-Typed LLP and others (see git history)++License: BSD-3-Clause (see file LICENSE)
+ bench/LongMnistBench.hs view
@@ -0,0 +1,18 @@+module Main (main) where++import Prelude++import Criterion.Main++import BenchMnistTools++main :: IO ()+main = defaultMain+         [ mnistBGroup1VTA 400+         , mnistBGroup1VTO 400+         , mnistBGroup2VTA 400+         , mnistBGroup2VTC 400+         , mnistBGroup2VTO 400+         , mnistBGroup2VTOZ 400+         , mnistBGroup2VTOX 400+         ]
+ bench/LongProdBench.hs view
@@ -0,0 +1,24 @@+module Main (main) where++import Prelude++import Control.DeepSeq+import Criterion.Main+import System.Random++import BenchProdTools++allxs :: [Double]+allxs = let xs = map (+ 0.55) $ randoms (mkStdGen 42)+        in deepseq (take 1000000 xs) xs++main :: IO ()+main =+  defaultMain  -- skips the tiny benchmarks+    [ bgroup1e4 allxs+{- OOMs, probably for a good reason (huge terms); TODO: diagnose and enable 1e5+    , bgroup1e5 allxs+    , bgroup1e6 allxs+    , bgroup1e7 allxs+    , bgroup5e7 allxs -}+    ]
+ bench/RealisticMnistBench.hs view
@@ -0,0 +1,21 @@+module Main (main) where++import Prelude++import Criterion.Main+import System.Random++import MnistData++import BenchMnistTools++main :: IO ()+main = defaultMain+         [ mnistBGroup1VTA 4000+         , mnistBGroup1VTO 4000+         , mnistBGroup2VTA 4000+         , mnistBGroup2VTC 4000+         , mnistBGroup2VTO 4000+         , mnistBGroup2VTOZ 4000+         , mnistBGroup2VTOX 4000+         ]
+ bench/ShortMnistForCI.hs view
@@ -0,0 +1,18 @@+module Main (main) where++import Prelude++import Criterion.Main++import BenchMnistTools++main :: IO ()+main = defaultMain+         [ mnistBGroup1VTA 40+         , mnistBGroup1VTO 40+         , mnistBGroup2VTA 40+         , mnistBGroup2VTC 40+         , mnistBGroup2VTO 40+         , mnistBGroup2VTOZ 40+         , mnistBGroup2VTOX 40+         ]
+ bench/ShortProdForCI.hs view
@@ -0,0 +1,20 @@+module Main (main) where++import Prelude++import Control.DeepSeq+import Criterion.Main+import System.Random++import BenchProdTools++allxs :: [Double]+allxs = let xs = map (+ 0.55) $ randoms (mkStdGen 42)+        in deepseq (take 50000000 xs) xs++main :: IO ()+main =+  defaultMain+    [ bgroup100 allxs+    , bgroup1000 allxs+    ]
+ bench/common/BenchMnistTools.hs view
@@ -0,0 +1,498 @@+{-# LANGUAGE TemplateHaskell #-}+{-# OPTIONS_GHC -Wno-missing-export-lists #-}+-- | A set of benchmarks using fully connected MNIST neural networks.+module BenchMnistTools where++import Prelude++import Control.Arrow ((***))+import Control.DeepSeq (NFData (..))+import Criterion.Main+import Data.Default qualified as Default+import Data.Proxy (Proxy (Proxy))+import GHC.Exts (WithDict)+import GHC.TypeLits (KnownNat)+import System.Random+import Test.Inspection+import Type.Reflection (Typeable)++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.OpsConcrete ()+import HordeAd.External.OptimizerTools++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Shaped.Shape++import MnistData+import MnistFcnnRanked1 qualified+import MnistFcnnRanked2 (XParams2)+import MnistFcnnRanked2 qualified++-- * Using lists of vectors, which is rank 1++type XParams widthHidden widthHidden2 r =+  X (MnistFcnnRanked1.ADFcnnMnist1Parameters+       Concrete widthHidden widthHidden2 r)++-- POPL differentiation, straight via the ADVal instance of RankedTensor,+-- which side-steps vectorization.+mnistTrainBench1VTA+  :: forall r. r ~ Double+  => String+  -> Int -> Int -> Double -> Int -> [MnistDataLinearR r]+  -> Benchmark+mnistTrainBench1VTA prefix widthHiddenInt widthHidden2Int+                    gamma batchSize xs =+  withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->+  withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters+                    Concrete widthHidden widthHidden2 r+      valsInit = fst $ randomValue 1 (mkStdGen 44)+      targetInit :: Concrete (XParams widthHidden widthHidden2 r)+      targetInit = toTarget @Concrete valsInit+  in do+    let f :: MnistDataLinearR Double+          -> ADVal Concrete (XParams widthHidden widthHidden2 Double)+          -> ADVal Concrete (TKScalar Double)+        f (glyph, label) adinputs =+          MnistFcnnRanked1.afcnnMnistLoss1+            widthHiddenSNat widthHidden2SNat+            (rconcrete glyph, rconcrete label) (fromTarget adinputs)+        chunk = take batchSize xs+        gradf c = fst $ sgdSTK knownSTK gamma f c targetInit+        name =+          prefix+          ++ unwords+               [ "v" ++ show (widthSTK+                              $ knownSTK @(XParams widthHidden widthHidden2 r))+               , "m0" ++ " =" ++ show (tsize knownSTK targetInit) ]+    bench name $ nf gradf chunk++mnistTestBench1VTA+  :: forall r. r ~ Double+  => String+  -> Int -> Int -> Double -> Int -> [MnistDataLinearR r]+  -> Benchmark+mnistTestBench1VTA prefix widthHiddenInt widthHidden2Int+                   _gamma batchSize xs =+  withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->+  withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters+                    Concrete widthHidden widthHidden2 r+      valsInit = fst $ randomValue 1 (mkStdGen 44)+      targetInit :: Concrete (XParams widthHidden widthHidden2 r)+      targetInit = toTarget @Concrete valsInit+      ftest :: [MnistDataLinearR r]+            -> MnistFcnnRanked1.ADFcnnMnist1Parameters+                 Concrete widthHidden widthHidden2 r+            -> r+      ftest = MnistFcnnRanked1.afcnnMnistTest1 widthHiddenSNat widthHidden2SNat+  in do+    let chunk = take batchSize xs+        score c = ftest c valsInit+        name =+          "test " ++ prefix+          ++ unwords+               [ "v" ++ show (widthSTK+                              $ knownSTK @(XParams widthHidden widthHidden2 r))+               , "m0" ++ " =" ++ show (tsize knownSTK targetInit) ]+    bench name $ whnf score chunk++mnistBGroup1VTA :: Int -> Benchmark+mnistBGroup1VTA chunkLength =+  env (do+    testData0 <- loadMnistData testGlyphsPath testLabelsPath  -- 10k total+    let testData = shuffle (mkStdGen 42) testData0+    return $! map mkMnistDataLinearR $ take chunkLength testData) $+  \ xs ->+  bgroup ("2-hidden-layer rank 1 VTA MNIST nn with samples: "+          ++ show chunkLength)+       [ mnistTestBench1VTA "30|10 " 30 10 0.02 chunkLength xs+       , mnistTrainBench1VTA "30|10 " 30 10 0.02 chunkLength xs+           -- toy width+       , mnistTestBench1VTA "300|100 " 300 100 0.02 chunkLength xs+       , mnistTrainBench1VTA "300|100 " 300 100 0.02 chunkLength xs+           -- ordinary width+       , mnistTestBench1VTA "500|150 " 500 150 0.02 chunkLength xs+       , mnistTrainBench1VTA "500|150 " 500 150 0.02 chunkLength xs+           -- another common width+       , mnistTrainBench1VTA "1500|500 " 1500 500 0.02 chunkLength xs+       ]++-- JAX differentiation, Ast term built and differentiated only once+-- and the result interpreted with different inputs in each gradient+-- descent iteration.+mnistTrainBench1VTO+  :: forall r. r ~ Double+  => String+  -> Int -> Int -> Double -> Int -> [MnistDataLinearR r]+  -> Benchmark+mnistTrainBench1VTO prefix widthHiddenInt widthHidden2Int+                    gamma batchSize xs =+  withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->+  withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters+                    Concrete widthHidden widthHidden2 r+      valsInit = fst $ randomValue 1 (mkStdGen 44)+      targetInit :: Concrete (XParams widthHidden widthHidden2 r)+      targetInit = toTarget @Concrete valsInit+  in do+{-      -- g is not enough to specialize to Double instead of to r,+        -- despite the declaration of r ~ Double above+        f :: ( MnistFcnnRanked1.ADFcnnMnist1Parameters+                 (AstTensor AstMethodLet FullSpan)+                 widthHidden widthHidden2 r+             , ( AstTensor AstMethodLet FullSpan (TKR 1 r)+               , AstTensor AstMethodLet FullSpan (TKR 1 r) ) )+          -> AstTensor AstMethodLet FullSpan (TKScalar r)+        f = \ (pars, (glyphR, labelR)) ->+          MnistFcnnRanked1.afcnnMnistLoss1+            widthHiddenSNat widthHidden2SNat+            (glyphR, labelR) pars+        g :: ( MnistFcnnRanked1.ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 Double, ( AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 1 Double) ) ) -> AstTensor AstMethodLet FullSpan (TKScalar Double)+        g = f+-}+    let dataInit = case xs of+          d : _ -> (rconcrete *** rconcrete) d+          [] -> error "empty test data"+        f :: ( MnistFcnnRanked1.ADFcnnMnist1Parameters+                 (AstTensor AstMethodLet FullSpan)+                 widthHidden widthHidden2 Double+             , ( AstTensor AstMethodLet FullSpan (TKR 1 Double)+               , AstTensor AstMethodLet FullSpan (TKR 1 Double) ) )+          -> AstTensor AstMethodLet FullSpan (TKScalar Double)+        f = \ (pars, (glyphR, labelR)) ->+          MnistFcnnRanked1.afcnnMnistLoss1 @_ @Double+            widthHiddenSNat widthHidden2SNat+            (glyphR, labelR) pars+        artRaw = gradArtifact f (valsInit, dataInit)+        art = simplifyArtifactGradient artRaw+        go :: [MnistDataLinearR r]+           -> Concrete (XParams widthHidden widthHidden2 r)+           -> Concrete (XParams widthHidden widthHidden2 r)+        go [] parameters = parameters+        go ((glyph, label) : rest) !parameters =+          let parametersAndInput =+                tpair parameters (tpair (rconcrete glyph) (rconcrete label))+              gradient = tproject1 $ fst+                         $ revInterpretArtifact art parametersAndInput Nothing+          in go rest (updateWithGradient gamma knownSTK parameters gradient)+        chunk = take batchSize xs+        gradf c = go c targetInit+        name =+          prefix+          ++ unwords+               [ "v" ++ show (widthSTK+                              $ knownSTK @(XParams widthHidden widthHidden2 r))+               , "m0" ++ " =" ++ show (tsize knownSTK targetInit) ]+    bench name $ nf gradf chunk++mnistTestBench1VTO+  :: forall r. r ~ Double+  => String+  -> Int -> Int -> Double -> Int -> [MnistDataLinearR r]+  -> Benchmark+mnistTestBench1VTO = mnistTestBench1VTA++mnistBGroup1VTO :: Int -> Benchmark+mnistBGroup1VTO chunkLength =+  env (do+    testData0 <- loadMnistData testGlyphsPath testLabelsPath  -- 10k total+    let testData = shuffle (mkStdGen 42) testData0+    return $! map mkMnistDataLinearR $ take chunkLength testData) $+  \ xs ->+  bgroup ("2-hidden-layer rank 1 VTO MNIST nn with samples: "+          ++ show chunkLength)+       [ mnistTestBench1VTO "30|10 " 30 10 0.02 chunkLength xs+       , mnistTrainBench1VTO "30|10 " 30 10 0.02 chunkLength xs+       , mnistTestBench1VTO "300|100 " 300 100 0.02 chunkLength xs+       , mnistTrainBench1VTO "300|100 " 300 100 0.02 chunkLength xs+       , mnistTestBench1VTO "500|150 " 500 150 0.02 chunkLength xs+       , mnistTrainBench1VTO "500|150 " 500 150 0.02 chunkLength xs+       , mnistTrainBench1VTO "1500|500 " 1500 500 0.02 chunkLength xs+       ]+++-- * Using matrices, which is rank 2++-- POPL differentiation, straight via the ADVal instance of RankedTensor,+-- which side-steps vectorization.+mnistTrainBench2VTA+  :: forall r. r ~ Double+  => String+  -> Int -> Int -> Double -> Int -> [MnistDataLinearR r]+  -> Benchmark+mnistTrainBench2VTA prefix widthHidden widthHidden2+                    gamma batchSize xs =+  withSNat widthHidden $ \(SNat @widthHidden) ->+  withSNat widthHidden2 $ \(SNat @widthHidden2) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue+            @(Concrete (X (MnistFcnnRanked2.ADFcnnMnist2ParametersShaped+                             Concrete widthHidden widthHidden2 r Float)))+            1 (mkStdGen 44)+  in do+{-    let f :: MnistDataLinearR r -> ADVal Concrete (XParams2 r Float)+          -> ADVal Concrete (TKScalar r)+        f (glyph, label) adinputs =+          MnistFcnnRanked2.afcnnMnistLoss2+            (rconcrete glyph, rconcrete label) (fromTarget adinputs) -}+    let f :: MnistDataLinearR Double -> ADVal Concrete (XParams2 Double Float)+          -> ADVal Concrete (TKScalar Double)+        f (glyph, label) adinputs =+          MnistFcnnRanked2.afcnnMnistLoss2+            (rconcrete glyph, rconcrete label) (fromTarget adinputs)+        chunk = take batchSize xs+        gradf c = fst $ sgd gamma f c targetInit+        name =+          prefix+          ++ unwords+               [ "v0 m" ++ show (widthSTK $ knownSTK @(XParams2 r Float))+               , "=" ++ show (tsize knownSTK targetInit) ]+    bench name $ nf gradf chunk++mnistTestBench2VTA+  :: forall r. r ~ Double+  => String+  -> Int -> Int -> Double -> Int -> [MnistDataLinearR r]+  -> Benchmark+mnistTestBench2VTA prefix widthHidden widthHidden2+                   _gamma batchSize xs =+  withSNat widthHidden $ \(SNat @widthHidden) ->+  withSNat widthHidden2 $ \(SNat @widthHidden2) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue+            @(Concrete (X (MnistFcnnRanked2.ADFcnnMnist2ParametersShaped+                             Concrete widthHidden widthHidden2 r Float)))+            1 (mkStdGen 44)+  in do+    let chunk = take batchSize xs+        score c = MnistFcnnRanked2.afcnnMnistTest2 c (fromTarget targetInit)+        name =+          "test " ++ prefix+          ++ unwords+               [ "v0 m" ++ show (widthSTK $ knownSTK @(XParams2 r Float))+               , "=" ++ show (tsize knownSTK targetInit) ]+    bench name $ whnf score chunk++mnistBGroup2VTA :: Int -> Benchmark+mnistBGroup2VTA chunkLength =+  env (do+    testData0 <- loadMnistData testGlyphsPath testLabelsPath  -- 10k total+    let testData = shuffle (mkStdGen 42) testData0+    return $! map mkMnistDataLinearR $ take chunkLength testData) $+  \ xs ->+  bgroup ("2-hidden-layer rank 2 VTA MNIST nn with samples: "+          ++ show chunkLength)+       [ mnistTestBench2VTA "30|10 "30 10 0.02 chunkLength xs+       , mnistTrainBench2VTA "30|10 " 30 10 0.02 chunkLength xs+       , mnistTestBench2VTA "300|100 " 300 100 0.02 chunkLength xs+       , mnistTrainBench2VTA "300|100 " 300 100 0.02 chunkLength xs+       , mnistTestBench2VTA "500|150 " 500 150 0.02 chunkLength xs+       , mnistTrainBench2VTA "500|150 " 500 150 0.02 chunkLength xs+       , mnistTrainBench2VTA "1500|500 " 1500 500 0.02 chunkLength xs+       ]++-- JAX differentiation, Ast term built and differentiated only once+-- and the result interpreted with different inputs in each gradient+-- descent iteration.++-- Only compilation time.+mnistTrainBench2VTC+  :: String+  -> Int -> Int+  -> Benchmark+mnistTrainBench2VTC prefix widthHidden widthHidden2 =+  bench prefix+  $ whnf (simplifyArtifactGradient . snd+          . MnistFcnnRanked2.mnistTrainBench2VTOGradient+              @Double (Proxy @Float) IgnoreIncomingCotangent+              1 (mkStdGen 44) widthHidden)+         widthHidden2++mnistBGroup2VTC :: Int -> Benchmark+mnistBGroup2VTC chunkLength =+  bgroup ("2-hidden-layer rank 2 VTC compilation MNIST nn with samples: "+          ++ show chunkLength)+       [ mnistTrainBench2VTC "30|10 " 30 10+       , mnistTrainBench2VTC "300|100 " 300 100+       , mnistTrainBench2VTC "500|150 " 500 150+       , mnistTrainBench2VTC "1500|500 " 1500 500+       ]++-- The same as above, but only runtime.+mnistTrainBench2VTO+  :: forall r. r ~ Double+  => String+  -> Double -> Int -> [MnistDataLinearR r]+  -> ( Concrete (XParams2 r Float)+     , AstArtifactRev+         (TKProduct+            (XParams2 r Float)+            (TKProduct (TKR2 1 (TKScalar Double))+                       (TKR2 1 (TKScalar Double))))+         (TKScalar r) )+  -> Benchmark+mnistTrainBench2VTO prefix gamma batchSize xs (targetInit, art) = do+    let go :: [MnistDataLinearR r] -> Concrete (XParams2 r Float)+           -> Concrete (XParams2 r Float)+        go [] parameters = parameters+        go ((glyph, label) : rest) !parameters =+          let parametersAndInput =+                tpair parameters (tpair (rconcrete glyph) (rconcrete label))+              gradient = tproject1 $ fst+                         $ revInterpretArtifact art parametersAndInput Nothing+          in go rest (updateWithGradient gamma knownSTK parameters gradient)+        chunk = take batchSize xs+        gradf c = go c targetInit+        name =+          prefix+          ++ unwords+               [ "v0 m" ++ show (widthSTK $ knownSTK @(XParams2 r Float))+               , "=" ++ show (tsize knownSTK targetInit) ]+    bench name $ nf gradf chunk++mnistBGroup2VTO :: Int -> Benchmark+mnistBGroup2VTO chunkLength =+  let (!targetInit, !artRaw) =+        MnistFcnnRanked2.mnistTrainBench2VTOGradient+          @Double (Proxy @Float) IgnoreIncomingCotangent+          1 (mkStdGen 44) 1500 500+      !art = simplifyArtifactGradient artRaw  -- no NFData for AST+  in env (do+    testData0 <- loadMnistData testGlyphsPath testLabelsPath  -- 10k total+    let testData = shuffle (mkStdGen 42) testData0+    return $! map mkMnistDataLinearR $ take chunkLength testData) $+  \ xs ->+   bgroup ("2-hidden-layer rank 2 VTO runtime MNIST nn with samples: "+           ++ show chunkLength)+     [ mnistTrainBench2VTO "1500|500 " 0.02 chunkLength xs (targetInit, art)+     ]++-- The same as above, but without simplifying the gradient.+mnistTrainBench2VTOZ+  :: forall r. r ~ Double+  => String+  -> Double -> Int -> [MnistDataLinearR r]+  -> ( Concrete (XParams2 r Float)+     , AstArtifactRev+         (TKProduct+            (XParams2 r Float)+            (TKProduct (TKR2 1 (TKScalar Double))+                       (TKR2 1 (TKScalar Double))))+         (TKScalar r) )+  -> Benchmark+mnistTrainBench2VTOZ prefix gamma batchSize xs (targetInit, art) = do+    let go :: [MnistDataLinearR r] -> Concrete (XParams2 r Float)+           -> Concrete (XParams2 r Float)+        go [] parameters = parameters+        go ((glyph, label) : rest) !parameters =+          let parametersAndInput =+                tpair parameters (tpair (rconcrete glyph) (rconcrete label))+              gradient = tproject1 $ fst+                         $ revInterpretArtifact art parametersAndInput Nothing+          in go rest (updateWithGradient gamma knownSTK parameters gradient)+        chunk = take batchSize xs+        gradf c = go c targetInit+        name =+          prefix+          ++ unwords+               [ "v0 m" ++ show (widthSTK $ knownSTK @(XParams2 r Float))+               , "=" ++ show (tsize knownSTK targetInit) ]+    bench name $ nf gradf chunk++mnistBGroup2VTOZ :: Int -> Benchmark+mnistBGroup2VTOZ chunkLength =+  let (!targetInit, !art) =+        MnistFcnnRanked2.mnistTrainBench2VTOGradient+          @Double (Proxy @Float) IgnoreIncomingCotangent+          1 (mkStdGen 44) 1500 500+  in env (do+    testData0 <- loadMnistData testGlyphsPath testLabelsPath  -- 10k total+    let testData = shuffle (mkStdGen 42) testData0+    return $! map mkMnistDataLinearR $ take chunkLength testData) $+  \ xs ->+   bgroup ("2-hidden-layer rank 2 VTOZ runtime MNIST nn with samples: "+           ++ show chunkLength)+     [ mnistTrainBench2VTO "1500|500 " 0.02 chunkLength xs (targetInit, art)+     ]++-- The same as above, but without any simplification, even the smart+-- constructors.+mnistTrainBench2VTOX+  :: forall r. r ~ Double+  => String+  -> Double -> Int -> [MnistDataLinearR r]+  -> ( Concrete (XParams2 r Float)+     , AstArtifactRev+         (TKProduct+            (XParams2 r Float)+            (TKProduct (TKR2 1 (TKScalar Double))+                       (TKR2 1 (TKScalar Double))))+         (TKScalar r) )+  -> Benchmark+mnistTrainBench2VTOX prefix gamma batchSize xs (targetInit, art) = do+    let go :: [MnistDataLinearR r] -> Concrete (XParams2 r Float)+           -> Concrete (XParams2 r Float)+        go [] parameters = parameters+        go ((glyph, label) : rest) !parameters =+          let parametersAndInput =+                tpair parameters (tpair (rconcrete glyph) (rconcrete label))+              gradient = tproject1 $ fst+                         $ revInterpretArtifact art parametersAndInput Nothing+          in go rest (updateWithGradient gamma knownSTK parameters gradient)+        chunk = take batchSize xs+        gradf c = go c targetInit+        name =+          prefix+          ++ unwords+               [ "v0 m" ++ show (widthSTK $ knownSTK @(XParams2 r Float))+               , "=" ++ show (tsize knownSTK targetInit) ]+    bench name $ nf gradf chunk++mnistBGroup2VTOX :: Int -> Benchmark+mnistBGroup2VTOX chunkLength =+  let (!targetInit, !art) =+        MnistFcnnRanked2.mnistTrainBench2VTOGradientX+          @Double (Proxy @Float) IgnoreIncomingCotangent+          1 (mkStdGen 44) 1500 500+  in env (do+    testData0 <- loadMnistData testGlyphsPath testLabelsPath  -- 10k total+    let testData = shuffle (mkStdGen 42) testData0+    return $! map mkMnistDataLinearR $ take chunkLength testData) $+  \ xs ->+   bgroup ("2-hidden-layer rank 2 VTOX runtime MNIST nn with samples: "+           ++ show chunkLength)+     [ mnistTrainBench2VTO "1500|500 " 0.02 chunkLength xs (targetInit, art)+     ]++{- TODO: re-enable once -fpolymorphic-specialisation works++-- This is expected to fail with -O0 and to pass with -O1+-- and -fpolymorphic-specialisation.+-- This prevents running benchmarks without optimization, which is a good thing.+--+-- The `Storable` is only needed for overloaded profiling, e.g., with+-- cabal bench longMnistBench -ftest_seq -w /home/mikolaj/r/ghc.HEAD/ghc/_build/stage1/bin/ghc --allow-newer --enable-optimization --enable-profiling --profiling-detail=none --ghc-options="-fprof-late-overloaded -fpolymorphic-specialisation" --benchmark-options='-n1 -m pattern "1 VTA MNIST nn with samples: 400/500|150 v" +RTS -pj'+inspect $ hasNoTypeClassesExcept 'mnistTrainBench1VTA [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Storable]+inspect $ hasNoTypeClassesExcept 'mnistTrainBench1VTO [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Storable,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Elt, ''Integral]+inspect $ hasNoTypeClassesExcept 'mnistTrainBench2VTA [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Storable]+inspect $ hasNoTypeClassesExcept 'mnistTrainBench2VTC [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default]+inspect $ hasNoTypeClassesExcept 'mnistTrainBench2VTO [''(~), ''GoodScalar, ''Show, ''Num, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default,       ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''WithDict, ''KnownShS, ''KnownSTK, ''KnownNat, ''Nested.Elt, ''Integral]+-- inspect $ coreOf 'mnistTrainBench1VTA+-}
+ bench/common/BenchProdTools.hs view
@@ -0,0 +1,268 @@+{-# LANGUAGE TemplateHaskell #-}+{-# OPTIONS_GHC -Wno-missing-export-lists #-}+-- {-# OPTIONS_GHC -ddump-stranal #-}+-- | A contrived benchmark: a product of a list of scalars.+module BenchProdTools where++import Prelude++import Control.DeepSeq (NFData (..))+import Criterion.Main+import Data.Default qualified as Default+import Data.Foldable qualified as Foldable+import Data.List (foldl1')+import Data.Proxy (Proxy (Proxy))+import GHC.Exts (IsList (..), WithDict)+import GHC.TypeLits (KnownNat)+import Test.Inspection+import Type.Reflection (Typeable)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.Ops++bgroup100, bgroup1000, bgroup1e4, bgroup1e5, bgroup1e6, bgroup1e7, bgroup5e7 :: [Double] -> Benchmark+bgroup100 = envProd 100 $ \args -> bgroup "100" $ benchProd args++bgroup1000 = envProd 1000 $ \args -> bgroup "1000" $ benchProd args++bgroup1e4 = envProd 1e4 $ \args -> bgroup "1e4" $ benchProd args++bgroup1e5 = envProd 1e5 $ \args -> bgroup "1e5" $ benchProd args++bgroup1e6 = envProd 1e6 $ \args -> bgroup "1e6" $ benchProd args++bgroup1e7 = envProd 1e7 $ \args -> bgroup "1e7" $ benchProd args++bgroup5e7 = envProd 5e7 $ \args -> bgroup "5e7" $ benchProd args+  -- 5e7 == 5 * 10^7 == 0.5 * 10^8 == 0.5e8++envProd :: r ~ Double+        => Rational+        -> (forall n.+            ( SNat n+            , [Concrete (TKScalar r)]+            , ListR n (Concrete (TKScalar r))+            , ListR n (Concrete (TKS '[] r))+            , Concrete (TKS '[n] r) )+            -> Benchmark)+        -> [r]+        -> Benchmark+envProd rat f allxs =+  let k = round rat+  in withSNat k $ \(SNat @k) ->+    env (return $!+      let l = take k allxs+          lt = map sscalar l+      in ( SNat @k+         , map Concrete l+         , fromList (map Concrete l)+         , fromList lt+         , sfromList . fromList $ lt) )+        (f @k)++benchProd :: r ~ Double+          => ( SNat n+             , [Concrete (TKScalar r)]+             , ListR n (Concrete (TKScalar r))+             , ListR n (Concrete (TKS '[] r))+             , Concrete (TKS '[n] r) )+          -> [Benchmark]+benchProd ~(snat, list, l, lt, t) = case snat of+  SNat ->+    [ bench "cgrad s MapAccum" $ nf (crevSMapAccum snat) t+    , bench "grad s MapAccum" $ nf (revSMapAccum snat) t+    , bench "cgrad scalar MapAccum" $ nf (crevScalarMapAccum snat) t+    , bench "grad scalar MapAccum" $ nf (revScalarMapAccum snat) t+    , bench "cgrad scalar list" $ nf crevScalarList list+    , bench "grad scalar list" $ nf revScalarList list+    , bench "cgrad scalar L" $ nf (crevScalarL snat) l+    , bench "grad scalar L" $ nf (revScalarL snat) l+    , bench "cgrad scalar R" $ nf (crevScalarR snat) l+    , bench "grad scalar R" $ nf (revScalarR snat) l+    , bench "cgrad scalar NotShared" $ nf (crevScalarNotShared snat) l+    , bench "cgrad s L" $ nf (crevSL snat) lt+    , bench "grad s L" $ nf (revSL snat) lt+    , bench "cgrad s R" $ nf (crevSR snat) lt+    , bench "grad s R" $ nf (revSR snat) lt+    , bench "cgrad s NotShared" $ nf (crevSNotShared snat) lt+    ]++-- Another variant, with foldl1' and indexing, would be a disaster.+-- We can define sproduct if this benchmark ends up used anywhere,+-- because the current codomain of gradientFromDelta rules out+-- low-level hacky pipeline tricks that could avoid indexing.+multSMapAccum :: (BaseTensor target, LetTensor target, GoodScalar r)+              => SNat n -> target (TKS '[n] r) -> target (TKS '[] r)+multSMapAccum SNat = sfold (*) (sscalar 1)+{-# SPECIALIZE multSMapAccum :: SNat n -> ADVal Concrete (TKS '[n] Double) -> ADVal Concrete (TKS '[] Double) #-}+{-# SPECIALIZE multSMapAccum :: SNat n -> AstTensor AstMethodLet FullSpan (TKS '[n] Double) -> AstTensor AstMethodLet FullSpan (TKS '[] Double) #-}++crevSMapAccum+  :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)+crevSMapAccum snat@SNat = cgrad (kfromS . multSMapAccum snat)++revSMapAccum+  :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)+revSMapAccum snat@SNat = grad (kfromS . multSMapAccum snat)++multScalarMapAccum :: forall target n r.+                      (BaseTensor target, GoodScalar r)+                   => SNat n -> target (TKS '[n] r) -> target (TKScalar r)+multScalarMapAccum snat@SNat  =+  tproject1+  . tmapAccumL (Proxy @target)+     snat+     (FTKScalar @r)+     (FTKScalar @Z1)+     (FTKScalar @r)+     (let g :: forall f. ADReady f+            => f (TKScalar r) -> f (TKScalar r)+            -> f (TKProduct (TKScalar r) TKUnit)+          g !acc !e = tpair (acc * e) tunit+      in g)+     1+{-# SPECIALIZE multScalarMapAccum :: SNat n -> ADVal Concrete (TKS '[n] Double) -> ADVal Concrete (TKScalar Double) #-}+{-# SPECIALIZE multScalarMapAccum :: SNat n -> AstTensor AstMethodLet FullSpan (TKS '[n] Double) -> AstTensor AstMethodLet FullSpan (TKScalar Double) #-}++crevScalarMapAccum+  :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)+crevScalarMapAccum snat@SNat = cgrad (multScalarMapAccum snat)++revScalarMapAccum+  :: SNat n -> Concrete (TKS '[n] Double) -> Concrete (TKS '[n] Double)+revScalarMapAccum snat@SNat = grad (multScalarMapAccum snat)++multScalarList :: (BaseTensor target, GoodScalar r)+               => [target (TKScalar r)] -> target (TKScalar r)+multScalarList = foldl1' (*)++crevScalarList+  :: [Concrete (TKScalar Double)] -> [Concrete (TKScalar Double)]+crevScalarList =+  cgrad multScalarList++revScalarList+  :: [Concrete (TKScalar Double)] -> [Concrete (TKScalar Double)]+revScalarList =+  grad multScalarList++multScalarL :: (BaseTensor target, GoodScalar r)+            => ListR n (target (TKScalar r)) -> target (TKScalar r)+multScalarL = foldl1' (*) . Foldable.toList++crevScalarL+  :: SNat n -> ListR n (Concrete (TKScalar Double))+  -> ListR n (Concrete (TKScalar Double))+crevScalarL snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $+  cgrad multScalarL++revScalarL+  :: SNat n -> ListR n (Concrete (TKScalar Double))+  -> ListR n (Concrete (TKScalar Double))+revScalarL snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $+  grad multScalarL++multScalarR :: (BaseTensor target, GoodScalar r)+            => ListR n (target (TKScalar r)) -> target (TKScalar r)+multScalarR = foldr1 (*)++crevScalarR+  :: SNat n -> ListR n (Concrete (TKScalar Double))+  -> ListR n (Concrete (TKScalar Double))+crevScalarR snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $+  cgrad multScalarR++revScalarR+  :: SNat n -> ListR n (Concrete (TKScalar Double))+  -> ListR n (Concrete (TKScalar Double))+revScalarR snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $+  grad multScalarR++multScalarNotShared :: (BaseTensor target, GoodScalar r)+                    => ListR n (ADVal target (TKScalar r))+                    -> ADVal target (TKScalar r)+multScalarNotShared = foldr1 multNotShared++crevScalarNotShared+  :: SNat n -> ListR n (Concrete (TKScalar Double))+  -> ListR n (Concrete (TKScalar Double))+crevScalarNotShared snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKScalar Double)) snat) $+  cgrad multScalarNotShared++multSL :: (BaseTensor target, GoodScalar r)+       => ListR n (target (TKS '[] r)) -> target (TKS '[] r)+multSL = foldl1' (*) . Foldable.toList++crevSL+  :: SNat n -> ListR n (Concrete (TKS '[] Double))+  -> ListR n (Concrete (TKS '[] Double))+crevSL snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $+  cgrad (kfromS . multSL)++revSL+  :: SNat n -> ListR n (Concrete (TKS '[] Double))+  -> ListR n (Concrete (TKS '[] Double))+revSL snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $+  grad (kfromS . multSL)++multSR :: (BaseTensor target, GoodScalar r)+       => ListR n (target (TKS '[] r)) -> target (TKS '[] r)+multSR = foldr1 (*)++crevSR+  :: SNat n -> ListR n (Concrete (TKS '[] Double))+  -> ListR n (Concrete (TKS '[] Double))+crevSR snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $+  cgrad (kfromS . multSR)++revSR+  :: SNat n -> ListR n (Concrete (TKS '[] Double))+  -> ListR n (Concrete (TKS '[] Double))+revSR snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $+  grad (kfromS . multSR)++multSNotShared :: (BaseTensor target, GoodScalar r)+               => ListR n (ADVal target (TKS '[] r))+               -> ADVal target (TKS '[] r)+multSNotShared = foldr1 multNotShared++crevSNotShared+  :: SNat n -> ListR n (Concrete (TKS '[] Double))+  -> ListR n (Concrete (TKS '[] Double))+crevSNotShared snat@SNat =+  withKnownSTK (stkOfListR (knownSTK @(TKS '[] Double)) snat) $+  cgrad (kfromS . multSNotShared)++{- TODO: re-enable once -fpolymorphic-specialisation works++-- KnownNat and AstSpan are tag types, so it's fine not to specialize+-- for them. Some of the other classes come from existential types,+-- some of which it's not advantageous to specialize.+--+-- This is expected to fail with -O0 and to pass with -O1+-- and -fpolymorphic-specialisation.+-- This prevents running benchmarks without optimization, which is a good thing.+inspect $ hasNoTypeClassesExcept 'crevScalarL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt]+inspect $ hasNoTypeClassesExcept 'revScalarL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]+inspect $ hasNoTypeClassesExcept 'crevScalarNotShared [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt]+inspect $ hasNoTypeClassesExcept 'crevSL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''Nested.Storable,    ''ShareTensor]+inspect $ hasNoTypeClassesExcept 'revSL [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]+inspect $ hasNoTypeClassesExcept 'crevSMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor]+inspect $ hasNoTypeClassesExcept 'revSMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]+inspect $ hasNoTypeClassesExcept 'crevScalarMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor]+inspect $ hasNoTypeClassesExcept 'revScalarMapAccum [''(~), ''KnownNat, ''WithDict, ''KnownShS, ''AdaptableTarget, ''RandomValue, ''KnownSTK, ''GoodScalar, ''Num, ''Show, ''Ord, ''Eq, ''Nested.PrimElt, ''Nested.KnownElt, ''Nested.NumElt, ''Typeable, ''IfDifferentiable, ''NFData, ''Default.Default, ''Nested.Elt, ''LetTensor, ''BaseTensor, ''ConvertTensor, ''Boolean, ''CommonTargetEqOrd, ''AllTargetShow, ''ShareTensor,      ''AstSpan, ''RealFloatH, ''Nested.FloatElt, ''Fractional, ''Floating, ''IntegralH, ''RealFrac, ''Real, ''Nested.Storable, ''Integral]+-- inspect $ coreOf 'revScalarL+-}
+ example/MnistCnnRanked2.hs view
@@ -0,0 +1,149 @@+{-# LANGUAGE OverloadedLists #-}+{-# OPTIONS_GHC -Wno-missing-export-lists #-}+-- | Ranked tensor-based implementation of Convolutional Neural Network+-- for classification of MNIST digits. Sports 2 hidden layers.+--+-- With the current CPU backend it's slow enough that it's hard to see+-- if it trains.+module MnistCnnRanked2 where++import Prelude++import Data.Vector.Generic qualified as V+import Data.Vector.Storable (Vector)+import GHC.TypeLits (type (*), type (+), type Div)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape++import HordeAd+import MnistData++-- | The differentiable type of all trainable parameters of this nn.+-- Shaped version, statically checking all dimension widths.+--+-- Due to subtraction complicating posititive number type inference,+-- @kh@ denotes kernel height minus one and analogously @kw@ is kernel+-- width minus one.+type ADCnnMnistParametersShaped+       (target :: Target) h w kh kw c_out n_hidden r =+  ( ( target (TKS '[c_out, 1, kh + 1, kw + 1] r)+    , target (TKS '[c_out] r) )+  , ( target (TKS '[c_out, c_out, kh + 1, kw + 1] r)+    , target (TKS '[c_out] r) )+  , ( target (TKS '[n_hidden, c_out * (h `Div` 4) * (w `Div` 4)] r)+    , target (TKS '[n_hidden] r) )+  , ( target (TKS '[SizeMnistLabel, n_hidden] r)+    , target (TKS '[SizeMnistLabel] r) )+  )++-- | The differentiable type of all trainable parameters of this nn.+-- Ranked version.+type ADCnnMnistParameters (target :: Target) r =+  ( ( target (TKR 4 r)+    , target (TKR 1 r) )+  , ( target (TKR 4 r)+    , target (TKR 1 r) )+  , ( target (TKR 2 r)+    , target (TKR 1 r) )+  , ( target (TKR 2 r)+    , target (TKR 1 r) ) )++-- | A single convolutional layer with @relu@ and @maxPool@.+convMnistLayerR+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r)  -- ^ @[c_out, c_in, kh + 1, kw + 1]@+  -> target (TKR 4 r)  -- ^ @[batch_size, c_in, h, w]@+  -> target (TKR 1 r)  -- ^ @[c_out]@+  -> target (TKR 4 r)  -- ^ @[batch_size, c_out, h \`Div\` 2, w \`Div\` 2]@+convMnistLayerR ker input bias =+  let (batch_size :$: _ :$: h :$: w :$: ZSR) = rshape input+      yConv = conv2dUnpadded ker input+      biasStretched = rtranspose [0, 3, 1, 2]+                      $ rreplicate batch_size $ rreplicate h $ rreplicate w bias+      yRelu = relu $ yConv + biasStretched+  in maxPool2dUnpadded 2 2 yRelu++-- | Composition of two convolutional layers.+convMnistTwoR+  :: (ADReady target, GoodScalar r, Differentiable r)+  => Int -> Int -> Int+  -> PrimalOf target (TKR 4 r)+       -- ^ input images @[batch_size, 1, SizeMnistHeight, SizeMnistWidth]@+  -> ADCnnMnistParameters target r  -- ^ parameters+  -> target (TKR 2 r)  -- ^ output classification @[SizeMnistLabel, batch_size]@+convMnistTwoR sizeMnistHeightI sizeMnistWidthI batch_size input+              ( (ker1, bias1), (ker2, bias2)+              , (weightsDense, biasesDense), (weightsReadout, biasesReadout) ) =+  let t1 = convMnistLayerR ker1 (rfromPrimal input) bias1+      t2 = convMnistLayerR ker2 t1 bias2+             -- [ batch_size, c_out+             -- , SizeMnistHeight `Div` 4, SizeMnistWidth `Div` 2 ]+      c_out = rwidth bias1+      m1 = rreshape (batch_size+                     :$: c_out * (sizeMnistHeightI `div` 4)+                               * (sizeMnistWidthI `div` 4)+                     :$: ZSR)+                    t2+      m2 = rtr m1+      denseLayer = weightsDense `rmatmul2` m2+                   + rtr (rreplicate batch_size biasesDense)+      denseRelu = relu denseLayer+  in weightsReadout `rmatmul2` denseRelu+     + rtr (rreplicate batch_size biasesReadout)++-- | The neural network composed with the SoftMax-CrossEntropy loss function.+convMnistLossFusedR+  :: (ADReady target, ADReady (PrimalOf target), GoodScalar r, Differentiable r)+  => Int  -- ^ batch_size+  -> ( PrimalOf target (TKR 3 r)+         -- ^ @[batch_size, SizeMnistHeight, SizeMnistWidth]@+     , PrimalOf target (TKR 2 r) )  -- ^ @[batch_size, SizeMnistLabel]@+  -> ADCnnMnistParameters target r  -- kh kw c_out n_hidden+  -> target (TKScalar r)+convMnistLossFusedR batch_size (glyphR, labelR) adparameters =+  let input = rreshape (batch_size+                        :$: 1+                        :$: sizeMnistHeightInt+                        :$: sizeMnistWidthInt+                        :$: ZSR)+                       glyphR+      result = convMnistTwoR sizeMnistHeightInt sizeMnistWidthInt+                             batch_size input adparameters+      targets = rtr labelR+      loss = lossSoftMaxCrossEntropyR targets result+  in kfromPrimal (recip $ kconcrete $ fromIntegral batch_size) * loss++-- | A function testing the neural network given testing set of inputs+-- and the trained parameters.+convMnistTestR+  :: forall target r.+     (target ~ Concrete, GoodScalar r, Differentiable r)+  => Int+  -> MnistDataBatchR r+  -> ADCnnMnistParameters Concrete r+  -> r+convMnistTestR 0 _ _ = 0+convMnistTestR batch_size (glyphR, labelR) testParams =+  let input :: target (TKR 4 r)+      input =+        rconcrete $ Nested.rreshape [ batch_size, 1+                                    , sizeMnistHeightInt, sizeMnistWidthInt ]+                                    glyphR+      outputR :: Concrete (TKR 2 r)+      outputR =+        let nn :: ADCnnMnistParameters target r+               -> target (TKR 2 r)  -- [SizeMnistLabel, batch_size]+            nn = convMnistTwoR sizeMnistHeightInt sizeMnistWidthInt+                               batch_size input+        in nn testParams+      outputs = map rtoVector $ runravelToList+                $ rtranspose [1, 0] outputR+      labels = map rtoVector+               $ runravelToList @_ @(TKScalar r)+               $ rconcrete labelR+      matchesLabels :: Vector r -> Vector r -> Int+      matchesLabels output label | V.maxIndex output == V.maxIndex label = 1+                                 | otherwise = 0+  in fromIntegral (sum (zipWith matchesLabels outputs labels))+     / fromIntegral batch_size
+ example/MnistCnnShaped2.hs view
@@ -0,0 +1,167 @@+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+{-# OPTIONS_GHC -Wno-missing-export-lists #-}+-- | Shaped tensor-based implementation of Convolutional Neural Network+-- for classification of MNIST digits. Sports 2 hidden layers.+--+-- With the current CPU backend it's slow enough that it's hard to see+-- if it trains.+module MnistCnnShaped2 where++import Prelude++import Data.Vector.Generic qualified as V+import Data.Vector.Storable (Vector)+import GHC.TypeLits (fromSNat, type (*), type (+), type (<=), type Div)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Shaped.Shape++import HordeAd+import MnistData++-- | The differentiable type of all trainable parameters of this nn.+-- Shaped version, statically checking all dimension widths.+--+-- Due to subtraction complicating posititive number type inference,+-- @kh@ denotes kernel height minus one and analogously @kw@ is kernel+-- width minus one.+type ADCnnMnistParametersShaped+       (target :: Target) h w kh kw c_out n_hidden r =+  ( ( target (TKS '[c_out, 1, kh + 1, kw + 1] r)+    , target (TKS '[c_out] r) )+  , ( target (TKS '[c_out, c_out, kh + 1, kw + 1] r)+    , target (TKS '[c_out] r) )+  , ( target (TKS '[n_hidden, c_out * (h `Div` 4) * (w `Div` 4) ] r)+    , target (TKS '[n_hidden] r) )+  , ( target (TKS '[SizeMnistLabel, n_hidden] r)+    , target (TKS '[SizeMnistLabel] r) )+  )++-- | A single convolutional layer with @relu@ and @maxPool@.+-- The @c_in@ type parameter is going to be alwayst 1, meaning grayscale,+-- but this function works for any @c_in@.+convMnistLayerS+  :: forall kh kw h w c_in c_out batch_size target r.+     ( 1 <= kh+     , 1 <= kw  -- wrongly reported as redundant due to plugins+     , ADReady target, GoodScalar r, Differentiable r )+  => SNat kh -> SNat kw -> SNat h -> SNat w+  -> SNat c_in -> SNat c_out -> SNat batch_size+  -> target (TKS '[c_out, c_in, kh + 1, kw + 1] r)+  -> target (TKS '[batch_size, c_in, h, w] r)+  -> target (TKS '[c_out] r)+  -> target (TKS '[batch_size, c_out, h `Div` 2, w `Div` 2] r)+convMnistLayerS SNat SNat SNat SNat SNat SNat SNat+                ker input bias =+  let yConv = conv2dUnpaddedS ker input+      biasStretched = stranspose @'[0, 3, 1, 2]+                      $ sreplicate {-@batch_size-}+                      $ sreplicate {-@h-}+                      $ sreplicate {-@w-} bias+      yRelu = reluS $ yConv + biasStretched+  in maxPool2dUnpaddedS @2 @2 yRelu++-- | Composition of two convolutional layers.+convMnistTwoS+  :: forall kh kw h w c_out n_hidden batch_size target r.+       -- @h@ and @w@ are fixed with MNIST data, but not with test data+     ( 1 <= kh  -- kernel height is large enough+     , 1 <= kw  -- kernel width is large enough+     , ADReady target, GoodScalar r, Differentiable r )+  => SNat kh -> SNat kw -> SNat h -> SNat w+  -> SNat c_out -> SNat n_hidden -> SNat batch_size+       -- ^ these boilerplate lines tie type parameters to the corresponding+       -- SNat value parameters denoting basic dimensions+  -> PrimalOf target (TKS '[batch_size, 1, h, w] r)  -- ^ input images+  -> ADCnnMnistParametersShaped target h w kh kw c_out n_hidden r+       -- ^ parameters+  -> target (TKS '[SizeMnistLabel, batch_size] r)  -- ^ output classification+convMnistTwoS kh@SNat kw@SNat h@SNat w@SNat+              c_out@SNat _n_hidden@SNat batch_size@SNat+              input+              ( (ker1, bias1), (ker2, bias2)+              , (weightsDense, biasesDense), (weightsReadout, biasesReadout) ) =+  assumeEquality @(Div (Div w 2) 2) @(Div w 4) $+  assumeEquality @(Div (Div h 2) 2) @(Div h 4) $+  let t1 = convMnistLayerS kh kw h w+                           (SNat @1) c_out batch_size+                           ker1 (sfromPrimal input) bias1+--      t2 :: target (TKS '[batch_size, c_out, h `Div` 4, w `Div` 4] r)+      t2 = convMnistLayerS kh kw (SNat @(h `Div` 2)) (SNat @(w `Div` 2))+                           c_out c_out batch_size+                           ker2 t1 bias2+--      m1 :: target (TKS '[batch_size, c_out * (h `Div` 4) * (w `Div` 4)] r)+      m1 = sreshape t2+      denseLayer = weightsDense `smatmul2` str m1+                   + str (sreplicate biasesDense)+  in weightsReadout `smatmul2` reluS denseLayer+     + str (sreplicate biasesReadout)++-- | The neural network composed with the SoftMax-CrossEntropy loss function.+convMnistLossFusedS+  :: forall kh kw h w c_out n_hidden batch_size target r.+     ( h ~ SizeMnistHeight, w ~ SizeMnistWidth+     , 1 <= kh+     , 1 <= kw+     , ADReady target, ADReady (PrimalOf target)+     , GoodScalar r, Differentiable r )+  => SNat kh -> SNat kw+  -> SNat c_out+  -> SNat n_hidden -> SNat batch_size+  -> ( PrimalOf target (TKS '[batch_size, h, w] r)+     , PrimalOf target (TKS '[batch_size, SizeMnistLabel] r) )+  -> ADCnnMnistParametersShaped target h w kh kw c_out n_hidden r+  -> target (TKScalar r)+convMnistLossFusedS kh@SNat kw@SNat+                    c_out@SNat n_hidden@SNat batch_size@SNat+                    (glyphS, labelS) adparameters =+  let input :: PrimalOf target (TKS '[batch_size, 1, h, w] r)+      input = sreshape glyphS+      result = convMnistTwoS kh kw (SNat @h) (SNat @w)+                             c_out n_hidden batch_size+                             input adparameters+      targets = str labelS+      loss = lossSoftMaxCrossEntropyS targets result+  in kfromPrimal (recip $ kconcrete $ fromInteger $ fromSNat batch_size) * loss++-- | A function testing the neural network given testing set of inputs+-- and the trained parameters.+convMnistTestS+  :: forall kh kw h w c_out n_hidden batch_size target r.+     ( h ~ SizeMnistHeight, w ~ SizeMnistWidth+     , 1 <= kh+     , 1 <= kw+     , target ~ Concrete+     , GoodScalar r, Differentiable r )+  => SNat kh -> SNat kw+  -> SNat c_out+  -> SNat n_hidden -> SNat batch_size+  -> MnistDataBatchS batch_size r+  -> ADCnnMnistParametersShaped target h w kh kw c_out n_hidden r+  -> r+convMnistTestS _ _ _ _ batch_size@SNat _ _+  | sNatValue batch_size == 0 = 0+convMnistTestS kh@SNat kw@SNat+               c_out@SNat n_hidden@SNat batch_size@SNat+               (glyphS, labelS) testParams =+  let input :: target (TKS '[batch_size, 1, h, w] r)+      input = sconcrete $ Nested.sreshape knownShS glyphS+      outputS :: Concrete (TKS '[SizeMnistLabel, batch_size] r)+      outputS =+        let nn :: ADCnnMnistParametersShaped target h w kh kw c_out n_hidden r+               -> target (TKS '[SizeMnistLabel, batch_size] r)+            nn = convMnistTwoS kh kw (SNat @h) (SNat @w)+                               c_out n_hidden batch_size+                               input+        in nn testParams+      outputs = map stoVector $ sunravelToList+                $ stranspose @'[1, 0] outputS+      labels = map stoVector+               $ sunravelToList @_ @_ @(TKScalar r)+               $ sconcrete labelS+      matchesLabels :: Vector r -> Vector r -> Int+      matchesLabels output label | V.maxIndex output == V.maxIndex label = 1+                                 | otherwise = 0+  in fromIntegral (sum (zipWith matchesLabels outputs labels))+     / fromInteger (fromSNat batch_size)
+ example/MnistData.hs view
@@ -0,0 +1,185 @@+{-# OPTIONS_GHC -Wno-missing-export-lists #-}+-- | Parsing and pre-processing of MNIST data.+module MnistData where++import Prelude++import Codec.Compression.GZip (decompress)+import Data.ByteString.Lazy qualified as LBS+import Data.IDX+import Data.List (sortBy)+import Data.List.NonEmpty qualified as NonEmpty+import Data.Maybe (fromMaybe)+import Data.Ord (comparing)+import Data.Vector.Generic qualified as V+import Data.Vector.Storable (Vector)+import Data.Vector.Storable qualified as VS+import Data.Vector.Unboxed qualified+import GHC.TypeLits (KnownNat, Nat, type (*))+import System.IO (IOMode (ReadMode), withBinaryFile)+import System.Random++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape++import HordeAd++type SizeMnistWidth = 28 :: Nat++sizeMnistWidth :: SNat SizeMnistWidth+sizeMnistWidth = SNat @SizeMnistWidth++sizeMnistWidthInt :: Int+sizeMnistWidthInt = sNatValue sizeMnistWidth++type SizeMnistHeight = SizeMnistWidth++sizeMnistHeight :: SNat SizeMnistHeight+sizeMnistHeight = SNat @SizeMnistHeight++sizeMnistHeightInt :: Int+sizeMnistHeightInt = valueOf @SizeMnistHeight++type SizeMnistGlyph = SizeMnistWidth * SizeMnistHeight++sizeMnistGlyphInt :: Int+sizeMnistGlyphInt = valueOf @SizeMnistGlyph++type SizeMnistLabel = 10 :: Nat++sizeMnistLabel :: SNat SizeMnistLabel+sizeMnistLabel = SNat @SizeMnistLabel++sizeMnistLabelInt :: Int+sizeMnistLabelInt = sNatValue sizeMnistLabel++type LengthTestData = 10000 :: Nat++-- Actually, a better representation, supported by @Data.IDX@,+-- is an integer label and a picture (the same vector as below).+-- Then we'd use @lossCrossEntropy@ that picks a component according+-- to the label instead of performing a dot product with scaling.+-- This results in much smaller Delta expressions.+-- Our library makes this easy to express and gradients compute fine.+-- OTOH, methods with only matrix operations and graphs can't handle that.+-- However, the goal of the exercise it to implement the same+-- neural net that backprop uses for comparative benchmarks.+-- Also, loss computation is not the bottleneck and the more general+-- mechanism that admits non-discrete target labels fuses nicely+-- with softMax. This also seems to be the standard or at least+-- a simple default in tutorial.+type MnistData r = (Vector r, Vector r)++type MnistDataLinearR r =+  ( Nested.Ranked 1 r+  , Nested.Ranked 1 r )++type MnistDataR r =+  ( Nested.Ranked 2 r+  , Nested.Ranked 1 r )++type MnistDataBatchR r =+  ( Nested.Ranked 3 r  -- [batch_size, SizeMnistHeight, SizeMnistWidth]+  , Nested.Ranked 2 r )  -- [batch_size, SizeMnistLabel]++type MnistDataS r =+  ( Nested.Shaped '[SizeMnistHeight, SizeMnistWidth] r+  , Nested.Shaped '[SizeMnistLabel] r )++type MnistDataBatchS batch_size r =+  ( Nested.Shaped '[batch_size, SizeMnistHeight, SizeMnistWidth] r+  , Nested.Shaped '[batch_size, SizeMnistLabel] r )++mkMnistDataLinearR :: Nested.PrimElt r+                   => MnistData r -> MnistDataLinearR r+mkMnistDataLinearR (input, target) =+  ( Nested.rfromVector+      (sizeMnistGlyphInt :$: ZSR) input+  , Nested.rfromVector+      (sizeMnistLabelInt :$: ZSR) target )++mkMnistDataR :: Nested.PrimElt r+             => MnistData r -> MnistDataR r+mkMnistDataR (input, target) =+  ( Nested.rfromVector+      (sizeMnistHeightInt :$: sizeMnistWidthInt :$: ZSR) input+  , Nested.rfromVector+      (sizeMnistLabelInt :$: ZSR) target )++mkMnistDataBatchR :: Nested.Elt r+                  => [MnistDataR r] -> MnistDataBatchR r+mkMnistDataBatchR l =+  let (inputs, targets) = unzip l+  in ( Nested.rfromListOuter $ NonEmpty.fromList inputs+     , Nested.rfromListOuter $ NonEmpty.fromList targets )++mkMnistDataS :: Nested.PrimElt r+             => MnistData r -> MnistDataS r+mkMnistDataS (input, target) =+  (Nested.sfromVector knownShS input, Nested.sfromVector knownShS target)+{-# SPECIALIZE mkMnistDataS :: MnistData Double -> MnistDataS Double #-}+{-# SPECIALIZE mkMnistDataS :: MnistData Float -> MnistDataS Float #-}++mkMnistDataBatchS :: forall batch_size r. (Nested.Elt r, KnownNat batch_size)+                  => [MnistDataS r] -> MnistDataBatchS batch_size r+mkMnistDataBatchS l =+  let (inputs, targets) = unzip l+  in ( Nested.sfromListOuter (SNat @batch_size)+       $ NonEmpty.fromList inputs+     , Nested.sfromListOuter (SNat @batch_size)+       $ NonEmpty.fromList targets )+{-# SPECIALIZE mkMnistDataBatchS :: forall batch_size. KnownNat batch_size => [MnistDataS Double] -> MnistDataBatchS batch_size Double #-}+{-# SPECIALIZE mkMnistDataBatchS :: forall batch_size. KnownNat batch_size => [MnistDataS Float] -> MnistDataBatchS batch_size Float #-}++readMnistData :: forall r. (VS.Storable r, Fractional r)+              => LBS.ByteString -> LBS.ByteString -> [MnistData r]+readMnistData glyphsBS labelsBS =+  let glyphs = fromMaybe (error "wrong MNIST glyphs file")+               $ decodeIDX glyphsBS+      labels = fromMaybe (error "wrong MNIST labels file")+               $ decodeIDXLabels labelsBS+      intData = fromMaybe (error "can't decode MNIST file into integers")+                $ labeledIntData labels glyphs+      f :: (Int, Data.Vector.Unboxed.Vector Int) -> MnistData r+      -- Copied from library backprop to enable comparison of results.+      -- I have no idea how this is different from @labeledDoubleData@, etc.+      f (labN, v) =+        let !vGlyph = V.map (\r -> fromIntegral r / 255) $ V.convert v+            !vLabel = V.generate sizeMnistLabelInt+                                 (\i -> if i == labN then 1 else 0)+        in (vGlyph, vLabel)+  in map f intData+{-# SPECIALIZE readMnistData :: LBS.ByteString -> LBS.ByteString -> [MnistData Double] #-}+{-# SPECIALIZE readMnistData :: LBS.ByteString -> LBS.ByteString -> [MnistData Float] #-}++trainGlyphsPath, trainLabelsPath, testGlyphsPath, testLabelsPath :: FilePath+trainGlyphsPath = "samplesData/train-images-idx3-ubyte.gz"+trainLabelsPath = "samplesData/train-labels-idx1-ubyte.gz"+testGlyphsPath  = "samplesData/t10k-images-idx3-ubyte.gz"+testLabelsPath  = "samplesData/t10k-labels-idx1-ubyte.gz"++loadMnistData :: (VS.Storable r, Fractional r)+              => FilePath -> FilePath -> IO [MnistData r]+loadMnistData glyphsPath labelsPath =+  withBinaryFile glyphsPath ReadMode $ \glyphsHandle ->+    withBinaryFile labelsPath ReadMode $ \labelsHandle -> do+      glyphsContents <- LBS.hGetContents glyphsHandle+      labelsContents <- LBS.hGetContents labelsHandle+      return $! readMnistData (decompress glyphsContents)+                              (decompress labelsContents)+{-# SPECIALIZE loadMnistData :: FilePath -> FilePath -> IO [MnistData Double] #-}+{-# SPECIALIZE loadMnistData :: FilePath -> FilePath -> IO [MnistData Float] #-}++-- Good enough for QuickCheck, so good enough for me.+shuffle :: StdGen -> [a] -> [a]+shuffle g l =+  let rnds = randoms g :: [Int]+      res = map fst $ sortBy (comparing snd) $ zip l rnds+  in foldr seq () res `seq` res++chunksOf :: Int -> [e] -> [[e]]+chunksOf n = go where+  go [] = []+  go l = let (chunk, rest) = splitAt n l+         in chunk : go rest
+ example/MnistFcnnRanked1.hs view
@@ -0,0 +1,110 @@+{-# OPTIONS_GHC -Wno-missing-export-lists #-}+-- | Implementation of fully connected neutral network for classification+-- of MNIST digits with sized lists of rank 1 tensors (vectors)+-- as the trainable parameters. Sports 2 hidden layers. No mini-batches.+-- This is an exotic and fundamentally inefficient way of implementing nns,+-- but it's valuable for comparative benchmarking.+module MnistFcnnRanked1 where++import Prelude++import Data.Vector.Generic qualified as V+import GHC.TypeLits (KnownNat, Nat)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape++import HordeAd+import HordeAd.Core.Ops (tfromListR)+import MnistData++-- | The differentiable type of all trainable parameters of this nn.+type ADFcnnMnist1Parameters+       (target :: Target) (widthHidden :: Nat) (widthHidden2 :: Nat) r =+  ( ( ListR widthHidden (target (TKS '[SizeMnistGlyph] r))+    , target (TKS '[widthHidden] r) )+  , ( ListR widthHidden2 (target (TKS '[widthHidden] Float))+    , target (TKS '[widthHidden2] r) )+  , ( ListR SizeMnistLabel (target (TKS '[widthHidden2] r))+    , target (TKS '[SizeMnistLabel] r) )+  )++-- | An ad-hoc matrix multiplication analogue for matrices represented+-- as lists of vectors.+listMatmul1+  :: forall target r w1 w2.+     (ADReady target, GoodScalar r, KnownNat w1)+  => target (TKS '[w1] r) -> ListR w2 (target (TKS '[w1] r))+  -> target (TKS '[w2] r)+{-# INLINE listMatmul1 #-}  -- this doesn't want to specialize+listMatmul1 x0 weights = tlet x0 $ \x ->+  let f :: target (TKS '[w1] r) -> target (TKS '[] r)+      f v = v `sdot0` x+  in tfromListR knownSTK $ f <$> weights++-- | Fully connected neural network for the MNIST digit classification task.+-- There are two hidden layers and both use the same activation function.+-- The output layer uses a different activation function.+-- The widths of the two hidden layers are @widthHidden@ and @widthHidden2@,+-- respectively.+afcnnMnist1 :: forall target r widthHidden widthHidden2.+               (ADReady target, GoodScalar r, Differentiable r)+            => (forall n. KnownNat n+                => target (TKS '[n] r) -> target (TKS '[n] r))+            -> (target (TKS '[SizeMnistLabel] r)+                -> target (TKS '[SizeMnistLabel] r))+            -> SNat widthHidden -> SNat widthHidden2+            -> target (TKS '[SizeMnistGlyph] r)+            -> ADFcnnMnist1Parameters target widthHidden widthHidden2 r+            -> target (TKR 1 r)+afcnnMnist1 factivationHidden factivationOutput SNat SNat+            datum ((hidden, bias), (hidden2, bias2), (readout, biasr)) =+  let hiddenLayer1 = listMatmul1 datum hidden + bias+      nonlinearLayer1 = factivationHidden hiddenLayer1+      hiddenLayer2 = scast (listMatmul1 (scast nonlinearLayer1) hidden2) + bias2+      nonlinearLayer2 = factivationHidden hiddenLayer2+      outputLayer = listMatmul1 nonlinearLayer2 readout + biasr+      result :: target (TKS '[SizeMnistLabel] r)+      result = factivationOutput outputLayer+  in rfromS result++-- | The neural network applied to concrete activation functions+-- and composed with the appropriate loss function.+afcnnMnistLoss1+  :: (ADReady target, GoodScalar r, Differentiable r)+  => SNat widthHidden -> SNat widthHidden2+  -> (target (TKR 1 r), target (TKR 1 r))+  -> ADFcnnMnist1Parameters target widthHidden widthHidden2 r+  -> target (TKScalar r)+afcnnMnistLoss1 widthHidden widthHidden2 (datum, target) adparams =+  let result = afcnnMnist1 logisticS softMax1S+                           widthHidden widthHidden2 (sfromR datum) adparams+  in lossCrossEntropyV target result+-- {-# SPECIALIZE afcnnMnistLoss1 :: (GoodScalar r, Differentiable r) => SNat widthHidden -> SNat widthHidden2 -> (AstTensor AstMethodLet FullSpan (TKR 1 r), AstTensor AstMethodLet FullSpan (TKR 1 r)) -> ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 r -> AstTensor AstMethodLet FullSpan (TKScalar r) #-}+{-# SPECIALIZE afcnnMnistLoss1 :: SNat widthHidden -> SNat widthHidden2 -> (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 1 Double)) -> ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 Double -> AstTensor AstMethodLet FullSpan (TKScalar Double) #-}+{-# SPECIALIZE afcnnMnistLoss1 :: SNat widthHidden -> SNat widthHidden2 -> (AstTensor AstMethodLet FullSpan (TKR 1 Float), AstTensor AstMethodLet FullSpan (TKR 1 Float)) -> ADFcnnMnist1Parameters (AstTensor AstMethodLet FullSpan) widthHidden widthHidden2 Float -> AstTensor AstMethodLet FullSpan (TKScalar Float) #-}+{-# SPECIALIZE afcnnMnistLoss1 :: SNat widthHidden -> SNat widthHidden2 -> (ADVal Concrete (TKR 1 Double), ADVal Concrete (TKR 1 Double)) -> ADFcnnMnist1Parameters (ADVal Concrete) widthHidden widthHidden2 Double -> ADVal Concrete (TKScalar Double) #-}+{-# SPECIALIZE afcnnMnistLoss1 :: SNat widthHidden -> SNat widthHidden2 -> (ADVal Concrete (TKR 1 Float), ADVal Concrete (TKR 1 Float)) -> ADFcnnMnist1Parameters (ADVal Concrete) widthHidden widthHidden2 Float -> ADVal Concrete (TKScalar Float) #-}++-- | A function testing the neural network given testing set of inputs+-- and the trained parameters.+afcnnMnistTest1+  :: forall target widthHidden widthHidden2 r.+     (target ~ Concrete, GoodScalar r, Differentiable r)+  => SNat widthHidden -> SNat widthHidden2+  -> [MnistDataLinearR r]+  -> ADFcnnMnist1Parameters target widthHidden widthHidden2 r+  -> r+afcnnMnistTest1 _ _ [] _ = 0+afcnnMnistTest1 widthHidden widthHidden2 dataList testParams =+  let matchesLabels :: MnistDataLinearR r -> Bool+      matchesLabels (glyph, label) =+        let glyph1 = rconcrete glyph+            nn :: ADFcnnMnist1Parameters target widthHidden widthHidden2 r+               -> target (TKR 1 r)+            nn = afcnnMnist1 logisticS softMax1S+                             widthHidden widthHidden2 (sfromR glyph1)+            v = rtoVector $ nn testParams+        in V.maxIndex v == V.maxIndex (Nested.rtoVector label)+  in fromIntegral (length (filter matchesLabels dataList))+     / fromIntegral (length dataList)
+ example/MnistFcnnRanked2.hs view
@@ -0,0 +1,184 @@+{-# OPTIONS_GHC -Wno-missing-export-lists #-}+-- | Ranked tensor-based implementation of fully connected neutral network+-- for classification of MNIST digits. Sports 2 hidden layers. No mini-batches,+-- so the maximum rank of tensors being used is 2.+module MnistFcnnRanked2 where++import Prelude++import Data.Proxy (Proxy (Proxy))+import Data.Vector.Generic qualified as V+import GHC.Exts (inline)+import GHC.TypeLits (Nat)+import System.Random++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.CarriersAst+import MnistData++-- | The differentiable type of all trainable parameters of this nn.+-- Shaped version, statically checking all dimension widths.+type ADFcnnMnist2ParametersShaped+       (target :: Target) (widthHidden :: Nat) (widthHidden2 :: Nat) r q =+  ( ( target (TKS '[widthHidden, SizeMnistGlyph] r)+    , target (TKS '[widthHidden] r) )+  , ( target (TKS '[widthHidden2, widthHidden] q)+    , target (TKS '[widthHidden2] r) )+  , ( target (TKS '[SizeMnistLabel, widthHidden2] r)+    , target (TKS '[SizeMnistLabel] r) )+  )++-- | The differentiable type of all trainable parameters of this nn.+type ADFcnnMnist2Parameters (target :: Target) r q =+  ( ( target (TKR 2 r)+    , target (TKR 1 r) )+  , ( target (TKR 2 q)+    , target (TKR 1 r) )+  , ( target (TKR 2 r)+    , target (TKR 1 r) )+  )++type XParams2 r q = X (MnistFcnnRanked2.ADFcnnMnist2Parameters Concrete r q)++-- | Fully connected neural network for the MNIST digit classification task.+-- There are two hidden layers and both use the same activation function.+-- The output layer uses a different activation function.+-- The widths of the two hidden layers are @widthHidden@ and @widthHidden2@,+-- respectively.+afcnnMnist2 :: ( ADReady target, GoodScalar r, Differentiable r+               , GoodScalar q, Differentiable q )+            => (target (TKR 1 r) -> target (TKR 1 r))+            -> (target (TKR 1 r) -> target (TKR 1 r))+            -> target (TKR 1 r)+            -> ADFcnnMnist2Parameters target r q+            -> target (TKR 1 r)+afcnnMnist2 factivationHidden factivationOutput+            datum ((hidden, bias), (hidden2, bias2), (readout, biasr)) =+  let hiddenLayer1 = rmatvecmul hidden datum + bias+      nonlinearLayer1 = factivationHidden hiddenLayer1+      hiddenLayer2 = rcast (rmatvecmul hidden2 (rcast nonlinearLayer1)) + bias2+      nonlinearLayer2 = factivationHidden hiddenLayer2+      outputLayer = rmatvecmul readout nonlinearLayer2 + biasr+  in factivationOutput outputLayer++-- | The neural network applied to concrete activation functions+-- and composed with the appropriate loss function.+afcnnMnistLoss2+  :: ( ADReady target, GoodScalar r, Differentiable r+     , GoodScalar q, Differentiable q )+  => (target (TKR 1 r), target (TKR 1 r)) -> ADFcnnMnist2Parameters target r q+  -> target (TKScalar r)+afcnnMnistLoss2 (datum, target) adparams =+  let result = inline afcnnMnist2 logistic softMax1 datum adparams+  in lossCrossEntropyV target result+{-# SPECIALIZE afcnnMnistLoss2 :: (ADVal Concrete (TKR 1 Double), ADVal Concrete (TKR 1 Double)) -> ADFcnnMnist2Parameters (ADVal Concrete) Double Float -> ADVal Concrete (TKScalar Double) #-}+{-# SPECIALIZE afcnnMnistLoss2 :: (ADVal Concrete (TKR 1 Float), ADVal Concrete (TKR 1 Float)) -> ADFcnnMnist2Parameters (ADVal Concrete) Float Float -> ADVal Concrete (TKScalar Float) #-}+{-# SPECIALIZE afcnnMnistLoss2 :: (ADVal Concrete (TKR 1 Double), ADVal Concrete (TKR 1 Double)) -> ADFcnnMnist2Parameters (ADVal Concrete) Double Double -> ADVal Concrete (TKScalar Double) #-}++-- | A function testing the neural network given testing set of inputs+-- and the trained parameters.+afcnnMnistTest2+  :: forall target r q.+     ( target ~ Concrete, GoodScalar r, Differentiable r+     , GoodScalar q, Differentiable q )+  => [MnistDataLinearR r]+  -> ADFcnnMnist2Parameters target r q+  -> r+afcnnMnistTest2 [] _ = 0+afcnnMnistTest2 dataList testParams =+  let matchesLabels :: MnistDataLinearR r -> Bool+      matchesLabels (glyph, label) =+        let glyph1 = rconcrete glyph+            nn :: ADFcnnMnist2Parameters target r q+               -> target (TKR 1 r)+            nn = inline afcnnMnist2 logistic softMax1 glyph1+            v = rtoVector $ nn testParams+        in V.maxIndex v == V.maxIndex (Nested.rtoVector label)+  in fromIntegral (length (filter matchesLabels dataList))+     / fromIntegral (length dataList)++-- | The loss function applied to randomly generated initial parameters+-- and wrapped in artifact generation. This is helpful to share code+-- between tests and benchmarks and to separate compile-time and run-time+-- for benchmarking (this part is considered compile-time).+mnistTrainBench2VTOGradient+  :: forall r q. ( GoodScalar r, Differentiable r+                 , GoodScalar q, Differentiable q )+  => Proxy q -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int+  -> ( Concrete (XParams2 r q)+     , AstArtifactRev+         (TKProduct+            (XParams2 r q)+            (TKProduct (TKR2 1 (TKScalar r))+                       (TKR2 1 (TKScalar r))))+         (TKScalar r) )+mnistTrainBench2VTOGradient Proxy cotangentHandling range seed widthHidden widthHidden2 =+  withSNat widthHidden $ \(SNat @widthHidden) ->+  withSNat widthHidden2 $ \(SNat @widthHidden2) ->+  -- Initial parameter generation is counted as part of compilation time.+  let targetInit =+        forgetShape $ fst+        $ randomValue @(Concrete (X (MnistFcnnRanked2.ADFcnnMnist2ParametersShaped+                                   Concrete widthHidden widthHidden2 r q)))+                      range seed+      ftk = tftk @Concrete (knownSTK @(XParams2 r q)) targetInit+      ftkData = FTKProduct (FTKR (sizeMnistGlyphInt :$: ZSR) FTKScalar)+                           (FTKR (sizeMnistLabelInt :$: ZSR) FTKScalar)+      f :: ( MnistFcnnRanked2.ADFcnnMnist2Parameters+               (AstTensor AstMethodLet FullSpan) r q+           , ( AstTensor AstMethodLet FullSpan (TKR 1 r)+             , AstTensor AstMethodLet FullSpan (TKR 1 r) ) )+        -> AstTensor AstMethodLet FullSpan (TKScalar r)+      f (pars, (glyphR, labelR)) =+        afcnnMnistLoss2 (glyphR, labelR) pars+      artRaw = revArtifactAdapt cotangentHandling f (FTKProduct ftk ftkData)+  in (targetInit, artRaw)+{-# SPECIALIZE mnistTrainBench2VTOGradient :: Proxy Float -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int -> ( Concrete (XParams2 Double Float), AstArtifactRev (TKProduct (XParams2 Double Float) (TKProduct (TKR2 1 (TKScalar Double)) (TKR2 1 (TKScalar Double)))) (TKScalar Double) ) #-}+{-# SPECIALIZE mnistTrainBench2VTOGradient :: Proxy Float -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int -> ( Concrete (XParams2 Float Float), AstArtifactRev (TKProduct (XParams2 Float Float) (TKProduct (TKR2 1 (TKScalar Float)) (TKR2 1 (TKScalar Float)))) (TKScalar Float) ) #-}+{-# SPECIALIZE mnistTrainBench2VTOGradient :: Proxy Double -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int -> ( Concrete (XParams2 Double Double), AstArtifactRev (TKProduct (XParams2 Double Double) (TKProduct (TKR2 1 (TKScalar Double)) (TKR2 1 (TKScalar Double)))) (TKScalar Double) ) #-}++-- | A version of 'mnistTrainBench2VTOGradient' without any simplification,+-- even the AST smart constructors. Intended for benchmarking.+mnistTrainBench2VTOGradientX+  :: forall r q. ( GoodScalar r, Differentiable r+                 , GoodScalar q, Differentiable q )+  => Proxy q -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int+  -> ( Concrete (XParams2 r q)+     , AstArtifactRev+         (TKProduct+            (XParams2 r q)+            (TKProduct (TKR2 1 (TKScalar r))+                       (TKR2 1 (TKScalar r))))+         (TKScalar r) )+mnistTrainBench2VTOGradientX Proxy cotangentHandling range seed widthHidden widthHidden2 =+  withSNat widthHidden $ \(SNat @widthHidden) ->+  withSNat widthHidden2 $ \(SNat @widthHidden2) ->+  -- Initial parameter generation is counted as part of compilation time.+  let targetInit =+        forgetShape $ fst+        $ randomValue @(Concrete (X (MnistFcnnRanked2.ADFcnnMnist2ParametersShaped+                                   Concrete widthHidden widthHidden2 r q)))+                      range seed+      ftk = tftk @Concrete (knownSTK @(XParams2 r q)) targetInit+      ftkData = FTKProduct (FTKR (sizeMnistGlyphInt :$: ZSR) FTKScalar)+                           (FTKR (sizeMnistLabelInt :$: ZSR) FTKScalar)+      f :: ( MnistFcnnRanked2.ADFcnnMnist2Parameters+               (AstTensor AstMethodLet FullSpan) r q+           , ( AstTensor AstMethodLet FullSpan (TKR 1 r)+             , AstTensor AstMethodLet FullSpan (TKR 1 r) ) )+        -> AstTensor AstMethodLet FullSpan (TKScalar r)+      f (((hidden, bias), (hidden2, bias2), (readout, biasr)), (glyphR, labelR)) =+        unAstNoSimplify+        $ afcnnMnistLoss2 (AstNoSimplify glyphR, AstNoSimplify labelR)+                          ( (AstNoSimplify hidden, AstNoSimplify bias)+                          , (AstNoSimplify hidden2, AstNoSimplify bias2)+                          , (AstNoSimplify readout, AstNoSimplify biasr) )+      artRaw = revArtifactAdapt cotangentHandling f (FTKProduct ftk ftkData)+  in (targetInit, artRaw)+{-# SPECIALIZE mnistTrainBench2VTOGradientX :: Proxy Float -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int -> ( Concrete (XParams2 Double Float), AstArtifactRev (TKProduct (XParams2 Double Float) (TKProduct (TKR2 1 (TKScalar Double)) (TKR2 1 (TKScalar Double)))) (TKScalar Double) ) #-}+{-# SPECIALIZE mnistTrainBench2VTOGradientX :: Proxy Float -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int -> ( Concrete (XParams2 Float Float), AstArtifactRev (TKProduct (XParams2 Float Float) (TKProduct (TKR2 1 (TKScalar Float)) (TKR2 1 (TKScalar Float)))) (TKScalar Float) ) #-}+{-# SPECIALIZE mnistTrainBench2VTOGradientX :: Proxy Double -> IncomingCotangentHandling -> Double -> StdGen -> Int -> Int -> ( Concrete (XParams2 Double Double), AstArtifactRev (TKProduct (XParams2 Double Double) (TKProduct (TKR2 1 (TKScalar Double)) (TKR2 1 (TKScalar Double)))) (TKScalar Double) ) #-}
+ example/MnistRnnRanked2.hs view
@@ -0,0 +1,155 @@+{-# OPTIONS_GHC -Wno-missing-export-lists #-}+-- | Ranked tensor-based implementation of Recurrent Neural Network+-- for classification of MNIST digits. Sports 2 hidden layers.+module MnistRnnRanked2 where++import Prelude hiding (foldl')++import Data.Kind (Type)+import Data.List (foldl')+import Data.Vector.Generic qualified as V+import Data.Vector.Storable (Vector)+import GHC.TypeLits (KnownNat, Nat, type (+))++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape++import HordeAd+import MnistData++-- | The differentiable type of all trainable parameters of this nn.+-- Shaped version, statically checking all dimension widths.+type ADRnnMnistParametersShaped (target :: Target) width r =+  ( LayerWeigthsRNNShaped target SizeMnistHeight width r+  , LayerWeigthsRNNShaped target width width r+  , ( target (TKS '[SizeMnistLabel, width] r)+    , target (TKS '[SizeMnistLabel] r) ) )++type LayerWeigthsRNNShaped :: Target -> Nat -> Nat -> Type -> Type+type LayerWeigthsRNNShaped target in_width out_width r =+  ( target (TKS '[out_width, in_width] r)   -- input weight+  , target (TKS '[out_width, out_width] r)  -- state weight+  , target (TKS '[out_width] r) )           -- bias++-- | The differentiable type of all trainable parameters of this nn.+type ADRnnMnistParameters target r =+  ( LayerWeigthsRNN target r+  , LayerWeigthsRNN target r+  , ( target (TKR 2 r)+    , target (TKR 1 r) ) )++type LayerWeigthsRNN (target :: Target) r =+  ( target (TKR 2 r)+  , target (TKR 2 r)+  , target (TKR 1 r) )++zeroStateR+  :: (BaseTensor target, GoodScalar r)+  => IShR n -> (target (TKR n r)  -- state+                    -> a)+  -> a+zeroStateR sh f = f (rrepl sh 0)++unrollLastR :: forall target state c w r n.+               (BaseTensor target, GoodScalar r, KnownNat n)+            => (state -> target (TKR n r) -> w -> (c, state))+            -> (state -> target (TKR (1 + n) r) -> w -> (c, state))+unrollLastR f s0 xs w =+  let g :: (c, state) -> target (TKR n r) -> (c, state)+      g (_, !s) x = f s x w+  in foldl' g (undefined, s0) (runravelToList xs)++-- | A single recurrent layer with @tanh@ activation function.+rnnMnistLayerR+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 2 r)  -- ^ in state, @[out_width, batch_size]@+  -> target (TKR 2 r)  -- ^ input, @[in_width, batch_size]@+  -> LayerWeigthsRNN target r  -- ^ parameters+  -> target (TKR 2 r)  -- ^ output state, @[out_width, batch_size]@+rnnMnistLayerR s x (wX, wS, b) = case rshape s of+  _out_width :$: batch_size :$: ZSR ->+    let y = wX `rmatmul2` x + wS `rmatmul2` s+            + rtr (rreplicate batch_size b)+    in tanh y++-- TODO: represent state as a pair to avoid appending; tlet now supports this.+-- | Composition of two recurrent layers.+rnnMnistTwoR+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 2 r)  -- initial state, @[2 * out_width, batch_size]@+  -> PrimalOf target (TKR 2 r)  -- @[sizeMnistHeight, batch_size]@+  -> ( LayerWeigthsRNN target r  -- sizeMnistHeight out_width+     , LayerWeigthsRNN target r )  -- out_width out_width+  -> ( target (TKR 2 r)  -- @[out_width, batch_size]@+     , target (TKR 2 r) )  -- final state, @[2 * out_width, batch_size]@+rnnMnistTwoR s' x ((wX, wS, b), (wX2, wS2, b2)) = case rshape s' of+  out_width_x_2 :$: _batch_size :$: ZSR ->+    let out_width = out_width_x_2 `div` 2+        s3 = tlet s' $ \s ->+          let s1 = rslice 0 out_width s+              s2 = rslice out_width out_width s+              vec1 = rnnMnistLayerR s1 (rfromPrimal x) (wX, wS, b)+              vec2 = rnnMnistLayerR s2 vec1 (wX2, wS2, b2)+          in rappend vec1 vec2+    in (rslice out_width out_width s3, s3)++-- | The two-layer recurrent nn with its state initialized to zero+-- and the result composed with a fully connected layer.+rnnMnistZeroR+  :: (ADReady target, GoodScalar r, Differentiable r)+  => Int  -- ^ batch_size+  -> PrimalOf target (TKR 3 r)+       -- ^ input data @[sizeMnistWidth, sizeMnistHeight, batch_size]@+  -> ADRnnMnistParameters target r  -- ^ parameters+  -> target (TKR 2 r)  -- ^ output classification @[SizeMnistLabel, batch_size]@+rnnMnistZeroR batch_size xs+              ((wX, wS, b), (wX2, wS2, b2), (w3, b3)) = case rshape b of+  out_width :$: ZSR ->+    let sh = 2 * out_width :$: batch_size :$: ZSR+        (out, _s) = zeroStateR sh (unrollLastR rnnMnistTwoR) xs+                                  ((wX, wS, b), (wX2, wS2, b2))+    in w3 `rmatmul2` out + rtr (rreplicate batch_size b3)++-- | The neural network composed with the SoftMax-CrossEntropy loss function.+rnnMnistLossFusedR+  :: (ADReady target, ADReady (PrimalOf target), GoodScalar r, Differentiable r)+  => Int+  -> (PrimalOf target (TKR 3 r), PrimalOf target (TKR 2 r))  -- batch_size+  -> ADRnnMnistParameters target r  -- SizeMnistHeight out_width+  -> target (TKScalar r)+rnnMnistLossFusedR batch_size (glyphR, labelR) adparameters =+  let xs = rtranspose [2, 1, 0] glyphR+      result = rnnMnistZeroR batch_size xs adparameters+      targets = rtr labelR+      loss = lossSoftMaxCrossEntropyR targets result+  in kfromPrimal (recip $ kconcrete $ fromIntegral batch_size) * loss++-- | A function testing the neural network given testing set of inputs+-- and the trained parameters.+rnnMnistTestR+  :: forall target r.+     (target ~ Concrete, GoodScalar r, Differentiable r)+  => Int+  -> MnistDataBatchR r  -- batch_size+  -> ADRnnMnistParameters target r+  -> r+rnnMnistTestR 0 _ _ = 0+rnnMnistTestR batch_size (glyphR, labelR) testParams =+  let input :: target (TKR 3 r)+      input = rconcrete $ Nested.rtranspose [2, 1, 0] glyphR+      outputR :: Concrete (TKR 2 r)+      outputR =+        let nn :: ADRnnMnistParameters target r  -- SizeMnistHeight out_width+               -> target (TKR 2 r)  -- [SizeMnistLabel, batch_size]+            nn = rnnMnistZeroR batch_size input+        in nn testParams+      outputs = map rtoVector $ runravelToList+                $ rtranspose [1, 0] outputR+      labels = map rtoVector+               $ runravelToList @_ @(TKScalar r)+               $ rconcrete labelR+      matchesLabels :: Vector r -> Vector r -> Int+      matchesLabels output label | V.maxIndex output == V.maxIndex label = 1+                                 | otherwise = 0+  in fromIntegral (sum (zipWith matchesLabels outputs labels))+     / fromIntegral batch_size
+ example/MnistRnnShaped2.hs view
@@ -0,0 +1,175 @@+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+{-# OPTIONS_GHC -Wno-missing-export-lists #-}+-- | Shaped tensor-based implementation of Recurrent Neural Network+-- for classification of MNIST digits. Sports 2 hidden layers.+module MnistRnnShaped2 where++import Prelude hiding (foldl')++import Data.Kind (Type)+import Data.List (foldl')+import Data.Vector.Generic qualified as V+import Data.Vector.Storable (Vector)+import GHC.TypeLits (KnownNat, Nat, fromSNat, type (*))++import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Shaped.Shape++import HordeAd+import MnistData++-- | The differentiable type of all trainable parameters of this nn.+-- Shaped version, statically checking all dimension widths.+type ADRnnMnistParametersShaped+       (target :: Target) sizeMnistHeight width r =+  ( LayerWeigthsRNNShaped target sizeMnistHeight width r+  , LayerWeigthsRNNShaped target width width r+  , ( target (TKS '[SizeMnistLabel, width] r)+    , target (TKS '[SizeMnistLabel] r) ) )++type LayerWeigthsRNNShaped :: Target -> Nat -> Nat -> Type -> Type+type LayerWeigthsRNNShaped target in_width out_width r =+  ( target (TKS '[out_width, in_width] r)   -- input weight+  , target (TKS '[out_width, out_width] r)  -- state weight+  , target (TKS '[out_width] r) )           -- bias++zeroStateS+  :: (BaseTensor target, KnownShS sh, GoodScalar r)+  => (target (TKS sh r)  -- state+      -> a)+  -> a+zeroStateS f = f (srepl 0)++unrollLastS :: forall target state c w r n sh.+               (BaseTensor target, KnownNat n, KnownShS sh, GoodScalar r)+            => (state -> target (TKS sh r) -> w -> (c, state))+            -> (state -> target (TKS (n ': sh) r) -> w -> (c, state))+unrollLastS f s0 xs w =+  let g :: (c, state) -> target (TKS sh r) -> (c, state)+      g (_, !s) x = f s x w+  in foldl' g (undefined, s0) (sunravelToList xs)++-- | A single recurrent layer with @tanh@ activation function.+rnnMnistLayerS+  :: (ADReady target, GoodScalar r, Differentiable r)+  => SNat in_width -> SNat out_width -> SNat batch_size+       -- ^ these boilerplate lines tie type parameters to the corresponding+       -- value parameters (@SNat@ below) denoting basic dimensions+  -> target (TKS '[out_width, batch_size] r)  -- state+  -> target (TKS '[in_width, batch_size] r)  -- input+  -> LayerWeigthsRNNShaped target in_width out_width r+  -> target (TKS '[out_width, batch_size] r)  -- output state+rnnMnistLayerS SNat SNat SNat+               s x (wX, wS, b) =+    let y = wX `smatmul2` x + wS `smatmul2` s+            + str (sreplicate {-@batch_size-} b)+    in tanh y++-- TODO: represent state as a pair to avoid appending; tlet now supports this.+-- | Composition of two recurrent layers.+rnnMnistTwoS+  :: (ADReady target, GoodScalar r, Differentiable r)+  => SNat out_width -> SNat batch_size -> SNat sizeMnistH+  -> target (TKS '[2 * out_width, batch_size] r)  -- initial state+  -> PrimalOf target (TKS '[sizeMnistH, batch_size] r)+  -> ( LayerWeigthsRNNShaped target sizeMnistH out_width r+     , LayerWeigthsRNNShaped target out_width out_width r )+  -> ( target (TKS '[out_width, batch_size] r)+     , target (TKS '[2 * out_width, batch_size] r) )  -- final state+rnnMnistTwoS out_width@SNat+             batch_size@SNat+             sizeMnistHeightHere@SNat+             s' x ((wX, wS, b), (wX2, wS2, b2)) =+    let s3 = tlet s' $ \s ->+          let s1 = sslice (SNat @0) out_width SNat s+              s2 = sslice out_width out_width SNat s+              vec1 = rnnMnistLayerS sizeMnistHeightHere+                                    out_width+                                    batch_size+                                    s1 (sfromPrimal x) (wX, wS, b)+              vec2 = rnnMnistLayerS out_width+                                    out_width+                                    batch_size+                                    s2 vec1 (wX2, wS2, b2)+          in sappend vec1 vec2+    in (sslice out_width out_width SNat s3, s3)++-- | The two-layer recurrent nn with its state initialized to zero+-- and the result composed with a fully connected layer.+rnnMnistZeroS+  :: (ADReady target, GoodScalar r, Differentiable r)+  => SNat out_width+  -> SNat batch_size+  -> SNat sizeMnistH -> SNat sizeMnistW+  -> PrimalOf target (TKS '[sizeMnistW, sizeMnistH, batch_size] r)+  -> ADRnnMnistParametersShaped target sizeMnistH out_width r+  -> target (TKS '[SizeMnistLabel, batch_size] r)+rnnMnistZeroS out_width@SNat+              batch_size@SNat+              sizeMnistHeightHere@SNat _sizeMnistWidthHere@SNat+              xs ((wX, wS, b), (wX2, wS2, b2), (w3, b3)) =+    let rnnMnistTwo = rnnMnistTwoS out_width batch_size sizeMnistHeightHere+        (out, _s) = zeroStateS (unrollLastS rnnMnistTwo) xs+                               ((wX, wS, b), (wX2, wS2, b2))+    in w3 `smatmul2` out + str (sreplicate {-@batch_size-} b3)++-- | The neural network composed with the SoftMax-CrossEntropy loss function.+rnnMnistLossFusedS+  :: forall target h w out_width batch_size r.+     ( h ~ SizeMnistHeight, w ~ SizeMnistWidth, Differentiable r+     , ADReady target, ADReady (PrimalOf target), GoodScalar r)+  => SNat out_width+  -> SNat batch_size+  -> ( PrimalOf target (TKS '[batch_size, h, w] r)+     , PrimalOf target (TKS '[batch_size, SizeMnistLabel] r) )+  -> ADRnnMnistParametersShaped target h out_width r+  -> target (TKScalar r)+rnnMnistLossFusedS out_width@SNat+                   batch_size@SNat+                   (glyphS, labelS) adparameters =+  let xs = stranspose @'[2, 1, 0] glyphS+      result = rnnMnistZeroS out_width+                             batch_size+                             (SNat @h) (SNat @w)+                             xs adparameters+      targets = str labelS+      loss = lossSoftMaxCrossEntropyS targets result+  in kfromPrimal (recip $ kconcrete $ fromInteger $ fromSNat batch_size) * loss++-- | A function testing the neural network given testing set of inputs+-- and the trained parameters.+rnnMnistTestS+  :: forall target h w out_width batch_size r.+     ( h ~ SizeMnistHeight, w ~ SizeMnistWidth+     , target ~ Concrete, Differentiable r, GoodScalar r )+  => SNat out_width+  -> SNat batch_size+  -> MnistDataBatchS batch_size r+  -> ADRnnMnistParametersShaped target h out_width r+  -> r+rnnMnistTestS out_width@SNat batch_size@SNat+              (glyphS, labelS) testParams =+  let -- input :: PrimalOf target (TKS '[sizeMnistW, sizeMnistH, batch_size] r)+      input = sconcrete+              $ Nested.stranspose (Permutation.makePerm @'[2, 1, 0]) glyphS+      outputS :: Concrete (TKS '[SizeMnistLabel, batch_size] r)+      outputS =+        let nn :: ADRnnMnistParametersShaped target h out_width r+               -> target (TKS '[SizeMnistLabel, batch_size] r)+            nn = rnnMnistZeroS out_width+                               batch_size+                               (SNat @h) (SNat @w)+                               input+        in nn testParams+      outputs = map stoVector $ sunravelToList+                $ stranspose @'[1, 0] outputS+      labels = map stoVector+               $ sunravelToList @_ @_ @(TKScalar r)+               $ sconcrete labelS+      matchesLabels :: Vector r -> Vector r -> Int+      matchesLabels output label | V.maxIndex output == V.maxIndex label = 1+                                 | otherwise = 0+  in fromIntegral (sum (zipWith matchesLabels outputs labels))+     / fromInteger (fromSNat batch_size)
+ horde-ad.cabal view
@@ -0,0 +1,468 @@+cabal-version:      3.6+-- The cabal-version field refers to the version of the .cabal specification,+-- and can be different from the cabal-install (the tool) version and the+-- Cabal (the library) version you are using. As such, the Cabal (the library)+-- version used must be equal or greater than the version stated in this field.+-- Starting from the specification version 2.2, the cabal-version field must be+-- the first thing in the cabal file.++-- Initial package description 'horde-ad' generated by+-- 'cabal init'. For further documentation, see:+--   http://haskell.org/cabal/users-guide/+--+-- The name of the package.+name:               horde-ad++-- The package version.+-- See the Haskell package versioning policy (PVP) for standards+-- guiding when and how versions should be incremented.+-- https://pvp.haskell.org+-- PVP summary:     +-+------- breaking API changes+--                  | | +----- non-breaking API additions+--                  | | | +--- code changes with no API change+version:            0.1.0.0++-- A short (one-line) description of the package.+synopsis:           Higher Order Reverse Derivatives Efficiently - Automatic Differentiation++-- A longer description of the package.+description:        An Automatic Differentiation library originally inspired by the paper "Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation" (POPL 2022). Compared to the paper and to classic taping AD Haskell packages, the library additionally efficiently supports array operations and generation of symbolic derivative programs, though the efficiency is confined to a narrowly typed class of source programs with limited higher-orderness. A detailed account of the extension is in the paper by Tom Smeding, Mikolaj Konarski, Simon Peyton Jones and Andrew Fitzgibbon available at http://arxiv.org/abs/2507.12640.++-- The license under which the package is released.+license:            BSD-3-Clause++-- The file containing the license text.+license-file:       LICENSE++-- The package author(s).+author:             Mikolaj Konarski and others++-- An email address to which users can send suggestions, bug reports, and patches.+maintainer:         mikolaj@well-typed.com++category:           Machine Learning, Tensors++-- A copyright notice.+-- copyright:+build-type:         Simple++-- Extra doc files to be distributed with the package, such as a CHANGELOG or a README.+extra-doc-files:    README.md,+                    CHANGELOG.md,+                    CREDITS.md++tested-with:        GHC ==9.10.2 || ==9.12.3++bug-reports:        https://github.com/Mikolaj/horde-ad/issues++source-repository head+  type:               git+  location:           https://github.com/Mikolaj/horde-ad.git++flag with_expensive_assertions+  description:        turn on expensive assertions of well-tested code+  default:            False+  manual:             True++flag release+  description:        remember to set before tagging for a release (to expose internal functions and types, disable tests that depend on not packaged data, etc.)+  default:            True+  manual:             True++flag test_seq+  description:        run all of the test suite in sequential mode+  default:            False+  manual:             True++common options+  default-language:   GHC2024+  default-extensions: StrictData, TypeFamilies, TypeFamilyDependencies,+                      FunctionalDependencies, RecordWildCards, MultiWayIf,+                      DefaultSignatures, PatternSynonyms, NoStarIsType,+                      TypeData, TypeAbstractions+  other-extensions:   UnboxedTuples, CPP, ViewPatterns, OverloadedLists,+                      DeriveAnyClass, UndecidableInstances, AllowAmbiguousTypes,+                      QuantifiedConstraints, TemplateHaskell, DerivingVia,+                      ImpredicativeTypes+  ghc-options:        -Wall -Wcompat -Wimplicit-prelude -Widentities -Wredundant-constraints -Wmissing-export-lists -Wpartial-fields -Wunused-packages -Wredundant-bang-patterns -Woperator-whitespace -Wredundant-strictness-flags -Wterm-variable-capture -Wmissed-specialisations+  ghc-options:        -Wmissing-poly-kind-signatures -Wmissing-role-annotations -Wno-unticked-promoted-constructors+  ghc-options:        -fprint-explicit-kinds+  -- This is supposedly risky, but prevents STG from creating thunks, #23848:+  ghc-options:        -fworker-wrapper-cbv+  -- Disabling this breaks performance:+  ghc-options:        -fexpose-all-unfoldings -fspecialise-aggressively -fsimpl-tick-factor=200+  -- Disabling hurts performance to some extent:+  ghc-options:        -fdicts-cheap -flate-dmd-anal+  -- ghc-options: -O2 -flate-specialise++  if flag(with_expensive_assertions)+    cpp-options:      -DWITH_EXPENSIVE_ASSERTIONS++  if flag(release)+    cpp-options:      -DEXPOSE_INTERNAL++  -- This needs to be after -O2.+  ghc-options:        -fno-ignore-asserts++common exe-options+  ghc-options:        -rtsopts+  -- Make GC more predictable in benchmarks.+  ghc-options:        "-with-rtsopts=-H1.5g -A0.75g -I0"++common exe-options-for-tests-that-reset-counters+  ghc-options:        -rtsopts -threaded+  ghc-options:        "-with-rtsopts=-H1.5g -A0.75g -I0"+    -- unconditionally seqential, because counter reset is not thread-safe++common exe-options-test+  ghc-options:        -rtsopts -threaded+  if flag(test_seq)+    ghc-options:      "-with-rtsopts=-H1.5g -A0.75g -I0"+  else+    -- Beware, -with-rtsopts is not cumulative!+    ghc-options:      "-with-rtsopts=-H1.5g -A0.75g -I0 -N"++library+    import:           options++    exposed-modules:  HordeAd+                      HordeAd.ADEngine+                      HordeAd.AstEngine+                      HordeAd.Core.Adaptor+                      HordeAd.Core.Ast+                      HordeAd.Core.AstEnv+                      HordeAd.Core.AstFreshId+                      HordeAd.Core.AstInline+                      HordeAd.Core.AstInterpret+                      HordeAd.Core.AstPrettyPrint+                      HordeAd.Core.AstSimplify+                      HordeAd.Core.AstTools+                      HordeAd.Core.AstTraverse+                      HordeAd.Core.AstVectorize+                      HordeAd.Core.CarriersADVal+                      HordeAd.Core.CarriersAst+                      HordeAd.Core.CarriersConcrete+                      HordeAd.Core.ConvertTensor+                      HordeAd.Core.Delta+                      HordeAd.Core.DeltaEval+                      HordeAd.Core.DeltaFreshId+                      HordeAd.Core.Ops+                      HordeAd.Core.OpsADVal+                      HordeAd.Core.OpsAst+                      HordeAd.Core.OpsConcrete+                      HordeAd.Core.TensorKind+                      HordeAd.Core.Types+                      HordeAd.Core.Unwind+                      HordeAd.External.CommonRankedOps+                      HordeAd.External.CommonShapedOps+                      HordeAd.External.Optimizer+                      HordeAd.External.OptimizerTools+                      HordeAd.OpsTensor+                      HordeAd.OpsTensorRanked+                      HordeAd.OpsTensorShaped+                      HordeAd.OpsTensorMixed++    hs-source-dirs:   src++    build-depends:+        assert-failure < 0.1.4+      , base >= 4.20.1 && < 4.21+          -- backport of https://gitlab.haskell.org/ghc/ghc/-/merge_requests/13498 is required for GHC 9.12; TODO: adjust base bounds once 9.12.3 is out+      , Boolean < 0.3+      , containers >= 0.6 && < 0.9+      , deepseq < 1.7+      , data-default < 0.9+      , dependent-enummap < 0.2+      , dependent-sum >= 0.7 && < 0.8+      , enummapset < 0.8+      , ghc-typelits-knownnat < 1+      , ghc-typelits-natnormalise < 1+      , ilist < 0.5+      , orthotope < 0.2+      , ox-arrays < 0.2+      , random >= 1.3.0 && < 1.4+      , some >= 1 && < 2+      , atomic-counter < 0.2+      , vector >= 0.13.2.0 && < 0.14++library exampleLibrary+    import:           options++    exposed-modules:  MnistData+                      MnistCnnRanked2+                      MnistCnnShaped2+                      MnistFcnnRanked1+                      MnistFcnnRanked2+                      MnistRnnRanked2+                      MnistRnnShaped2++    hs-source-dirs:   example++    build-depends:+      , horde-ad:horde-ad++      , base+      , bytestring < 0.13+      , ghc-typelits-knownnat+      , ghc-typelits-natnormalise+      , mnist-idx < 0.2+      , ox-arrays+      , random+      , vector+      , zlib < 0.8++library benchCommonLibrary+    import:           options++    exposed-modules:  BenchMnistTools+                      BenchProdTools++    hs-source-dirs:   bench/common++    build-depends:+      , horde-ad:exampleLibrary+      , horde-ad:horde-ad++      , base+      , criterion < 1.7+      , deepseq < 1.7+      , data-default < 0.9+      , inspection-testing < 0.7+      , ox-arrays+      , random++benchmark longProdBench+    import:           options, exe-options++    type:             exitcode-stdio-1.0++    hs-source-dirs:   bench++    main-is:          LongProdBench.hs++    build-depends:+      , horde-ad:benchCommonLibrary++      , base+      , criterion+      , deepseq+      , random++benchmark shortProdForCI+    import:           options, exe-options++    type:             exitcode-stdio-1.0++    hs-source-dirs:   bench++    main-is:          ShortProdForCI.hs++    build-depends:+      , horde-ad:benchCommonLibrary++      , base+      , criterion+      , deepseq+      , random++-- The data files for this are not included in the cabal package,+-- hence 'buildable: False'.+benchmark longMnistBench+    import:           options, exe-options++    if flag(release)+      buildable: False+    else+      buildable: True++    type:             exitcode-stdio-1.0++    hs-source-dirs:   bench++    main-is:          LongMnistBench.hs++    build-depends:+      , horde-ad:benchCommonLibrary++      , base+      , criterion++-- The data files for this are not included in the cabal package,+-- hence 'buildable: False'.+benchmark shortMnistForCI+    import:           options, exe-options++    if flag(release)+      buildable: False+    else+      buildable: True++    type:             exitcode-stdio-1.0++    hs-source-dirs:   bench++    main-is:          ShortMnistForCI.hs++    build-depends:+      , horde-ad:benchCommonLibrary++      , base+      , criterion++-- takes forever, so mostly to invoke with --benchmark-options='-n 1 +RTS -s'+benchmark realisticMnistBench+    import:           options, exe-options++-- Disabled so that `cabal bench` terminates in finite time.+--    if flag(release)+    buildable: False+--    else+--      buildable: True++    type:             exitcode-stdio-1.0++    hs-source-dirs:   bench++    main-is:          RealisticMnistBench.hs++    build-depends:+      , horde-ad:benchCommonLibrary+      , horde-ad:exampleLibrary++      , base+      , criterion+      , random++library testToolLibrary+    import:           options++    exposed-modules:  CrossTesting+                      EqEpsilon+                      Shared++    hs-source-dirs:   test/tool++    build-depends:+      , horde-ad:horde-ad++      , base+      , ox-arrays+      , tasty >= 1.0 && < 1.6+      , tasty-hunit < 0.11+      , vector++library testCommonLibrary+    import:           options++    exposed-modules:  TestAdaptorSimplified+                      TestConvSimplified+                      TestGatherSimplified+                      TestHighRankSimplified+                      TestMnistCNNR+                      TestMnistCNNS+                      TestMnistFCNNR+                      TestMnistPP+                      TestMnistRNNR+                      TestMnistRNNS+                      TestRevFwdFold++    hs-source-dirs:   test/simplified++    -- Other library packages from which modules are imported.+    build-depends:+      , horde-ad:horde-ad+      , horde-ad:exampleLibrary+      , horde-ad:testToolLibrary++      , assert-failure+      , base+      , ghc-typelits-knownnat+      , ghc-typelits-natnormalise+      , ox-arrays+      , random+      , tasty >= 1.0 && < 1.6+      , tasty-hunit < 0.11+      , tasty-quickcheck < 0.12++-- The data files for this are not included in the cabal package,+-- hence 'buildable: False'.+test-suite fullTest+    import:           options, exe-options-for-tests-that-reset-counters++    if flag(release)+      buildable: False+    else+      buildable: True++    type:             exitcode-stdio-1.0++    hs-source-dirs:   test++    main-is:          FullTest.hs++    build-depends:+      , horde-ad:testCommonLibrary+      , horde-ad:testToolLibrary++      , base+      , tasty >= 1.0++test-suite CAFlessTest+    import:           options, exe-options-for-tests-that-reset-counters++    if flag(release)+      buildable: False+    else+      buildable: True++    type:             exitcode-stdio-1.0++    hs-source-dirs:   test++    main-is:          CAFlessTest.hs++    build-depends:+      , horde-ad:testCommonLibrary+      , horde-ad:testToolLibrary++      , base+      , tasty >= 1.0++test-suite parallelTest+    import:           options, exe-options-test++    if flag(release)+      buildable: False+    else+      buildable: True++    type:             exitcode-stdio-1.0++    hs-source-dirs:   test++    main-is:          ParallelTest.hs++    build-depends:+      , horde-ad:testCommonLibrary+      , horde-ad:testToolLibrary++      , base+      , tasty >= 1.0++-- These tests don't require any dataset and are short, so can be enabled+-- in the Hackage package.+test-suite minimalTest+    import:           options, exe-options-for-tests-that-reset-counters++    type:             exitcode-stdio-1.0++    hs-source-dirs:   test++    main-is:          MinimalTest.hs++    build-depends:+      , horde-ad:testCommonLibrary+      , horde-ad:testToolLibrary++      , base+      , tasty >= 1.0
+ src/HordeAd.hs view
@@ -0,0 +1,50 @@+-- | A subset of the API of the horde-ad library that should be+-- sufficient for general use.+-- To (ab)use implementation details or to access ready tools for specific+-- applications, such as neural networks operating on MNIST data,+-- you may need some extra imports not covered here.+--+-- If you are interested mainly in using multi-dimensional arrays,+-- start with "HordeAd.OpsTensor" (or use directly+-- https://hackage.haskell.org/package/ox-arrays+-- that doesn't have the capability nor the overhead of term rewriting+-- over the symbolic represenation of array operations).+-- For Automatic Differentiation, start with "HordeAd.ADEngine".+module HordeAd+  ( -- * The main array API+    module HordeAd.OpsTensor+  , module HordeAd.OpsTensorRanked+  , module HordeAd.OpsTensorShaped+  , module HordeAd.OpsTensorMixed+  , module HordeAd.Core.ConvertTensor+    -- * The main AD API+  , module HordeAd.ADEngine+  , module HordeAd.AstEngine+    -- * Additional support types and operations+  , module HordeAd.Core.Ast+  , module HordeAd.Core.CarriersADVal+  , module HordeAd.Core.CarriersConcrete+  , module HordeAd.Core.TensorKind+  , module HordeAd.Core.Types+  , module HordeAd.External.CommonRankedOps+  , module HordeAd.External.CommonShapedOps+  , module HordeAd.External.Optimizer+  ) where++import Prelude ()++import HordeAd.ADEngine+import HordeAd.AstEngine+import HordeAd.Core.Ast+import HordeAd.Core.CarriersADVal+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.External.CommonRankedOps+import HordeAd.External.CommonShapedOps+import HordeAd.External.Optimizer+import HordeAd.OpsTensor+import HordeAd.OpsTensorMixed+import HordeAd.OpsTensorRanked+import HordeAd.OpsTensorShaped
+ src/HordeAd/ADEngine.hs view
@@ -0,0 +1,603 @@+{-# OPTIONS_GHC -Wno-orphans #-}+-- | The implementation of reverse derivative and forward derivative+-- calculation for an objective function on values of complicated types,+-- e.g., nested tuples of tensors.+--+-- The objective function can be defined as a sufficiently polymorphic+-- Haskell function that uses numeric classes as well as the multi-dimensional+-- tensor operation listed in "HordeAd.OpsTensor". To obtain symbolic+-- derivatives (derivative code that can be executed many times without+-- performing AD again), the user needs an objective function polymorphic+-- enough so that it can be instantiated to the 'HordeAd.Core.Ast.AstTensor'+-- type (nested in tuples, etc., for some extra flexibility).+-- For non-symbolic derivatives, the ability to instantiate to the+-- `HordeAd.Core.CarriersADVal.ADVal` type of dual numbers is enough.+-- See the classes these types are instances of to gauge the breadth+-- of the offered respective APIs.+module HordeAd.ADEngine+  ( -- * Symbolic reverse derivative adaptors+    grad, vjp+  , gradArtifact, vjpArtifact+  , gradInterpretArtifact, vjpInterpretArtifact+    -- * Symbolic forward derivative adaptors+  , jvp, jvpArtifact, jvpInterpretArtifact+    -- * Non-symbolic reverse derivative adaptors+  , cgrad, cvjp+    -- * Non-symbolic forward derivative adaptors+  , cjvp+    -- * Internal machinery for symbolic adaptors+  , IncomingCotangentHandling(..)+  , revArtifactAdapt, revArtifactDelta+  , revProduceArtifactWithoutInterpretation, revInterpretArtifact+  , fwdArtifactAdapt, fwdArtifactDelta, fwdInterpretArtifact+    -- * Internal machinery for non-symbolic adaptors+  , cfwdBoth+  ) where++import Prelude++import HordeAd.AstEngine+import HordeAd.Core.Adaptor+import HordeAd.Core.Ast+import HordeAd.Core.AstEnv+import HordeAd.Core.AstInterpret+import HordeAd.Core.CarriersADVal+import HordeAd.Core.CarriersAst+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Delta+import HordeAd.Core.DeltaEval+import HordeAd.Core.Ops+import HordeAd.Core.OpsADVal+import HordeAd.Core.OpsAst+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind++-- * Symbolic reverse derivative adaptors++-- | This simplified version of the symbolic reverse derivative operation+-- sets the incoming cotangent @dt@ to be 1 and assumes the codomain+-- of the function to be differentiated is a scalar.+--+-- We don't enforce (e.g., by quantifcation) that the objective function+-- is closed, because we evaluate the result of the differentiation+-- down to concrete arrays and so there's no risk of "perturbation confusion"+-- between different levels of differentiation if it's done multiple times.+-- For simplicity of the type signature, the resulting value is converted from+-- the type of concrete contangents to the type of concrete input parameters.+grad+  :: forall src r tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan (TKScalar r) )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> Value src  -- morally Value (ADTensorKind src)+{-# INLINE grad #-}+grad f vals = revMaybe f vals Nothing++-- | This version of the symbolic reverse derivative operation+-- explicitly takes the sensitivity parameter (the incoming cotangent).+-- It also permits an arbitrary (nested tuple+) type of the domain+-- and arbitrary (nested pair) tensor kind of the codomain+-- of the function to be differentiated. The downside of the generality+-- is that if the function doesn't have an explicit type signature,+-- the type to which this operation is instantiated often has to be spelled+-- in full via explicit type applications to aid type reconstruction.+-- For simplicity of the type signature, the resulting value is converted from+-- the type of concrete contangents to the type of concrete input parameters.+vjp+  :: forall src ztgt tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> Concrete (ADTensorKind ztgt)+  -> Value src  -- morally Value (ADTensorKind src)+{-# INLINE vjp #-}+vjp f vals dt = revMaybe f vals (Just dt)++-- | Compute the reverse derivative not for a specific input, but as symbolic+-- function from inputs to the gradient value.+-- The function is represented as an "artifact", which is the gradient+-- AST term together with the variable corresponding to the input.+gradArtifact+  :: forall src r tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan (TKScalar r) )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> AstArtifactRev (X src) (TKScalar r)+       -- ^ the artifact containing the symbolic code of the derivative+{-# INLINE gradArtifact #-}+gradArtifact f vals0 =+  let xftk = tftkG (knownSTK @(X src)) $ unConcrete $ toTarget vals0+  in revArtifactAdapt IgnoreIncomingCotangent f xftk++-- | Compute the reverse derivative not for a specific input, but as symbolic+-- function from inputs and incoming cotangents to the gradient value.+-- The function is represented as an "artifact", which is the gradient+-- AST term together with variables corresponding to the input and cotangent.+vjpArtifact+  :: forall src ztgt tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> AstArtifactRev (X src) ztgt+       -- ^ the artifact containing the symbolic code of the derivative+{-# INLINE vjpArtifact #-}+vjpArtifact f vals0 =+  let xftk = tftkG (knownSTK @(X src)) $ unConcrete $ toTarget vals0+  in revArtifactAdapt UseIncomingCotangent f xftk++-- | Interpret the "artifact" as a function from a concrete tensor+-- to a concrete tensor (possibly adapted, e.g., from horde-ad nested pairs+-- to Haskell n-tuples).+gradInterpretArtifact+  :: forall x r avals.+     (X avals ~ ADTensorKind x, AdaptableTarget Concrete avals)+  => AstArtifactRev x (TKScalar r)+       -- ^ the artifact containing the symbolic code of the derivative+  -> Concrete x+  -> avals+{-# INLINE gradInterpretArtifact #-}+gradInterpretArtifact AstArtifactRev{..} parameters =+  let xftk = varNameToFTK artVarDomainRev+      azftk = varNameToFTK artVarDtRev+                -- STKScalar @(ADTensorScalar r) or STKScalar @Z1+      oneAtF = treplTarget 1 azftk+      env = extendEnv artVarDtRev oneAtF+            $ extendEnv artVarDomainRev parameters emptyEnv+  in if tftkG (ftkToSTK xftk) (unConcrete parameters) == xftk+     then fromTarget $ interpretAstPrimal env artDerivativeRev+     else error "gradInterpretArtifact: reverse derivative parameters must have the same shape as the domain of the objective function"++-- | Interpret the "artifact" as a function from concrete tensors+-- to a concrete tensor (possibly adapted, e.g., from horde-ad nested pairs+-- to Haskell n-tuples).+vjpInterpretArtifact+  :: forall x z avals.+     (X avals ~ ADTensorKind x, AdaptableTarget Concrete avals)+  => AstArtifactRev x z+       -- ^ the artifact containing the symbolic code of the derivative+  -> Concrete x+  -> Concrete (ADTensorKind z)+  -> avals+{-# INLINE vjpInterpretArtifact #-}+vjpInterpretArtifact AstArtifactRev{..} parameters dt =+  let xftk = varNameToFTK artVarDomainRev+      azftk = varNameToFTK artVarDtRev+      env = extendEnv artVarDtRev dt+            $ extendEnv artVarDomainRev parameters emptyEnv+  in if tftkG (ftkToSTK xftk) (unConcrete parameters) == xftk+     then if tftkG (ftkToSTK azftk) (unConcrete dt) == azftk+          then fromTarget $ interpretAstPrimal env artDerivativeRev+          else error "vjpInterpretArtifact: reverse derivative incoming cotangent must have the same shape as the codomain of the objective function"+     else error "vjpInterpretArtifact: reverse derivative parameters must have the same shape as the domain of the objective function"+++-- * Symbolic reverse derivative adaptors' internal machinery++revMaybe+  :: forall src ztgt tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> Maybe (Concrete (ADTensorKind ztgt))+  -> Value src  -- morally Value (ADTensorKind src)+{-# INLINE revMaybe #-}+revMaybe f vals0 mdt =+  let valsTarget = toTarget vals0+      xftk = tftkG (knownSTK @(X src)) $ unConcrete valsTarget+      cotangentHandling =+        maybe IgnoreIncomingCotangent (const UseIncomingCotangent) mdt+      artifactRaw = revArtifactAdapt cotangentHandling f xftk+      artifact = simplifyArtifactGradient artifactRaw+  in fromTarget $ fromADTensorKindShared (ftkToSTK xftk)+     $ fst $ revInterpretArtifact artifact valsTarget mdt++revArtifactAdapt+  :: forall src ztgt tgt.+     ( AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => IncomingCotangentHandling+  -> (src -> tgt)  -- ^ the objective function+  -> FullShapeTK (X src)+  -> AstArtifactRev (X src) ztgt+       -- ^ the artifact containing the symbolic code of the derivative+{-# INLINE revArtifactAdapt #-}+revArtifactAdapt cotangentHandling f xftk =+  let g :: AstTensor AstMethodLet FullSpan (X src) -> tgt+      g !arg = simplifyInline $ ttlet arg $ f . fromTarget+                                  -- fromTarget requires duplicable+  in revProduceArtifact cotangentHandling g emptyEnv xftk++revInterpretArtifact+  :: forall x z.+     AstArtifactRev x z+       -- ^ the artifact containing the symbolic code of the derivative+  -> Concrete x+  -> Maybe (Concrete (ADTensorKind z))+  -> (Concrete (ADTensorKind x), Concrete z)+{-# INLINE revInterpretArtifact #-}+revInterpretArtifact AstArtifactRev{..} parameters mdt =+  let azftk = varNameToFTK artVarDtRev+      env = extendEnv artVarDomainRev parameters emptyEnv+      envDt = case mdt of+        Nothing ->+          let oneAtF = treplTarget 1 azftk+          in extendEnv artVarDtRev oneAtF env+        Just dt ->+          if tftkG (ftkToSTK azftk) (unConcrete dt) == azftk+          then extendEnv artVarDtRev dt env+          else error "revInterpretArtifact: reverse derivative incoming cotangent must have the same shape as the codomain of the objective function"+      gradient = interpretAstPrimal envDt artDerivativeRev+      primal = interpretAstPrimal env artPrimalRev+  in (gradient, primal)+++-- * Symbolic reverse derivative adaptors' testing-only internal machinery++revArtifactDelta+  :: forall src ztgt tgt.+     ( AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => IncomingCotangentHandling+  -> (src -> tgt)  -- ^ the objective function+  -> FullShapeTK (X src)+  -> (AstArtifactRev (X src) ztgt, Delta (AstRaw PrimalSpan) ztgt)+       -- ^ the artifact containing the symbolic code of the derivative+{-# INLINE revArtifactDelta #-}+revArtifactDelta cotangentHandling f xftk =+  let g :: AstTensor AstMethodLet FullSpan (X src) -> tgt+      g !arg = ttlet arg $ f . fromTarget+  in revArtifactFromForwardPass cotangentHandling+                                (forwardPassByInterpretation g emptyEnv) xftk++revProduceArtifactWithoutInterpretation+  :: forall x z.+     IncomingCotangentHandling+  -> (ADVal (AstRaw PrimalSpan) x -> ADVal (AstRaw PrimalSpan) z)+  -> FullShapeTK x+  -> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)+       -- ^ the artifact containing the symbolic code of the derivative+{-# INLINE revProduceArtifactWithoutInterpretation #-}+revProduceArtifactWithoutInterpretation cotangentHandling f xftk =+  -- No simplification performed to let individual tests decide.+  revArtifactFromForwardPass cotangentHandling+                             (forwardPassByApplication f)+                             xftk++forwardPassByApplication+  :: forall x z.+     (ADVal (AstRaw PrimalSpan) x -> ADVal (AstRaw PrimalSpan) z)+  -> AstTensor AstMethodShare PrimalSpan x+  -> AstVarName FullSpan x+  -> AstTensor AstMethodLet FullSpan x+  -> ADVal (AstRaw PrimalSpan) z+{-# INLINE forwardPassByApplication #-}+forwardPassByApplication g astVarPrimal var _astVar =+  let deltaInputs = generateDeltaInputs $ varNameToFTK var+      varInputs = dDnotShared (AstRaw astVarPrimal) deltaInputs+  in g varInputs+++-- * Symbolic forward derivative adaptors++-- | The forward derivative operation takes the perturbation parameter+-- by convention. It permits an arbitrary (nested tuple+)+-- type of the domain and arbitrary (nested pair) tensor kind of the codomain+-- of the function to be differentiated. The generality sometimes makes it+-- necessary to suppy type hints when applying this operation.+jvp+  :: forall src ztgt tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> Value src  -- morally (ADTensorKind src)+  -> Concrete (ADTensorKind ztgt)+{-# INLINE jvp #-}+jvp f vals0 ds =+  let valsTarget = toTarget vals0+      xftk = tftkG (knownSTK @(X src)) $ unConcrete valsTarget+      artifactRaw = fwdArtifactAdapt f xftk+      artifact = simplifyArtifactDerivative artifactRaw+  in fst $ fwdInterpretArtifact artifact valsTarget+         $ toADTensorKindShared xftk (toTarget ds)+       -- the shapes of vals0 vs ds are checked in fwdInterpretArtifact++-- | Compute the forward derivative not for a specific input, but as symbolic+-- function from inputs and perturbation to the derivative value.+-- The function is represented as an "artifact", which is the derivative+-- AST term together with variables corresponding to the input and perturbation.+jvpArtifact+  :: forall src ztgt tgt.+     ( X src ~ X (Value src), KnownSTK (X src)+     , AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , AdaptableTarget Concrete (Value src)+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> Value src+  -> AstArtifactFwd (X src) ztgt+       -- ^ the artifact containing the symbolic code of the derivative+{-# INLINE jvpArtifact #-}+jvpArtifact f vals0 =+  let xftk = tftkG (knownSTK @(X src)) $ unConcrete $ toTarget vals0+  in fwdArtifactAdapt f xftk++-- | Interpret the "artifact" as a function from concrete tensors+-- to a concrete tensor.+jvpInterpretArtifact+  :: forall x z.+     AstArtifactFwd x z+       -- ^ the artifact containing the symbolic code of the derivative+  -> Concrete x+  -> Concrete (ADTensorKind x)+  -> Concrete (ADTensorKind z)+{-# INLINE jvpInterpretArtifact #-}+jvpInterpretArtifact art parameters = fst . fwdInterpretArtifact art parameters+  -- the shapes of parameters vs ds are checked in fwdInterpretArtifact+++-- * Symbolic forward derivative adaptors' internal machinery++fwdArtifactAdapt+  :: forall src ztgt tgt.+     ( AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> FullShapeTK (X src)+  -> AstArtifactFwd (X src) ztgt+       -- ^ the artifact containing the symbolic code of the derivative+{-# INLINE fwdArtifactAdapt #-}+fwdArtifactAdapt f xftk =+  let g :: AstTensor AstMethodLet FullSpan (X src) -> tgt+      g !arg = simplifyInline $ ttlet arg $ f . fromTarget+                                  -- fromTarget requires duplicable+  in fwdProduceArtifact g emptyEnv xftk++fwdInterpretArtifact+  :: forall x z.+     AstArtifactFwd x z+       -- ^ the artifact containing the symbolic code of the derivative+  -> Concrete x+  -> Concrete (ADTensorKind x)+  -> (Concrete (ADTensorKind z), Concrete z)+{-# INLINE fwdInterpretArtifact #-}+fwdInterpretArtifact AstArtifactFwd{..} parameters ds =+  let xftk = varNameToFTK artVarDomainFwd+      xstk = ftkToSTK xftk+      env = extendEnv artVarDomainFwd parameters emptyEnv+      envD = extendEnv artVarDsFwd ds env+  in if tftkG xstk (unConcrete parameters) == xftk+     then if tftkG (adSTK xstk) (unConcrete ds) == adFTK xftk+          then let derivative = interpretAstPrimal envD artDerivativeFwd+                   primal = interpretAstPrimal env artPrimalFwd+               in (derivative, primal)+          else error "fwdInterpretArtifact: forward derivative perturbation must have the same shape as the domain of the objective function"+     else error "fwdInterpretArtifact: forward derivative input must have the same shape as the domain of the objective function"+++-- * Symbolic forward derivative adaptors' testing-only internal machinery++fwdArtifactDelta+  :: forall src ztgt tgt.+     ( AdaptableTarget (AstTensor AstMethodLet FullSpan) src+     , tgt ~ AstTensor AstMethodLet FullSpan ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> FullShapeTK (X src)+  -> (AstArtifactFwd (X src) ztgt, Delta (AstRaw PrimalSpan) ztgt)+       -- ^ the artifact containing the symbolic code of the derivative+{-# INLINE fwdArtifactDelta #-}+fwdArtifactDelta f xftk =+  let g :: AstTensor AstMethodLet FullSpan (X src) -> tgt+      g !arg = ttlet arg $ f . fromTarget+  in fwdArtifactFromForwardPass (forwardPassByInterpretation g emptyEnv) xftk+++-- * Non-symbolic reverse derivative adaptors++-- We are inlining these functions because they take function arguments+-- and are not too large. However, because they are called in many places,+-- we break the inline chain not far from the top, to avoid exe blowup.+--+-- | This simplified version of the concrete (non-symbolic)+-- reverse derivative operation sets the incoming cotangent @dt@ to be 1+-- and assumes the codomain of the function to be differentiated is a scalar.+cgrad+  :: forall src r tgt.+     ( X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal Concrete) src+     , AdaptableTarget Concrete (DValue src)+     , tgt ~ ADVal Concrete (TKScalar r) )+  => (src -> tgt)  -- ^ the objective function+  -> DValue src+  -> DValue src  -- morally DValue (ADTensorKind src)+{-# INLINE cgrad #-}+cgrad f vals = crevMaybe f vals Nothing++-- | This more general version of the concrete (non-symbolic)+-- reverse derivative operation additionally takes the sensitivity parameter+-- (the incoming cotangent).+cvjp+  :: forall src ztgt tgt.+     ( X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal Concrete) src+     , AdaptableTarget Concrete (DValue src)+     , tgt ~ ADVal Concrete ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> DValue src+  -> Concrete (ADTensorKind ztgt)+  -> DValue src  -- morally DValue (ADTensorKind src)+{-# INLINE cvjp #-}+cvjp f vals dt = crevMaybe f vals (Just dt)+++-- * Non-symbolic reverse derivative adaptors' internal machinery++crevMaybe+  :: forall src ztgt tgt.+     ( X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal Concrete) src+     , AdaptableTarget Concrete (DValue src)+     , tgt ~ ADVal Concrete ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> DValue src+  -> Maybe (Concrete (ADTensorKind ztgt))+  -> DValue src  -- morally DValue (ADTensorKind src)+{-# INLINE crevMaybe #-}+crevMaybe f vals0 mdt =+  let valsTarget = toTarget vals0+      g :: ADVal Concrete (X src) -> tgt+      g = f . fromTarget+      xftk = tftkG (knownSTK @(X src)) $ unConcrete valsTarget+  in fromTarget $ fromADTensorKindShared (ftkToSTK xftk)+     $ fst $ crevOnParams mdt g xftk valsTarget+++-- * Non-symbolic forward derivative adaptors++-- | Concrete (non-symbolic) forward derivative operation. It always takes+-- the perturbation parameter, by convention.+cjvp+  :: forall src ztgt tgt.+     ( X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal Concrete) src+     , AdaptableTarget Concrete (DValue src)+     , tgt ~ ADVal Concrete ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> DValue src+  -> DValue src  -- morally (ADTensorKind src)+  -> Concrete (ADTensorKind ztgt)+{-# INLINE cjvp #-}+cjvp f vals ds = fst $ cfwdBoth f vals ds+++-- * Non-symbolic forward derivative adaptors' internal machinery++cfwdBoth+  :: forall src ztgt tgt.+     ( X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal Concrete) src+     , AdaptableTarget Concrete (DValue src)+     , tgt ~ ADVal Concrete ztgt )+  => (src -> tgt)  -- ^ the objective function+  -> DValue src+  -> DValue src  -- morally (ADTensorKind src)+  -> (Concrete (ADTensorKind ztgt), Concrete ztgt)+{-# INLINE cfwdBoth #-}+cfwdBoth f vals0 ds =+  let valsTarget = toTarget vals0+      xftk = tftkG (knownSTK @(X src)) $ unConcrete valsTarget+      g :: ADVal Concrete (X src) -> tgt+      g = f . fromTarget+      dsTarget = toTarget ds+  in if tftkG (ftkToSTK xftk) (unConcrete dsTarget) == xftk+     then cfwdOnParams xftk valsTarget g+          $ toADTensorKindShared xftk dsTarget+     else error "cfwdBoth: forward derivative input must have the same shape as the perturbation argument"++++++-- This specialization is not possible where the functions are defined,+-- due to dependency cycles, but it's possible here:+{-# SPECIALIZE gradientFromDelta :: FullShapeTK x -> FullShapeTK z -> Concrete (ADTensorKind z) -> Delta Concrete z -> Concrete (ADTensorKind x) #-}+{-# SPECIALIZE evalRev :: FullShapeTK y -> EvalState Concrete -> Concrete (ADTensorKind y) -> Delta Concrete y -> EvalState Concrete #-}+{-# SPECIALIZE evalRevFTK :: EvalState Concrete -> Concrete (ADTensorKind y) -> Delta Concrete y -> EvalState Concrete #-}+-- RULE left-hand side too complicated to desugar:+-- {-# SPECIALIZE evalRevSame :: y ~ ADTensorKind y => EvalState Concrete -> Concrete (ADTensorKind y) -> Delta Concrete y -> EvalState Concrete #-}+{-# SPECIALIZE evalRevFromnMap :: EvalState Concrete -> EvalState Concrete #-}++{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKScalar Double) -> Delta Concrete (TKScalar Double) -> EvalState Concrete #-}+{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKScalar Float) -> Delta Concrete (TKScalar Float) -> EvalState Concrete #-}+{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKR n Double) -> Delta Concrete (TKR n Double) -> EvalState Concrete #-}+{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKR n Float) -> Delta Concrete (TKR n Float) -> EvalState Concrete #-}+{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKS sh Double) -> Delta Concrete (TKS sh Double) -> EvalState Concrete #-}+{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKS sh Float) -> Delta Concrete (TKS sh Float) -> EvalState Concrete #-}+{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKX sh Double) -> Delta Concrete (TKX sh Double) -> EvalState Concrete #-}+{-# SPECIALIZE evalRevSame :: EvalState Concrete -> Concrete (TKX sh Float) -> Delta Concrete (TKX sh Float) -> EvalState Concrete #-}+++-- These and all other SPECIALIZE pragmas are needed due to the already+-- mostly fixed issues #21286 and others, even just to compare+-- the output with them and without.+-- This is needed for all three AstSpan values, to handle recursive calls+-- from interpretAstDual, etc.+{-# SPECIALIZE interpretAst+  :: AstEnv (ADVal Concrete)+  -> AstTensor AstMethodLet PrimalSpan y+  -> ADVal Concrete y #-}+{-# SPECIALIZE interpretAst+  :: AstEnv (ADVal (AstRaw PrimalSpan))+  -> AstTensor AstMethodLet PrimalSpan y+  -> ADVal (AstRaw PrimalSpan) y #-}+{-# SPECIALIZE interpretAst+  :: AstEnv Concrete+  -> AstTensor AstMethodLet PrimalSpan y+  -> Concrete y #-}+{-# SPECIALIZE interpretAst+  :: AstEnv (ADVal Concrete)+  -> AstTensor AstMethodLet DualSpan y+  -> ADVal Concrete y #-}+{-# SPECIALIZE interpretAst+  :: AstEnv (ADVal (AstRaw PrimalSpan))+  -> AstTensor AstMethodLet DualSpan y+  -> ADVal (AstRaw PrimalSpan) y #-}+{-# SPECIALIZE interpretAst+  :: AstEnv Concrete+  -> AstTensor AstMethodLet DualSpan y+  -> Concrete y #-}+{-# SPECIALIZE interpretAst+  :: AstEnv (ADVal Concrete)+  -> AstTensor AstMethodLet FullSpan y+  -> ADVal Concrete y #-}+{-# SPECIALIZE interpretAst+  :: AstEnv (ADVal (AstRaw PrimalSpan))+  -> AstTensor AstMethodLet FullSpan y+  -> ADVal (AstRaw PrimalSpan) y #-}+{-# SPECIALIZE interpretAst+  :: AstEnv Concrete+  -> AstTensor AstMethodLet FullSpan y+  -> Concrete y #-}++{-# SPECIALIZE interpretAstPrimal+  :: AstEnv (ADVal Concrete)+  -> AstTensor AstMethodLet PrimalSpan y+  -> Concrete y #-}+{-# SPECIALIZE interpretAstPrimal+  :: AstEnv (ADVal (AstRaw PrimalSpan))+  -> AstTensor AstMethodLet PrimalSpan y+  -> AstRaw PrimalSpan y #-}+{-# SPECIALIZE interpretAstPrimal+  :: AstEnv Concrete+  -> AstTensor AstMethodLet PrimalSpan y+  -> Concrete y #-}++{-# SPECIALIZE interpretAstBool+  :: AstEnv (ADVal Concrete)+  -> AstBool AstMethodLet+  -> Bool #-}+{-# SPECIALIZE interpretAstBool+  :: AstEnv (ADVal (AstRaw PrimalSpan))+  -> AstBool AstMethodLet+  -> AstBool AstMethodShare #-}+{-# SPECIALIZE interpretAstBool+  :: AstEnv Concrete+  -> AstBool AstMethodLet+  -> Bool #-}
+ src/HordeAd/AstEngine.hs view
@@ -0,0 +1,149 @@+-- | Predefined common functions for simplification and pretty-printing of AST.+module HordeAd.AstEngine+  ( -- * The joint inlining and simplification term transformations+    simplifyArtifact, simplifyArtifactGradient, simplifyArtifactDerivative+  , simplifyInline, simplifyInlineContract, simplifyInlineContractNoExpand+    -- * Pretty-printing terms in a few useful configurations+  , printAstVarName+  , printAstSimple, printAstPretty, printAstPrettyButNested+  , printArtifactSimple, printArtifactPretty+  , printArtifactPrimalSimple, printArtifactPrimalPretty+  ) where++import Prelude++import Data.EnumMap.Strict qualified as EM+import Data.IntMap.Strict (IntMap)+import Data.IntMap.Strict qualified as IM++import HordeAd.Core.Ast+import HordeAd.Core.AstInline+import HordeAd.Core.AstPrettyPrint+import HordeAd.Core.AstTraverse++-- * The joint inlining and simplification term transformation++-- | Simplify the whole reverse derivative artifact (which includes+-- also the primal value computed during the differentiation process).+simplifyArtifact :: forall x z.+                    AstArtifactRev x z -> AstArtifactRev x z+simplifyArtifact art =+  let !der = simplifyInlineContract $ artDerivativeRev art in+  let !prim = simplifyInlineContract $ artPrimalRev art+  in art {artDerivativeRev = der, artPrimalRev = prim}++-- | Simplify only the gradient in the reverse derivative artifact.+simplifyArtifactGradient :: forall x z.+                            AstArtifactRev x z -> AstArtifactRev x z+simplifyArtifactGradient art =+  art { artDerivativeRev =+        simplifyInlineContract $ artDerivativeRev art }++-- | Simplify only the derivative in the forward derivative artifact.+simplifyArtifactDerivative :: forall x z.+                              AstArtifactFwd x z -> AstArtifactFwd x z+simplifyArtifactDerivative art =+  art { artDerivativeFwd =+        simplifyInlineContract $ artDerivativeFwd art }++-- | A mixture of simplification and inlining to use when the resultng+-- term is not yet supposed to be interpreted using a computational backed,+-- but rather to be stored and later composed with other terms.+simplifyInline+  :: forall z s. AstSpan s+  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z+simplifyInline =+  simplifyAst . expandAst . snd . inlineAst EM.empty+  . simplifyAst . expandAst . snd . inlineAst EM.empty+  . simplifyAst++-- | A mixture of simplification, inlining and recognition of additional+-- backend-specific primitives, to be used just before a term+-- is interpreted as a value in the computational backend.+simplifyInlineContract+  :: forall z s. AstSpan s+  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z+simplifyInlineContract =+  contractAst . expandAst . snd . inlineAst EM.empty+  . simplifyAst . expandAst . snd . inlineAst EM.empty+  . simplifyAst++simplifyInlineContractNoExpand+  :: forall z s. AstSpan s+  => AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z+simplifyInlineContractNoExpand =+  contractAst . simplifyAst . snd . inlineAst EM.empty+  . simplifyAst . snd . inlineAst EM.empty+  . simplifyAst+++-- * Pretty-printing terms in a few useful configurations++printAstVarRename :: AstSpan s+                  => IntMap String -> AstVarName s y -> String+printAstVarRename renames var =+  printAstVar (defaulPrintConfig {varRenames = renames}) var ""++printAstSimpleRename :: AstSpan s+                     => IntMap String -> AstTensor ms s y -> String+printAstSimpleRename renames t =+  printAst+    (defaulPrintConfig {loseRoudtrip = False, varRenames = renames}) 0 t ""++printAstPrettyRename :: AstSpan s+                     => IntMap String -> AstTensor ms s y -> String+printAstPrettyRename renames t =+  printAst (defaulPrintConfig {varRenames = renames}) 0 t ""++printAstVarName :: AstSpan s+                => AstVarName s y -> String+printAstVarName var =+  printAstVar defaulPrintConfig var ""++-- | Print an AST term in a form close to being able to roundtrip,+-- including explicit sharing preservation.+printAstSimple :: AstSpan s+               => AstTensor ms s y -> String+printAstSimple t =+  printAst (defaulPrintConfig {loseRoudtrip = False}) 0 t ""++-- | Print an AST term in a readable form that does not roundtrip,+-- and where Haskell @let@ (sharing on Haskell heap) is used instead+-- of explicit sharing of subterms.+printAstPretty :: AstSpan s+               => AstTensor ms s y -> String+printAstPretty t =+  printAst defaulPrintConfig 0 t ""++printAstPrettyButNested :: AstSpan s+                        => AstTensor ms s y -> String+printAstPrettyButNested t =+  printAst (defaulPrintConfig {ignoreNestedLambdas = False}) 0 t ""++printArtifactSimple :: AstArtifactRev x z -> String+printArtifactSimple AstArtifactRev{..} =+  let nDt = fromEnum (varNameToAstVarId artVarDtRev) - 100000000+      renames = IM.singleton nDt "dret"+      varsPP = [ printAstVarRename renames artVarDtRev+               , printAstVarRename renames artVarDomainRev ]+  in "\\" ++ unwords varsPP+          ++ " -> " ++ printAstSimpleRename renames artDerivativeRev++printArtifactPretty :: AstArtifactRev x z -> String+printArtifactPretty AstArtifactRev{..} =+  let nDt = fromEnum (varNameToAstVarId artVarDtRev) - 100000000+      renames = IM.singleton nDt "dret"+      varsPP = [ printAstVarRename renames artVarDtRev+               , printAstVarRename renames artVarDomainRev ]+  in "\\" ++ unwords varsPP+          ++ " -> " ++ printAstPrettyRename renames artDerivativeRev++printArtifactPrimalSimple :: AstArtifactRev x z -> String+printArtifactPrimalSimple AstArtifactRev{..} =+  "\\" ++ printAstVarName artVarDomainRev+       ++ " -> " ++ printAstSimple artPrimalRev++printArtifactPrimalPretty :: AstArtifactRev x z -> String+printArtifactPrimalPretty AstArtifactRev{..} =+  "\\" ++ printAstVarName artVarDomainRev+       ++ " -> " ++ printAstPretty artPrimalRev
+ src/HordeAd/Core/Adaptor.hs view
@@ -0,0 +1,510 @@+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Adaptors for working with types of collections of tensors,+-- e.g., tuples, sized lists and user types of statically known size,+-- as long as they have the proper instances defined.+-- The collections are used as representations of the domains+-- of objective functions that become the codomains of the reverse+-- derivative functions and also to handle multiple arguments+-- and results of fold-like operations.+module HordeAd.Core.Adaptor+  ( AdaptableTarget(..), TermValue(..), DualNumberValue(..)+  , ForgetShape(..), RandomValue(..)+  , stkOfListR+    -- * Helper classes and types+  , Tups, NoShapeTensorKind+  ) where++import Prelude++import Data.Default+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:))+import Data.Vector.Generic qualified as V+import Data.Vector.Strict qualified as Data.Vector+import GHC.TypeLits (KnownNat, OrderingI (..), cmpNat, type (-), type (<=?))+import System.Random++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.Ast+import HordeAd.Core.CarriersADVal+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.OpsAst ()+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * Adaptor classes++-- Inspired by adaptors from @tomjaguarpaw's branch.+--+-- | The class that makes it possible to treat @vals@ (e.g., a tuple of tensors)+-- as a @target@-based (e.g., concrete or symbolic) value+-- of tensor kind @X vals@.+class AdaptableTarget (target :: Target) vals where+  type X vals :: TK  -- ^ what tensor kind represents the collection+  toTarget :: vals -> target (X vals)+    -- ^ represent a collection of tensors+  fromTarget :: target (X vals) -> vals+    -- ^ recovers a collection of tensors from its canonical representation;+    --   requires a duplicable argument++-- | An embedding of a concrete collection of tensors to a non-concrete+-- counterpart of the same shape and containing the same data.+class TermValue vals where+  type Value vals = result | result -> vals+    -- ^ a helper type, with the same general shape,+    -- but possibly more concrete, e.g., arrays instead of terms,+    -- where the injectivity is crucial to limit the number+    -- of type applications the library user has to supply+  fromValue :: Value vals -> vals  -- ^ an embedding++-- | An embedding of a concrete collection of tensors to a non-concrete+-- counterpart of the same shape and containing the same data.+-- This variant is possible to define more often, but the associated+-- type family is not injective.+class DualNumberValue vals where+  type DValue vals+    -- ^ a helper type, with the same general shape,+    -- but possibly more concrete, e.g., arrays instead of terms,+    -- where the injectivity is hard to obtain, but is not so important,+    -- because the type is not used in the best pipeline+  fromDValue :: DValue vals -> vals  -- ^ an embedding++-- | A helper class for for converting all tensors inside a type+-- from shaped to ranked. It's useful when a collection of parameters+-- is defined as shaped tensor for 'RandomValue' but then is going+-- to be used as ranked tensor to make type reconstruction easier.+class ForgetShape vals where+  type NoShape vals+  forgetShape :: vals -> NoShape vals++-- | A helper class for randomly generating initial parameters.+-- Only instance for collections of shaped tensors and scalars are possible,+-- because only then the shapes of the tensors to generate are known+-- from their types.+class RandomValue vals where+  randomValue :: Double -> StdGen -> (vals, StdGen)+++-- * Base instances++instance AdaptableTarget target (target y) where+  type X (target y) = y+  toTarget = id+  fromTarget t = t+  {-# SPECIALIZE instance AdaptableTarget Concrete (Concrete (TKS sh Double)) #-}+  {-# SPECIALIZE instance AdaptableTarget Concrete (Concrete (TKS sh Float)) #-}+    -- a failed attempt to specialize without -fpolymorphic-specialisation++instance KnownSTK y+         => TermValue (AstTensor AstMethodLet FullSpan y) where+  type Value (AstTensor AstMethodLet FullSpan y) = Concrete y+  fromValue t = tconcrete (tftkG (knownSTK @y) $ unConcrete t) t++instance (BaseTensor target, BaseTensor (PrimalOf target), KnownSTK y)+         => DualNumberValue (target y) where+  type DValue (target y) = Concrete y+  fromDValue t = tfromPrimal (knownSTK @y)+                 $ tconcrete (tftkG (knownSTK @y) $ unConcrete t) t++instance ForgetShape (target (TKScalar r)) where+  type NoShape (target (TKScalar r)) = target (TKScalar r)+  forgetShape = id++instance ForgetShape (target (TKR n r)) where+  type NoShape (target (TKR n r)) = target (TKR n r)+  forgetShape = id++instance (KnownShS sh, GoodScalar r, ConvertTensor target)+         => ForgetShape (target (TKS sh r)) where+  type NoShape (target (TKS sh r)) = target (TKR (Rank sh) r)+  forgetShape = rfromS++instance ForgetShape (target (TKX sh r)) where+  type NoShape (target (TKX sh r)) = target (TKX sh r)+  forgetShape = id++type family NoShapeTensorKind tk where+  NoShapeTensorKind (TKScalar r) = TKScalar r+  NoShapeTensorKind (TKR2 n r) = TKR2 n r+  NoShapeTensorKind (TKS2 sh r) = TKR2 (Rank sh) r+  NoShapeTensorKind (TKX2 sh r) = TKX2 sh r+  NoShapeTensorKind (TKProduct y z) =+    TKProduct (NoShapeTensorKind y) (NoShapeTensorKind z)++instance ( ForgetShape (target a)+         , ForgetShape (target b)+         , target (NoShapeTensorKind a) ~ NoShape (target a)+         , target (NoShapeTensorKind b) ~ NoShape (target b)+         , BaseTensor target, LetTensor target )+         => ForgetShape (target (TKProduct a b)) where+  type NoShape (target (TKProduct a b)) =+    target (NoShapeTensorKind (TKProduct a b))+  forgetShape ab =+    ttlet ab $ \abShared ->+      tpair (forgetShape (tproject1 abShared))+            (forgetShape (tproject2 abShared))++instance forall r target. (GoodScalar r, BaseTensor target)+         => RandomValue (target (TKScalar r)) where+  randomValue range g =+    ifDifferentiable @r+      (let (r, g2) = random g+           m = 2 * realToFrac range * (r - 0.5)+       in (tkconcrete m, g2))+      (tkconcrete def, g)++instance forall sh r target. (KnownShS sh, GoodScalar r, BaseTensor target)+         => RandomValue (target (TKS sh r)) where+  randomValue range g =+    ifDifferentiable @r+      (let createRandomVector :: Int -> StdGen -> target (TKS sh r)+           createRandomVector n seed =+             srepl (2 * realToFrac range)+             * (tsconcrete+                  (Nested.sfromVector knownShS (V.fromListN n (randoms seed)))+                - srepl 0.5)+           (g1, g2) = splitGen g+           arr = createRandomVector (shsSize (knownShS @sh)) g1+       in (arr, g2))+      (srepl def, g)+   where srepl = tsconcrete . Nested.sreplicateScal knownShS+  -- {-# SPECIALIZE instance (KnownShS sh, GoodScalar r, Fractional r, Random r) => RandomValue (Concrete (TKS sh r)) #-}+  {-# SPECIALIZE instance KnownShS sh => RandomValue (Concrete (TKS sh Double)) #-}+  {-# SPECIALIZE instance KnownShS sh => RandomValue (Concrete (TKS sh Float)) #-}++instance (RandomValue (target a), RandomValue (target b), BaseTensor target)+         => RandomValue (target (TKProduct a b)) where+  randomValue range g =+    let (v1, g1) = randomValue range g+        (v2, g2) = randomValue range g1+    in (tpair v1 v2, g2)++-- These instances are messy and hard to use, but we probably can't do better.+instance DualNumberValue Double where+  type DValue Double = Concrete (TKScalar Double)+  fromDValue (Concrete d) = d++instance DualNumberValue Float where+  type DValue Float = Concrete (TKScalar Float)+  fromDValue (Concrete d) = d++instance TermValue (Concrete (TKScalar Double)) where+  type Value (Concrete (TKScalar Double)) = Double+  fromValue = Concrete++instance TermValue (Concrete (TKScalar Float)) where+  type Value (Concrete (TKScalar Float)) = Float+  fromValue = Concrete+++-- * Compound instances++instance (BaseTensor target, ConvertTensor target, GoodScalar r)+         => AdaptableTarget target [target (TKScalar r)] where+  type X [target (TKScalar r)] = TKR 1 r+  toTarget l = if null l+               then trconcrete Nested.remptyArray+               else trfromVector $ V.fromList $ map rfromK l+  fromTarget = map kfromR . trunravelToList+                              -- inefficient, but we probably can't do better++instance (BaseTensor target, ConvertTensor target, GoodScalar r)+         => AdaptableTarget target+                            (Data.Vector.Vector (target (TKScalar r))) where+  type X (Data.Vector.Vector (target (TKScalar r))) = TKR 1 r+  toTarget v = if V.null v+               then trconcrete Nested.remptyArray+               else trfromVector $ V.map rfromK v+  fromTarget =+    V.fromList . map kfromR . trunravelToList+                                -- inefficient, but we probably can't do better++type family Tups n t where+  Tups 0 t = TKUnit+  Tups n t = TKProduct t (Tups (n - 1) t)++stkOfListR :: forall t n.+              SingletonTK t -> SNat n -> SingletonTK (Tups n t)+stkOfListR _ (SNat' @0) = stkUnit+stkOfListR stk SNat =+  gcastWith (unsafeCoerceRefl :: (1 <=? n) :~: True) $+  gcastWith (unsafeCoerceRefl :: Tups n t :~: TKProduct t (Tups (n - 1) t)) $+  STKProduct stk (stkOfListR stk (SNat @(n - 1)))++instance (BaseTensor target, KnownNat n, AdaptableTarget target a)+         => AdaptableTarget target (ListR n a) where+  type X (ListR n a) = Tups n (X a)+  toTarget ZR = tkconcrete Z1+  toTarget ((:::) @n1 a rest) =+    gcastWith (unsafeCoerceRefl+               :: X (ListR n a) :~: TKProduct (X a) (X (ListR n1 a))) $+    let a1 = toTarget a+        rest1 = toTarget rest+    in tpair a1 rest1+  fromTarget tups = case SNat @n of+    SNat' @0 -> ZR+    _ ->+      gcastWith (unsafeCoerceRefl :: (1 <=? n) :~: True) $+      gcastWith (unsafeCoerceRefl+                 :: X (ListR n a) :~: TKProduct (X a) (X (ListR (n - 1) a))) $+      let (a1, rest1) = (tproject1 tups, tproject2 tups)+          a = fromTarget a1+          rest = fromTarget rest1+      in (a ::: rest)+  {-# SPECIALIZE instance (KnownNat n, AdaptableTarget (AstTensor AstMethodLet FullSpan) a) => AdaptableTarget (AstTensor AstMethodLet FullSpan) (ListR n a) #-}+  {-# SPECIALIZE instance (KnownNat n, AdaptableTarget (ADVal Concrete) a) => AdaptableTarget (ADVal Concrete) (ListR n a) #-}++instance TermValue a => TermValue [a] where+  type Value [a] = [Value a]+  fromValue = map fromValue++instance TermValue a => TermValue (Data.Vector.Vector a) where+  type Value (Data.Vector.Vector a) = Data.Vector.Vector (Value a)+  fromValue = V.map fromValue++instance TermValue a => TermValue (ListR n a) where+  type Value (ListR n a) = ListR n (Value a)+  fromValue ZR = ZR+  fromValue (a ::: rest) = fromValue a ::: fromValue rest++instance DualNumberValue a => DualNumberValue [a] where+  type DValue [a] = [DValue a]+  fromDValue = map fromDValue++instance DualNumberValue a => DualNumberValue (Data.Vector.Vector a) where+  type DValue (Data.Vector.Vector a) = Data.Vector.Vector (DValue a)+  fromDValue = V.map fromDValue++instance DualNumberValue a => DualNumberValue (ListR n a) where+  type DValue (ListR n a) = ListR n (DValue a)+  fromDValue ZR = ZR+  fromDValue (a ::: rest) = fromDValue a ::: fromDValue rest++instance ForgetShape [a] where+  type NoShape [a] = [a]+  forgetShape = id++instance ForgetShape (Data.Vector.Vector a) where+  type NoShape (Data.Vector.Vector a) = Data.Vector.Vector a+  forgetShape = id++instance ForgetShape a => ForgetShape (ListR n a) where+  type NoShape (ListR n a) = ListR n (NoShape a)+  forgetShape ZR = ZR+  forgetShape (a ::: rest) = forgetShape a ::: forgetShape rest++instance (RandomValue a, KnownNat n) => RandomValue (ListR n a) where+  randomValue range g = case cmpNat (Proxy @n) (Proxy @0)  of+    LTI -> error "randomValue: impossible"+    EQI -> (ZR, g)+    GTI -> gcastWith (unsafeCoerceRefl :: (1 <=? n) :~: True) $+           let (v, g1) = randomValue range g+               (rest, g2) = randomValue @(ListR (n - 1) a) range g1+           in (v ::: rest, g2)+++-- * Tuple instances++instance ( BaseTensor target+         , AdaptableTarget target a+         , AdaptableTarget target b )+         => AdaptableTarget target (a, b) where+  type X (a, b) = TKProduct (X a) (X b)+  toTarget (a, b) =+    let a1 = toTarget a+        b1 = toTarget b+    in tpair a1 b1+  fromTarget ab =+    let a = fromTarget $ tproject1 ab+        b = fromTarget $ tproject2 ab+    in (a, b)+  {-# SPECIALIZE instance (AdaptableTarget (AstTensor AstMethodLet FullSpan) a, AdaptableTarget (AstTensor AstMethodLet FullSpan) b) => AdaptableTarget (AstTensor AstMethodLet FullSpan) (a, b) #-}++instance (TermValue a, TermValue b) => TermValue (a, b) where+  type Value (a, b) = (Value a, Value b)+  fromValue (va, vb) = (fromValue va, fromValue vb)++instance (DualNumberValue a, DualNumberValue b) => DualNumberValue (a, b) where+  type DValue (a, b) = (DValue a, DValue b)+  fromDValue (va, vb) = (fromDValue va, fromDValue vb)++instance ( ForgetShape a+         , ForgetShape b ) => ForgetShape (a, b) where+  type NoShape (a, b) = (NoShape a, NoShape b)+  forgetShape (a, b) = (forgetShape a, forgetShape b)++instance ( RandomValue a+         , RandomValue b ) => RandomValue (a, b) where+  randomValue range g =+    let (v1, g1) = randomValue range g+        (v2, g2) = randomValue range g1+    in ((v1, v2), g2)++instance ( BaseTensor target+         , AdaptableTarget target a+         , AdaptableTarget target b+         , AdaptableTarget target c )+         => AdaptableTarget target (a, b, c) where+  type X (a, b, c) = TKProduct (TKProduct (X a) (X b)) (X c)+  toTarget (a, b, c) =+    let a1 = toTarget a+        b1 = toTarget b+        c1 = toTarget c+    in tpair (tpair a1 b1) c1+  fromTarget abc =+    let a = fromTarget $ tproject1 $ tproject1 abc+        b = fromTarget $ tproject2 $ tproject1 abc+        c = fromTarget $ tproject2 abc+    in (a, b, c)+  {-# SPECIALIZE instance (AdaptableTarget (AstTensor AstMethodLet FullSpan) a, AdaptableTarget (AstTensor AstMethodLet FullSpan) b, AdaptableTarget (AstTensor AstMethodLet FullSpan) c) => AdaptableTarget (AstTensor AstMethodLet FullSpan) (a, b, c) #-}++instance (TermValue a, TermValue b, TermValue c)+         => TermValue (a, b, c) where+  type Value (a, b, c) = (Value a, Value b, Value c)+  fromValue (va, vb, vc) = (fromValue va, fromValue vb, fromValue vc)++instance (DualNumberValue a, DualNumberValue b, DualNumberValue c)+         => DualNumberValue (a, b, c) where+  type DValue (a, b, c) = (DValue a, DValue b, DValue c)+  fromDValue (va, vb, vc) = (fromDValue va, fromDValue vb, fromDValue vc)++instance ( ForgetShape a+         , ForgetShape b+         , ForgetShape c ) => ForgetShape (a, b, c) where+  type NoShape (a, b, c) = (NoShape a, NoShape b, NoShape c)+  forgetShape (a, b, c) = (forgetShape a, forgetShape b, forgetShape c)++instance ( RandomValue a+         , RandomValue b+         , RandomValue c ) => RandomValue (a, b, c) where+  randomValue range g =+    let (v1, g1) = randomValue range g+        (v2, g2) = randomValue range g1+        (v3, g3) = randomValue range g2+    in ((v1, v2, v3), g3)++instance ( BaseTensor target+         , AdaptableTarget target a+         , AdaptableTarget target b+         , AdaptableTarget target c+         , AdaptableTarget target d)+         => AdaptableTarget target (a, b, c, d) where+  type X (a, b, c, d) = TKProduct (TKProduct (X a) (X b))+                                  (TKProduct (X c) (X d))+  toTarget (a, b, c, d) =+    let a1 = toTarget a+        b1 = toTarget b+        c1 = toTarget c+        d1 = toTarget d+    in  tpair (tpair a1 b1) (tpair c1 d1)+  fromTarget abcd =+    let a = fromTarget $ tproject1 $ tproject1 abcd+        b = fromTarget $ tproject2 $ tproject1 abcd+        c = fromTarget $ tproject1 $ tproject2 abcd+        d = fromTarget $ tproject2 $ tproject2 abcd+    in (a, b, c, d)+  {-# SPECIALIZE instance (AdaptableTarget (AstTensor AstMethodLet FullSpan) a, AdaptableTarget (AstTensor AstMethodLet FullSpan) b, AdaptableTarget (AstTensor AstMethodLet FullSpan) c, AdaptableTarget (AstTensor AstMethodLet FullSpan) d) => AdaptableTarget (AstTensor AstMethodLet FullSpan) (a, b, c, d) #-}++instance (TermValue a, TermValue b, TermValue c, TermValue d)+         => TermValue (a, b, c, d) where+  type Value (a, b, c, d) = (Value a, Value b, Value c, Value d)+  fromValue (va, vb, vc, vd) =+    (fromValue va, fromValue vb, fromValue vc, fromValue vd)++instance ( DualNumberValue a, DualNumberValue b, DualNumberValue c+         , DualNumberValue d )+         => DualNumberValue (a, b, c, d) where+  type DValue (a, b, c, d) = (DValue a, DValue b, DValue c, DValue d)+  fromDValue (va, vb, vc, vd) =+    (fromDValue va, fromDValue vb, fromDValue vc, fromDValue vd)++instance ( ForgetShape a+         , ForgetShape b+         , ForgetShape c+         , ForgetShape d ) => ForgetShape (a, b, c, d) where+  type NoShape (a, b, c, d) =+    (NoShape a, NoShape b, NoShape c, NoShape d)+  forgetShape (a, b, c, d) =+    (forgetShape a, forgetShape b, forgetShape c, forgetShape d)++instance ( RandomValue a+         , RandomValue b+         , RandomValue c+         , RandomValue d ) => RandomValue (a, b, c, d) where+  randomValue range g =+    let (v1, g1) = randomValue range g+        (v2, g2) = randomValue range g1+        (v3, g3) = randomValue range g2+        (v4, g4) = randomValue range g3+    in ((v1, v2, v3, v4), g4)++instance ( BaseTensor target+         , AdaptableTarget target a+         , AdaptableTarget target b+         , AdaptableTarget target c+         , AdaptableTarget target d+         , AdaptableTarget target e)+         => AdaptableTarget target (a, b, c, d, e) where+  type X (a, b, c, d, e) = TKProduct (TKProduct (TKProduct (X a) (X b)) (X c))+                                     (TKProduct (X d) (X e))+  toTarget (a, b, c, d, e) =+    let a1 = toTarget a+        b1 = toTarget b+        c1 = toTarget c+        d1 = toTarget d+        e1 = toTarget e+    in tpair (tpair (tpair a1 b1) c1) (tpair d1 e1)+  fromTarget abcde =+    let a = fromTarget $ tproject1 $ tproject1 $ tproject1 abcde+        b = fromTarget $ tproject2 $ tproject1 $ tproject1 abcde+        c = fromTarget $ tproject2 $ tproject1 abcde+        d = fromTarget $ tproject1 $ tproject2 abcde+        e = fromTarget $ tproject2 $ tproject2 abcde+    in (a, b, c, d, e)+  {-# SPECIALIZE instance (AdaptableTarget (AstTensor AstMethodLet FullSpan) a, AdaptableTarget (AstTensor AstMethodLet FullSpan) b, AdaptableTarget (AstTensor AstMethodLet FullSpan) c, AdaptableTarget (AstTensor AstMethodLet FullSpan) d, AdaptableTarget (AstTensor AstMethodLet FullSpan) e) => AdaptableTarget (AstTensor AstMethodLet FullSpan) (a, b, c, d, e) #-}++instance (TermValue a, TermValue b, TermValue c, TermValue d, TermValue e)+         => TermValue (a, b, c, d, e) where+  type Value (a, b, c, d, e) = (Value a, Value b, Value c, Value d, Value e)+  fromValue (va, vb, vc, vd, ve) =+    (fromValue va, fromValue vb, fromValue vc, fromValue vd, fromValue ve)++instance ( DualNumberValue a, DualNumberValue b, DualNumberValue c+         , DualNumberValue d, DualNumberValue e )+         => DualNumberValue (a, b, c, d, e) where+  type DValue (a, b, c, d, e) =+    (DValue a, DValue b, DValue c, DValue d, DValue e)+  fromDValue (va, vb, vc, vd, ve) =+    (fromDValue va, fromDValue vb, fromDValue vc, fromDValue vd, fromDValue ve)++instance ( ForgetShape a+         , ForgetShape b+         , ForgetShape c+         , ForgetShape d+         , ForgetShape e ) => ForgetShape (a, b, c, d, e) where+  type NoShape (a, b, c, d, e) =+    (NoShape a, NoShape b, NoShape c, NoShape d, NoShape e)+  forgetShape (a, b, c, d, e) =+    (forgetShape a, forgetShape b, forgetShape c, forgetShape d, forgetShape e)++instance ( RandomValue a+         , RandomValue b+         , RandomValue c+         , RandomValue d+         , RandomValue e ) => RandomValue (a, b, c, d, e) where+  randomValue range g =+    let (v1, g1) = randomValue range g+        (v2, g2) = randomValue range g1+        (v3, g3) = randomValue range g2+        (v4, g4) = randomValue range g3+        (v5, g5) = randomValue range g4+    in ((v1, v2, v3, v4, v5), g5)
+ src/HordeAd/Core/Ast.hs view
@@ -0,0 +1,498 @@+{-# LANGUAGE ViewPatterns #-}+-- | AST of corresponding to the horde-ad operations specified+-- in the 'HordeAd.Core.Ops.BaseTensor' class and others.+-- The AST is essential for efficient handling of second order operations+-- such as build and map via BOT (bulk-operation transformation),+-- and fold and mapAccum via symbolic nested derivatives.+-- It also permits producing reusable reverse derivative terms,+-- which can be simplified, fused, inlined once and then+-- interpreted many times.+--+-- Note that @Ast*@ modules rarely depend on @Ops*@ and @Carriers*@ modules+-- (except for "HordeAd.Core.AstInterpret" and "HordeAd.Core.AstEnv"+-- that describe how to go from @Ast*@ to @Ops*@). Similarly, @Ops*@+-- and @Carriers*@ modules rarely depend on @Ast*@ modules+-- (except for "HordeAd.Core.OpsAst" and "HordeAd.Core.CarriersAst"+-- that describe how to define @Ops*@ in terms of @Ast*@).+-- Syntax is relatively separated from semantics and they meet+-- in the interpreter ("HordeAd.Core.AstInterpret")+-- and in the semantic model constructed from syntax ("HordeAd.Core.OpsAst").+--+-- (A copy of the text above is in "HordeAd.Core.Ops".)+module HordeAd.Core.Ast+  ( -- * The AstSpan tags and class+    AstSpanType(..), AstSpan(..), sameAstSpan+    -- * Variables and related types+  , AstVarId, intToAstVarId+  , AstInt, IntVarName, pattern AstIntVar+  , AstVarName, mkAstVarName, varNameToAstVarId, varNameToFTK, varNameToBounds+  , astVar+  , AstArtifactRev(..), AstArtifactFwd(..)+  , AstIxS, AstVarListS, pattern AstLeqInt+    -- * AST+  , AstMethodOfSharing(..), AstTensor(..)+  , AstHFun(..)+  , AstBool(..), OpCodeNum1(..), OpCode1(..), OpCode2(..), OpCodeIntegral2(..)+  ) where++import Prelude hiding (foldl')++import Data.Dependent.EnumMap.Strict qualified as DMap+import Data.Functor.Const+import Data.Int (Int64)+import Data.Kind (Type)+import Data.Some+import Data.Type.Equality (TestEquality (..), (:~:) (Refl))+import Data.Vector.Strict qualified as Data.Vector+import GHC.TypeLits (type (+), type (<=))+import Type.Reflection (Typeable, typeRep)++import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (Init)++import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * The AstSpan tags and class++-- | A kind (a type intended to be promoted) marking whether an AST term+-- is supposed to denote the primal part of a dual number, the dual part+-- or the whole dual number. It's mainly used to index the terms+-- of the AstTensor type  and related GADTs.+type data AstSpanType = PrimalSpan | DualSpan | FullSpan++class Typeable s => AstSpan (s :: AstSpanType) where+  fromPrimal :: AstTensor ms PrimalSpan y -> AstTensor ms s y+  fromDual :: AstTensor ms DualSpan y -> AstTensor ms s y+  primalPart :: AstTensor ms s y -> AstTensor ms PrimalSpan y+  dualPart :: AstTensor ms s y -> AstTensor ms DualSpan y++-- These are weak instance and we can't move them to AstSimplify,+-- because it's too late and also astPrimalPart only works on AstMethodLet.+instance AstSpan PrimalSpan where+  fromPrimal = id+  fromDual t = AstPrimalPart $ AstFromDual t  -- this is primal zero+  primalPart t = t+  dualPart t = AstDualPart $ AstFromPrimal t  -- this is dual zero++instance AstSpan DualSpan where+  fromPrimal t = AstDualPart $ AstFromPrimal t  -- this is dual zero+  fromDual = id+  primalPart t = AstPrimalPart $ AstFromDual t  -- this is primal zero+  dualPart t = t++instance AstSpan FullSpan where+  fromPrimal = AstFromPrimal+  fromDual = AstFromDual+  primalPart (AstFromPrimal t) = t+  primalPart t = AstPrimalPart t+  dualPart (AstFromDual t) = t+  dualPart t = AstDualPart t++sameAstSpan :: forall s1 s2. (AstSpan s1, AstSpan s2) => Maybe (s1 :~: s2)+sameAstSpan = testEquality (typeRep @s1) (typeRep @s2)+++-- * Variables and related types++newtype AstVarId = AstVarId Int+ deriving (Eq, Ord, Show, Enum)++intToAstVarId :: Int -> AstVarId+intToAstVarId = AstVarId++type role AstVarName phantom nominal+data AstVarName :: AstSpanType -> TK -> Type where+  AstVarName :: forall s y.+                FullShapeTK y -> Int64 -> Int64 -> AstVarId+             -> AstVarName s y++instance Eq (AstVarName s y) where+  AstVarName _ _ _ varId1 == AstVarName _ _ _ varId2 = varId1 == varId2++instance Show (AstVarName s y) where+  showsPrec d (AstVarName _ _ _ varId) =+    showsPrec d varId  -- less verbose, more readable++instance DMap.Enum1 (AstVarName s) where+  type Enum1Info (AstVarName s) = Some FtkAndBounds+  fromEnum1 (AstVarName ftk minb maxb varId) =+    (fromEnum varId, Some (FtkAndBounds ftk minb maxb))+  toEnum1 varIdInt (Some (FtkAndBounds ftk minb maxb)) =+    Some $ AstVarName ftk minb maxb $ toEnum varIdInt++type role FtkAndBounds nominal+data FtkAndBounds y = FtkAndBounds (FullShapeTK y) Int64 Int64++instance TestEquality (AstVarName s) where+  testEquality (AstVarName ftk1 _ _ _) (AstVarName ftk2 _ _ _) =+    matchingFTK ftk1 ftk2++mkAstVarName :: forall s y.+                FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId+             -> AstVarName s y+mkAstVarName ftk Nothing = AstVarName ftk (-1000000000) 1000000000+mkAstVarName ftk (Just (minb, maxb)) = AstVarName ftk minb maxb++varNameToAstVarId :: AstVarName s y -> AstVarId+varNameToAstVarId (AstVarName _ _ _ varId) = varId++varNameToFTK :: AstVarName s y -> FullShapeTK y+varNameToFTK (AstVarName ftk _ _ _) = ftk++varNameToBounds :: AstVarName s y -> Maybe (Int64, Int64)+varNameToBounds (AstVarName _ minb maxb _) =+  if minb == -1000000000 && maxb == 1000000000+  then Nothing+  else Just (minb, maxb)++astVar :: AstSpan s+       => AstVarName s y -> AstTensor ms s y+astVar (AstVarName (FTKScalar @r) lb ub _)+  | lb == ub+  , Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =+    fromPrimal $ AstConcreteK lb+astVar varName = AstVar varName++-- | The reverse derivative artifact.+type role AstArtifactRev nominal nominal+data AstArtifactRev x z = AstArtifactRev+  { artVarDtRev      :: AstVarName PrimalSpan (ADTensorKind z)+  , artVarDomainRev  :: AstVarName PrimalSpan x+  , artDerivativeRev :: AstTensor AstMethodLet PrimalSpan (ADTensorKind x)+  , artPrimalRev     :: AstTensor AstMethodLet PrimalSpan z+  }+ deriving Show++-- | The forward derivative artifact.+type role AstArtifactFwd nominal nominal+data AstArtifactFwd x z = AstArtifactFwd+  { artVarDsFwd      :: AstVarName PrimalSpan (ADTensorKind x)+  , artVarDomainFwd  :: AstVarName PrimalSpan x+  , artDerivativeFwd :: AstTensor AstMethodLet PrimalSpan (ADTensorKind z)+  , artPrimalFwd     :: AstTensor AstMethodLet PrimalSpan z+  }+ deriving Show++-- | This is the (arbitrarily) chosen representation of terms denoting+-- integers in the indexes of tensor operations.+type AstInt ms = AstTensor ms PrimalSpan (TKScalar Int64)+-- ~ IntOf (AstTensor ms FullSpan)++type IntVarName = AstVarName PrimalSpan (TKScalar Int64)++pattern AstIntVar :: IntVarName -> AstInt ms+pattern AstIntVar var <- AstVar var++-- Data invariant: the var names have bounds of the form (0, k - 1),+-- where the corresponding dimension in sh is k. This is never checked.+type AstVarListS sh = ListS sh (Const IntVarName)++-- There's no data invariant here. The shape matches rather the argument+-- of indexing (or gather) than the indexes.+type AstIxS ms sh = IxS sh (AstInt ms)++pattern AstLeqInt :: AstInt ms -> AstInt ms -> AstBool ms+pattern AstLeqInt t u <- (matchAstLeqInt -> Just (t, u))+  where AstLeqInt t u = AstLeqK t u++matchAstLeqInt :: AstBool ms -> Maybe (AstInt ms, AstInt ms)+matchAstLeqInt (AstLeqK @r t u)+  | Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =+      Just (t, u)+matchAstLeqInt _ = Nothing+++-- * AST++type data AstMethodOfSharing = AstMethodShare | AstMethodLet++-- | AST for tensors that are meant to be differentiated.+type role AstTensor nominal nominal nominal+data AstTensor :: AstMethodOfSharing -> AstSpanType -> Target where+  -- General operations, for scalar, ranked, shared and other tensors at once+  AstPair :: forall y z ms s.+             AstTensor ms s y -> AstTensor ms s z+          -> AstTensor ms s (TKProduct y z)+  AstProject1 :: forall y z ms s.+                 AstTensor ms s (TKProduct y z) -> AstTensor ms s y+  AstProject2 :: forall y z ms s.+                 AstTensor ms s (TKProduct y z) -> AstTensor ms s z+  AstFromVector :: forall y k ms s.+                   SNat k -> SingletonTK y+                -> Data.Vector.Vector (AstTensor ms s y)+                -> AstTensor ms s (BuildTensorKind k y)+  AstSum :: forall y k ms s.+            SNat k -> SingletonTK y+         -> AstTensor ms s (BuildTensorKind k y)+         -> AstTensor ms s y+  AstReplicate :: forall y k ms s.+                  SNat k -> SingletonTK y+               -> AstTensor ms s y+               -> AstTensor ms s (BuildTensorKind k y)+  AstMapAccumRDer+    :: forall accy by ey k ms s.+       SNat k+    -> FullShapeTK by+    -> FullShapeTK ey+    -> AstHFun s s+               (TKProduct accy ey) (TKProduct accy by)+    -> AstHFun s s+               (TKProduct (ADTensorKind (TKProduct accy ey))+                          (TKProduct accy ey))+               (ADTensorKind (TKProduct accy by))+    -> AstHFun s s+               (TKProduct (ADTensorKind (TKProduct accy by))+                          (TKProduct accy ey))+               (ADTensorKind (TKProduct accy ey))+    -> AstTensor ms s accy+    -> AstTensor ms s (BuildTensorKind k ey)+    -> AstTensor ms s (TKProduct accy (BuildTensorKind k by))+  AstMapAccumLDer+    :: forall accy by ey k ms s.+       SNat k+    -> FullShapeTK by+    -> FullShapeTK ey+    -> AstHFun s s+               (TKProduct accy ey) (TKProduct accy by)+    -> AstHFun s s+               (TKProduct (ADTensorKind (TKProduct accy ey))+                          (TKProduct accy ey))+               (ADTensorKind (TKProduct accy by))+    -> AstHFun s s+               (TKProduct (ADTensorKind (TKProduct accy by))+                          (TKProduct accy ey))+               (ADTensorKind (TKProduct accy ey))+    -> AstTensor ms s accy+    -> AstTensor ms s (BuildTensorKind k ey)+    -> AstTensor ms s (TKProduct accy (BuildTensorKind k by))+  AstApply :: (AstSpan s1, AstSpan s)+           => AstHFun s1 s x z -> AstTensor ms s1 x -> AstTensor ms s z+  AstVar :: AstVarName s y -> AstTensor ms s y+  AstCond :: forall y ms s.+             AstBool ms -> AstTensor ms s y -> AstTensor ms s y+          -> AstTensor ms s y+  AstBuild1 :: forall y k ms s.+               SNat k -> SingletonTK y+            -> (IntVarName, AstTensor ms s y)+            -> AstTensor ms s (BuildTensorKind k y)++  -- Sharing-related operations, mutually exclusive via AstMethodOfSharing+  AstLet :: forall y z s s2. AstSpan s+         => AstVarName s y -> AstTensor AstMethodLet s y+         -> AstTensor AstMethodLet s2 z+         -> AstTensor AstMethodLet s2 z+  AstShare :: AstVarName s y -> AstTensor AstMethodShare s y+           -> AstTensor AstMethodShare s y+  AstToShare :: AstTensor AstMethodLet s y+             -> AstTensor AstMethodShare s y++  -- Explicit dual numbers handling, eliminated in interpretation to ADVal+  AstPrimalPart :: forall y ms.+                   AstTensor ms FullSpan y -> AstTensor ms PrimalSpan y+  AstDualPart :: forall y ms.+                 AstTensor ms FullSpan y -> AstTensor ms DualSpan y+  AstFromPrimal :: forall y ms.+                   AstTensor ms PrimalSpan y -> AstTensor ms FullSpan y+  AstFromDual :: forall y ms.+                 AstTensor ms DualSpan y -> AstTensor ms FullSpan y++  -- Scalar arithmetic (to avoid the slowness of indexes as 1-element tensors)+  AstPlusK :: GoodScalar r+           => AstTensor ms s (TKScalar r)+           -> AstTensor ms s (TKScalar r)+           -> AstTensor ms s (TKScalar r)+  AstTimesK :: GoodScalar r+            => AstTensor ms s (TKScalar r)+            -> AstTensor ms s (TKScalar r)+            -> AstTensor ms s (TKScalar r)+  AstN1K :: GoodScalar r+         => OpCodeNum1 -> AstTensor ms s (TKScalar r)+         -> AstTensor ms s (TKScalar r)+  AstR1K :: (RealFloatH r, Nested.FloatElt r, GoodScalar r)+         => OpCode1 -> AstTensor ms s (TKScalar r)+         -> AstTensor ms s (TKScalar r)+  AstR2K :: (RealFloatH r, Nested.FloatElt r, GoodScalar r)+         => OpCode2 -> AstTensor ms s (TKScalar r)+         -> AstTensor ms s (TKScalar r)+         -> AstTensor ms s (TKScalar r)+  AstI2K :: (IntegralH r, Nested.IntElt r, GoodScalar r)+         => OpCodeIntegral2 -> AstTensor ms s (TKScalar r)+         -> AstTensor ms s (TKScalar r)+         -> AstTensor ms s (TKScalar r)+  AstConcreteK :: GoodScalar r+               => r -> AstTensor ms PrimalSpan (TKScalar r)+  AstFloorK :: (GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2)+            => AstTensor ms PrimalSpan (TKScalar r1)+            -> AstTensor ms PrimalSpan (TKScalar r2)+  AstFromIntegralK :: (GoodScalar r1, Integral r1, GoodScalar r2)+                   => AstTensor ms PrimalSpan (TKScalar r1)+                   -> AstTensor ms PrimalSpan (TKScalar r2)+  AstCastK :: (GoodScalar r1, RealFrac r1, RealFrac r2, GoodScalar r2)+           => AstTensor ms s (TKScalar r1) -> AstTensor ms s (TKScalar r2)++  -- Shaped arithmetic+  AstPlusS :: GoodScalar r+           => AstTensor ms s (TKS sh r)+           -> AstTensor ms s (TKS sh r)+           -> AstTensor ms s (TKS sh r)+  AstTimesS :: GoodScalar r+            => AstTensor ms s (TKS sh r)+            -> AstTensor ms s (TKS sh r)+            -> AstTensor ms s (TKS sh r)+  AstN1S :: GoodScalar r+         => OpCodeNum1 -> AstTensor ms s (TKS sh r)+         -> AstTensor ms s (TKS sh r)+  AstR1S :: (RealFloatH r, Nested.FloatElt r, GoodScalar r)+         => OpCode1 -> AstTensor ms s (TKS sh r)+         -> AstTensor ms s (TKS sh r)+  AstR2S :: (RealFloatH r, Nested.FloatElt r, GoodScalar r)+         => OpCode2 -> AstTensor ms s (TKS sh r)+         -> AstTensor ms s (TKS sh r)+         -> AstTensor ms s (TKS sh r)+  AstI2S :: (IntegralH r, Nested.IntElt r, GoodScalar r)+         => OpCodeIntegral2 -> AstTensor ms s (TKS sh r)+         -> AstTensor ms s (TKS sh r)+         -> AstTensor ms s (TKS sh r)+  AstConcreteS :: GoodScalar r+               => Nested.Shaped sh r -> AstTensor ms PrimalSpan (TKS sh r)+  AstFloorS :: (GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2)+            => AstTensor ms PrimalSpan (TKS sh r1)+            -> AstTensor ms PrimalSpan (TKS sh r2)+  AstFromIntegralS :: (GoodScalar r1, Integral r1, GoodScalar r2)+                   => AstTensor ms PrimalSpan (TKS sh r1)+                   -> AstTensor ms PrimalSpan (TKS sh r2)+  AstCastS :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+           => AstTensor ms s (TKS sh r1)+           -> AstTensor ms s (TKS sh r2)++  -- Shaped tensor operations+  AstIndexS :: forall shm shn x s ms.+               ShS shn+            -> AstTensor ms s (TKS2 (shm ++ shn) x) -> AstIxS ms shm+            -> AstTensor ms s (TKS2 shn x)+    -- out of bounds indexing is permitted and the results is def (==0)+  AstScatterS :: forall shm shn shp x s ms.+                 ShS shn -> AstTensor ms s (TKS2 (shm ++ shn) x)+              -> (AstVarListS shm, AstIxS ms shp)+              -> AstTensor ms s (TKS2 (shp ++ shn) x)+    -- out of bounds indexing is permitted and the results is def (==0)+  AstGatherS :: forall shm shn shp x s ms.+                ShS shn -> AstTensor ms s (TKS2 (shp ++ shn) x)+             -> (AstVarListS shm, AstIxS ms shp)+             -> AstTensor ms s (TKS2 (shm ++ shn) x)+    -- out of bounds indexing is permitted and the results is def (==0)+  AstMinIndexS :: forall n sh r r2 ms. (GoodScalar r, GoodScalar r2)+               => AstTensor ms PrimalSpan (TKS (n ': sh) r)+               -> AstTensor ms PrimalSpan (TKS (Init (n ': sh)) r2)+  AstMaxIndexS :: forall n sh r r2 ms. (GoodScalar r, GoodScalar r2)+               => AstTensor ms PrimalSpan (TKS (n ': sh) r)+               -> AstTensor ms PrimalSpan (TKS (Init (n ': sh)) r2)+  AstIotaS :: forall n r ms. GoodScalar r+           => SNat n -> AstTensor ms PrimalSpan (TKS '[n] r)+  AstAppendS :: forall m n sh x ms s.+                AstTensor ms s (TKS2 (m ': sh) x)+             -> AstTensor ms s (TKS2 (n ': sh) x)+             -> AstTensor ms s (TKS2 ((m + n) ': sh) x)+  AstSliceS :: SNat i -> SNat n -> SNat k+            -> AstTensor ms s (TKS2 (i + n + k ': sh) x)+            -> AstTensor ms s (TKS2 (n ': sh) x)+  AstReverseS :: forall n sh x ms s.+                 AstTensor ms s (TKS2 (n ': sh) x)+              -> AstTensor ms s (TKS2 (n ': sh) x)+  AstTransposeS :: (Permutation.IsPermutation perm, Rank perm <= Rank sh)+                => Permutation.Perm perm -> AstTensor ms s (TKS2 sh x)+                -> AstTensor ms s (TKS2 (Permutation.PermutePrefix perm sh) x)+  AstReshapeS :: Product sh ~ Product sh2+              => ShS sh2+              -> AstTensor ms s (TKS2 sh x) -> AstTensor ms s (TKS2 sh2 x)++  -- Conversions+  AstConvert :: TKConversion a b -> AstTensor ms s a -> AstTensor ms s b++  -- Backend-specific primitives+  AstSum0S :: AstTensor ms s (TKS2 sh x)+           -> AstTensor ms s (TKS2 '[] x)+  AstDot0S :: GoodScalar r+           => AstTensor ms s (TKS sh r) -> AstTensor ms s (TKS sh r)+           -> AstTensor ms s (TKS '[] r)+  AstDot1InS :: forall sh n r ms s. GoodScalar r+             => ShS sh -> SNat n+             -> AstTensor ms s (TKS (sh ++ '[n]) r)+             -> AstTensor ms s (TKS (sh ++ '[n]) r)+             -> AstTensor ms s (TKS sh r)+  AstMatmul2S :: GoodScalar r+              => SNat m -> SNat n -> SNat p+              -> AstTensor ms s (TKS '[m, n] r)+              -> AstTensor ms s (TKS '[n, p] r)+              -> AstTensor ms s (TKS '[m, p] r)++deriving instance Show (AstTensor ms s y)+  -- for this to work, AstConcreteS can't take a Concrete;+  -- an alternative might be @Has Show (AstTensor ms s)@, but then we'd need+  -- to write @has@ before we apply @show@ and we'd weaken @AllTargetShow@++type role AstHFun nominal nominal nominal nominal+data AstHFun s s2 x z where+  AstLambda :: ~(AstVarName s x)+            -> ~(AstTensor AstMethodLet s2 z)+            -> AstHFun s s2 x z+    -- ^ The function body can't have any free variables outside those+    -- listed in the first component of the pair; this reflects+    -- the quantification in 'HordeAd.Core.Ops.rrev'+    -- and prevents "perturbation confusion".+    --+    -- The constructor is non-strict in order not to pre-compute+    -- higher derivatives (e.g., inside folds) that are never going to be used.+    -- As a side effect, all lambdas (closed functions) are processed+    -- lazily, which makes no harm, since they have no outside free variables+    -- and so can't easiliy induce leaks by retaining outside values (e.g.,+    -- big environments from which values for the variables would be drawn).+    -- The cost of computing a reverse derivative of a fold nested inside+    -- the function argument n times is reduced by the laziness from 20^n+    -- to under 2^n (old experimental results). Note, however,+    -- that if the n-th forward and reverse derivative is taken,+    -- the laziness is defeated.++deriving instance Show (AstHFun s s2 x z)++type role AstBool nominal+data AstBool ms where+  AstBoolConst :: Bool -> AstBool ms+  AstBoolNot :: AstBool ms -> AstBool ms+  AstBoolAnd :: AstBool ms -> AstBool ms -> AstBool ms+  -- There are existential variables here.+  AstLeqK :: forall r ms. GoodScalar r+          => AstTensor ms PrimalSpan (TKScalar r)+          -> AstTensor ms PrimalSpan (TKScalar r)+          -> AstBool ms+  AstLeqS :: forall sh r ms. GoodScalar r+          => AstTensor ms PrimalSpan (TKS sh r)+          -> AstTensor ms PrimalSpan (TKS sh r)+          -> AstBool ms+deriving instance Show (AstBool ms)++data OpCodeNum1 =+    NegateOp | AbsOp | SignumOp+ deriving (Show, Eq)++data OpCode1 =+    RecipOp+  | ExpOp | LogOp | SqrtOp+  | SinOp | CosOp | TanOp | AsinOp | AcosOp | AtanOp+  | SinhOp | CoshOp | TanhOp | AsinhOp | AcoshOp | AtanhOp+ deriving (Show, Eq)++data OpCode2 =+    DivideOp+  | PowerOp | LogBaseOp+  | Atan2Op+ deriving (Show, Eq)++data OpCodeIntegral2 =+    QuotOp | RemOp+ deriving (Show, Eq)
+ src/HordeAd/Core/AstEnv.hs view
@@ -0,0 +1,68 @@+{-# LANGUAGE QuantifiedConstraints #-}+-- | The environment datatype and operations for creating and accessing it.+module HordeAd.Core.AstEnv+  ( AstEnv, emptyEnv, showsPrecAstEnv+  , extendEnv, extendEnvI, extendEnvVarsS+  ) where++import Prelude++import Data.Coerce (coerce)+import Data.Dependent.EnumMap.Strict (DEnumMap)+import Data.Dependent.EnumMap.Strict qualified as DMap+import Data.Dependent.Sum+import Data.Foldable qualified as Foldable+import Data.Kind (Type)+import Text.Show (showListWith)++import Data.Array.Nested.Shaped.Shape++import HordeAd.Core.Ast+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- | The environment that keeps values assigned to variables+-- during interpretation.+type AstEnv :: Target -> Type+type AstEnv target = DEnumMap (AstVarName FullSpan) target+  -- We can't easily index over span and tensor kind at once,+  -- so instead we represent PrimalSpan values as FullSpan+  -- (dual number) values with zero dual component and DualSpan values+  -- as FullSpan values with zero primal component.++emptyEnv :: AstEnv target+emptyEnv = DMap.empty++showsPrecAstEnv+  :: (forall y. KnownSTK y => Show (target y))+  => Int -> AstEnv target -> ShowS+showsPrecAstEnv d demap =+  showParen (d > 10) $+    showString "fromList "+    . showListWith+        (\(k :=> target) ->+           withKnownSTK (ftkToSTK $ varNameToFTK k) $+           showsPrec 2 k . showString " :=> " . showsPrec 1 target)+        (DMap.toList demap)++extendEnv :: forall target s y.+             AstVarName s y -> target y -> AstEnv target+          -> AstEnv target+extendEnv var !t !env =+  let var2 :: AstVarName FullSpan y+      var2 = coerce var  -- only FullSpan variables permitted in env; see above+  in DMap.insertWithKey (\_ _ _ -> error $ "extendEnv: duplicate " ++ show var)+                        var2 t env++extendEnvI :: BaseTensor target+           => IntVarName -> IntOf target -> AstEnv target+           -> AstEnv target+extendEnvI var !i !env = extendEnv var (tfromPrimal STKScalar i) env++extendEnvVarsS :: forall target sh. BaseTensor target+               => AstVarListS sh -> IxSOf target sh -> AstEnv target+               -> AstEnv target+extendEnvVarsS vars !ix !env =+  let assocs = zip (listsToList vars) (Foldable.toList ix)+  in foldr (uncurry extendEnvI) env assocs
+ src/HordeAd/Core/AstFreshId.hs view
@@ -0,0 +1,184 @@+-- | Operations that (impurely, via a strictly increasing thread-safe counter)+-- generate fresh variables and sometimes also produce AST terms+-- by applying functions to such variables. This module encapsulates+-- the impurity, though some functions are in IO and they are used+-- with @unsafePerformIO@ outside, so some of the impurity escapes+-- and is encapsulated elsewhere.+module HordeAd.Core.AstFreshId+  ( funToAstIO, funToAst, funToAst2, fun1ToAst+  , funToAstRevIO, funToAstFwdIO+  , funToAstIntVarIO, funToAstIntVar, funToAstI+  , funToVarsIxS, funToAstIxS+    -- * Low level counter manipulation to be used only in sequential tests+  , resetVarCounter+  ) where++import Prelude++import Control.Concurrent.Counter (Counter, add, new, set)+import Data.Int (Int64)+import GHC.Exts (IsList (..))+import System.IO.Unsafe (unsafePerformIO)++import Data.Array.Nested.Shaped.Shape++import HordeAd.Core.Ast+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- | A counter that is impure but only in the most trivial way+-- (only ever incremented by one).+unsafeAstVarCounter :: Counter+{-# NOINLINE unsafeAstVarCounter #-}+unsafeAstVarCounter = unsafePerformIO (new 100000001)++-- | Only for tests, e.g., to ensure `show` applied to terms has stable length.+-- Tests that use this tool need to be run sequentially+-- to avoid variable confusion.+resetVarCounter :: IO ()+resetVarCounter = set unsafeAstVarCounter 100000001++unsafeGetFreshAstVarId :: IO AstVarId+{-# INLINE unsafeGetFreshAstVarId #-}+unsafeGetFreshAstVarId =+  intToAstVarId <$> add unsafeAstVarCounter 1++unsafeGetFreshAstVarName :: FullShapeTK y -> Maybe (Int64, Int64)+                         -> IO (AstVarName s y)+{-# INLINE unsafeGetFreshAstVarName #-}+unsafeGetFreshAstVarName ftk bounds =+  mkAstVarName ftk bounds+  . intToAstVarId <$> add unsafeAstVarCounter 1++funToAstIO2 :: forall y z s s2 ms. AstSpan s+            => FullShapeTK y -> Maybe (Int64, Int64)+            -> (AstTensor ms s y -> AstTensor ms s2 z)+            -> IO (AstVarName s y, AstTensor ms s2 z)+{-# INLINE funToAstIO2 #-}+funToAstIO2 ftk bounds f = do+  freshId <- unsafeGetFreshAstVarName ftk bounds+  let !x = f (astVar freshId)+  return (freshId, x)+-- Warning: adding a bang before freshId breaks fragile tests.+-- Probably GHC then optimizes differently and less predictably+-- and so changes results between -O0 vs -O1 and possibly also+-- between different GHC versions and between local vs CI setup.++funToAst2 :: AstSpan s+          => FullShapeTK y -> Maybe (Int64, Int64)+          -> (AstTensor ms s y -> AstTensor ms s2 z)+          -> (AstVarName s y, AstTensor ms s2 z)+{-# NOINLINE funToAst2 #-}+funToAst2 ftk bounds = unsafePerformIO . funToAstIO2 ftk bounds++funToAstIO :: forall y z s ms. AstSpan s+           => FullShapeTK y+           -> (AstTensor ms s y -> AstTensor ms s z)+           -> IO (AstVarName s y, AstTensor ms s z)+{-# INLINE funToAstIO #-}+funToAstIO ftk = funToAstIO2 ftk Nothing++funToAst :: AstSpan s+         => FullShapeTK y -> Maybe (Int64, Int64)+         -> (AstTensor ms s y -> AstTensor ms s z)+         -> (AstVarName s y, AstTensor ms s z)+{-# NOINLINE funToAst #-}+funToAst ftk bounds = unsafePerformIO . funToAstIO2 ftk bounds++fun1ToAstIO :: FullShapeTK y -> (AstVarName s y -> AstTensor ms s y)+            -> IO (AstTensor ms s y)+{-# INLINE fun1ToAstIO #-}+fun1ToAstIO ftk f = do+  !freshId <- unsafeGetFreshAstVarName ftk Nothing+  return $! f freshId++fun1ToAst :: FullShapeTK y -> (AstVarName s y -> AstTensor ms s y)+          -> AstTensor ms s y+{-# NOINLINE fun1ToAst #-}+fun1ToAst ftk = unsafePerformIO . fun1ToAstIO ftk++funToAstRevIO :: forall x.+                 FullShapeTK x+              -> IO ( AstVarName PrimalSpan x+                    , AstTensor AstMethodShare PrimalSpan x+                    , AstVarName FullSpan x+                    , AstTensor AstMethodLet FullSpan x )+{-# INLINE funToAstRevIO #-}+funToAstRevIO ftk = do+  !freshId <- unsafeGetFreshAstVarId+  let varPrimal :: AstVarName PrimalSpan x+      varPrimal = mkAstVarName ftk Nothing freshId+      var :: AstVarName FullSpan x+      var = mkAstVarName ftk Nothing freshId+      astVarPrimal :: AstTensor AstMethodShare PrimalSpan x+      !astVarPrimal = astVar varPrimal+      astVarD :: AstTensor AstMethodLet FullSpan x+      !astVarD = astVar var+  return (varPrimal, astVarPrimal, var, astVarD)++funToAstFwdIO :: forall x.+                 FullShapeTK x+              -> IO ( AstVarName PrimalSpan (ADTensorKind x)+                    , AstTensor AstMethodShare PrimalSpan (ADTensorKind x)+                    , AstVarName PrimalSpan x+                    , AstTensor AstMethodShare PrimalSpan x+                    , AstVarName FullSpan x+                    , AstTensor AstMethodLet FullSpan x )+{-# INLINE funToAstFwdIO #-}+funToAstFwdIO ftk = do+  !freshIdDs <- unsafeGetFreshAstVarId+  !freshId <- unsafeGetFreshAstVarId+  let varPrimalD :: AstVarName PrimalSpan (ADTensorKind x)+      varPrimalD = mkAstVarName (adFTK ftk) Nothing freshIdDs+      varPrimal :: AstVarName PrimalSpan x+      varPrimal = mkAstVarName ftk Nothing freshId+      var :: AstVarName FullSpan x+      var = mkAstVarName ftk Nothing freshId+      astVarPrimalD :: AstTensor AstMethodShare PrimalSpan (ADTensorKind x)+      !astVarPrimalD = astVar varPrimalD+      astVarPrimal :: AstTensor AstMethodShare PrimalSpan x+      !astVarPrimal = astVar varPrimal+      astVarD :: AstTensor AstMethodLet FullSpan x+      !astVarD = astVar var+  return (varPrimalD, astVarPrimalD, varPrimal, astVarPrimal, var, astVarD)++funToAstIntVarIO :: Maybe (Int64, Int64) -> ((IntVarName, AstInt ms) -> a)+                 -> IO a+{-# INLINE funToAstIntVarIO #-}+funToAstIntVarIO bounds f = do+  !varName <- unsafeGetFreshAstVarName (FTKScalar @Int64) bounds+  return $! f (varName, astVar varName)++funToAstIntVar :: Maybe (Int64, Int64) -> ((IntVarName, AstInt ms) -> a) -> a+{-# NOINLINE funToAstIntVar #-}+funToAstIntVar bounds = unsafePerformIO . funToAstIntVarIO bounds++funToAstI :: Maybe (Int64, Int64) -> (AstInt ms -> t) -> (IntVarName, t)+{-# NOINLINE funToAstI #-}+funToAstI bounds f = unsafePerformIO . funToAstIntVarIO bounds+                     $ \ (!var, !i) -> let !x = f i in (var, x)++funToVarsIxIOS+  :: forall sh a ms.+     ShS sh -> ((AstVarListS sh, AstIxS ms sh) -> a) -> IO a+{-# INLINE funToVarsIxIOS #-}+funToVarsIxIOS sh f = withKnownShS sh $ do+  let freshBound n =+        unsafeGetFreshAstVarName (FTKScalar @Int64)+                                 (Just (0, fromIntegral n - 1))+  !varList <- mapM freshBound $ shsToList sh+  let !vars = fromList varList+  let !ix = fromList $ map astVar varList+  return $! f (vars, ix)++funToVarsIxS+  :: ShS sh -> ((AstVarListS sh, AstIxS ms sh) -> a) -> a+{-# NOINLINE funToVarsIxS #-}+funToVarsIxS sh = unsafePerformIO . funToVarsIxIOS sh++funToAstIxS+  :: ShS sh -> (AstIxS ms sh -> AstIxS ms sh2)+  -> (AstVarListS sh, AstIxS ms sh2)+{-# NOINLINE funToAstIxS #-}+funToAstIxS sh f = unsafePerformIO $ funToVarsIxIOS sh+                   $ \ (!vars, !ix) -> let !x = f ix in (vars, x)
+ src/HordeAd/Core/AstInline.hs view
@@ -0,0 +1,522 @@+-- | Inlining and global sharing elimination.+module HordeAd.Core.AstInline+  ( -- * Inlining+    inlineAst+    -- * Translation of global sharing to local lets+  , unshareAstTensor+    -- * Helper classes and types+  , AstMemo+  ) where++import Prelude++import Control.Arrow (second)+import Data.Dependent.EnumMap.Strict (DEnumMap)+import Data.Dependent.EnumMap.Strict qualified as DMap+import Data.Dependent.Sum (DSum (..))+import Data.EnumMap.Strict qualified as EM+import Data.Foldable qualified as Foldable+import Data.List (mapAccumR)+import Data.Some+import Data.Type.Equality ((:~:) (Refl))+import GHC.Exts (IsList (..))++import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)+import Data.Array.Nested.Shaped.Shape++import HordeAd.Core.Ast (AstBool, AstTensor)+import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))+import HordeAd.Core.Ast qualified as Ast+import HordeAd.Core.AstSimplify (substituteAst)+import HordeAd.Core.AstTools+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * The pass that inlines lets with the bottom-up strategy++type AstMemo = EM.EnumMap AstVarId Int++-- | This inlines 'HordeAd.Core.Ast.AstLet', traversing the term bottom-up.+inlineAst+  :: forall s y. AstSpan s+  => AstMemo -> AstTensor AstMethodLet s y+  -> (AstMemo, AstTensor AstMethodLet s y)+inlineAst memo v0 = case v0 of+  Ast.AstPair t1 t2 ->+    let (memo2, v1) = inlineAst memo t1+        (memo3, v2) = inlineAst memo2 t2+    in (memo3, Ast.AstPair v1 v2)+  Ast.AstProject1 t -> second Ast.AstProject1 (inlineAst memo t)+  Ast.AstProject2 t -> second Ast.AstProject2 (inlineAst memo t)+  Ast.AstFromVector snat stk l ->+    let (memo2, l2) = mapAccumR inlineAst memo l+    in (memo2, Ast.AstFromVector snat stk l2)+  Ast.AstSum snat stk v -> second (Ast.AstSum snat stk) (inlineAst memo v)+  Ast.AstReplicate snat stk v ->+    second (Ast.AstReplicate snat stk) (inlineAst memo v)+  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->+    let (memo1, f2) = inlineAstHFun memo f+        (memo2, df2) = inlineAstHFun memo1 df+        (memo3, rf2) = inlineAstHFun memo2 rf+        (memo4, acc02) = inlineAst memo3 acc0+        (memo5, es2) = inlineAst memo4 es+    in (memo5, Ast.AstMapAccumRDer k bftk eftk f2 df2 rf2 acc02 es2)+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    let (memo1, f2) = inlineAstHFun memo f+        (memo2, df2) = inlineAstHFun memo1 df+        (memo3, rf2) = inlineAstHFun memo2 rf+        (memo4, acc02) = inlineAst memo3 acc0+        (memo5, es2) = inlineAst memo4 es+    in (memo5, Ast.AstMapAccumLDer k bftk eftk f2 df2 rf2 acc02 es2)+  Ast.AstApply t ll ->+    let (memo1, t2) = inlineAstHFun memo t+        (memo2, ll2) = inlineAst memo1 ll+    in (memo2, Ast.AstApply t2 ll2)+  Ast.AstVar var ->+    let f Nothing = Just 1+        f (Just count) = Just $ count + 1+    in (EM.alter f (varNameToAstVarId var) memo, v0)+  Ast.AstCond b a2 a3 ->+    -- This is a place where our inlining may increase code size+    -- by enlarging both branches due to not considering number of syntactic+    -- occurrences, but only dynamic occurrences. Tensor expressions+    -- in conditionals are problematic and special enough+    -- that we can let it be until problems are encountered in the wild.+    -- See https://github.com/VMatthijs/CHAD/blob/main/src/Count.hs#L88-L152.+    let (memo1, b1) = inlineAstBool memo b+        (memoA2, t2) = inlineAst EM.empty a2+        (memoA3, t3) = inlineAst EM.empty a3+        memo4 = EM.unionWith max memoA2 memoA3+        memo5 = EM.unionWith (+) memo1 memo4+    in (memo5, Ast.AstCond b1 t2 t3)+  Ast.AstBuild1 k stk (var, v) ->+    let (memoV0, !v2) = inlineAst EM.empty v+        memo1 = EM.unionWith+                  (\c1 c0 -> c1 + sNatValue k * c0) memo memoV0+    in (memo1, Ast.AstBuild1 k stk (var, v2))++  Ast.AstLet var u v ->+    -- We assume there are no nested lets with the same variable, hence+    -- the delete and hence var couldn't appear in memo, so we can make+    -- the recursive call for v with memo intact, to record extra occurrences+    -- of other variables without the costly summing of maps.+    let vv = varNameToAstVarId var+        (memo1, v2) = inlineAst memo v+        memo1NoVar = EM.delete vv memo1+        (memo2, u2) = inlineAst memo1NoVar u+    in case EM.findWithDefault 0 vv memo1 of+      0 -> (memo1, v2)+      1 -> (memo2, substituteAst u2 var v2)+      count | astIsSmall (count < 10) u ->+        let (memoU0, u0) = inlineAst EM.empty u+            memo3 = EM.unionWith (\c1 c0 -> c1 + count * c0) memo1NoVar memoU0+                      -- u is small, so the union is fast+        in (memo3, substituteAst u0 var v2)+      _ -> (memo2, Ast.AstLet var u2 v2)++  Ast.AstPrimalPart a -> second Ast.AstPrimalPart $ inlineAst memo a+  Ast.AstDualPart a -> second Ast.AstDualPart $ inlineAst memo a+  Ast.AstFromPrimal a -> second Ast.AstFromPrimal $ inlineAst memo a+  Ast.AstFromDual a -> second Ast.AstFromDual $ inlineAst memo a++  Ast.AstPlusK u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstPlusK u2 v3)+  Ast.AstTimesK u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstTimesK u2 v3)+  Ast.AstN1K opCode u ->+    let (memo2, u2) = inlineAst memo u+    in (memo2, Ast.AstN1K opCode u2)+  Ast.AstR1K opCode u ->+    let (memo2, u2) = inlineAst memo u+    in (memo2, Ast.AstR1K opCode u2)+  Ast.AstR2K opCode u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstR2K opCode u2 v3)+  Ast.AstI2K opCode u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstI2K opCode u2 v3)+  Ast.AstConcreteK{} -> (memo, v0)+  Ast.AstFloorK a -> second Ast.AstFloorK $ inlineAst memo a+  Ast.AstFromIntegralK a -> second Ast.AstFromIntegralK $ inlineAst memo a+  Ast.AstCastK a -> second Ast.AstCastK $ inlineAst memo a++  Ast.AstPlusS u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstPlusS u2 v3)+  Ast.AstTimesS u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstTimesS u2 v3)+  Ast.AstN1S opCode u ->+    let (memo2, u2) = inlineAst memo u+    in (memo2, Ast.AstN1S opCode u2)+  Ast.AstR1S opCode u ->+    let (memo2, u2) = inlineAst memo u+    in (memo2, Ast.AstR1S opCode u2)+  Ast.AstR2S opCode u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstR2S opCode u2 v3)+  Ast.AstI2S opCode u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstI2S opCode u2 v3)+  Ast.AstConcreteS{} -> (memo, v0)+  Ast.AstFloorS a -> second Ast.AstFloorS $ inlineAst memo a+  Ast.AstFromIntegralS v -> second Ast.AstFromIntegralS $ inlineAst memo v+  Ast.AstCastS v -> second Ast.AstCastS $ inlineAst memo v++  Ast.AstIndexS @sh1 shn v ix ->+    let (memo1, v2) = inlineAst memo v+        (memo2, ix2) = mapAccumR inlineAst memo1 (Foldable.toList ix)+    in withKnownShS (shsFromIxS ix)+       (memo2, Ast.AstIndexS @sh1 shn v2 (fromList ix2))+  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->+    let (memo1, v2) = inlineAst memo v+        (memoI0, ix2) = mapAccumR inlineAst EM.empty (Foldable.toList ix)+        count = shsSize (shsFromIxS ix) * shsSize shn+        memo2 = EM.unionWith (\c1 c0 -> c1 + count * c0) memo1 memoI0+        !ix3 = withKnownShS (shsFromIxS ix) $ fromList ix2+    in (memo2, Ast.AstScatterS @shm @shn @shp shn v2 (vars, ix3))+  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->+    let (memo1, v2) = inlineAst memo v+        (memoI0, ix2) = mapAccumR inlineAst EM.empty (Foldable.toList ix)+        count = shsSize (shsFromListS vars) * shsSize shn+        memo2 = EM.unionWith (\c1 c0 -> c1 + count * c0) memo1 memoI0+        !ix3 = withKnownShS (shsFromIxS ix) $ fromList ix2+    in (memo2, Ast.AstGatherS @shm @shn @shp shn v2 (vars, ix3))+  Ast.AstMinIndexS a -> second Ast.AstMinIndexS $ inlineAst memo a+  Ast.AstMaxIndexS a -> second Ast.AstMaxIndexS $ inlineAst memo a+  Ast.AstIotaS{} -> (memo, v0)+  Ast.AstAppendS x y ->+    let (memo1, t1) = inlineAst memo x+        (memo2, t2) = inlineAst memo1 y+    in (memo2, Ast.AstAppendS t1 t2)+  Ast.AstSliceS i n k v -> second (Ast.AstSliceS i n k) (inlineAst memo v)+  Ast.AstReverseS v -> second Ast.AstReverseS (inlineAst memo v)+  Ast.AstTransposeS perm v -> second (Ast.AstTransposeS perm) $ inlineAst memo v+  Ast.AstReshapeS sh v -> second (Ast.AstReshapeS sh) (inlineAst memo v)++  Ast.AstConvert c v -> second (Ast.AstConvert c) $ inlineAst memo v++  Ast.AstSum0S v -> second Ast.AstSum0S (inlineAst memo v)+  Ast.AstDot0S u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstDot0S u2 v3)+  Ast.AstDot1InS m n u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstDot1InS m n u2 v3)+  Ast.AstMatmul2S m n p u v ->+    let (memo2, u2) = inlineAst memo u+        (memo3, v3) = inlineAst memo2 v+    in (memo3, Ast.AstMatmul2S m n p u2 v3)++inlineAstHFun+  :: AstSpan s2+  => AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)+inlineAstHFun memo v0 = case v0 of+  Ast.AstLambda var l ->+    -- No other free variables in l, so no outside lets can reach there,+    -- so we don't need to pass the information from v upwards.+    (memo, Ast.AstLambda var (snd $ inlineAst EM.empty l))++inlineAstBool :: AstMemo -> AstBool AstMethodLet+              -> (AstMemo, AstBool AstMethodLet)+inlineAstBool memo v0 = case v0 of+  Ast.AstBoolConst{} -> (memo, v0)+  Ast.AstBoolNot arg ->+    let (memo2, arg2) = inlineAstBool memo arg+    in (memo2, Ast.AstBoolNot arg2)+  Ast.AstBoolAnd arg1 arg2 ->+    let (memo1, b1) = inlineAstBool memo arg1+        (memo2, b2) = inlineAstBool memo1 arg2+    in (memo2, Ast.AstBoolAnd b1 b2)+  Ast.AstLeqK arg1 arg2 ->+    let (memo1, r1) = inlineAst memo arg1+        (memo2, r2) = inlineAst memo1 arg2+    in (memo2, Ast.AstLeqK r1 r2)+  Ast.AstLeqS arg1 arg2 ->+    let (memo1, r1) = inlineAst memo arg1+        (memo2, r2) = inlineAst memo1 arg2+    in (memo2, Ast.AstLeqS r1 r2)+++-- * Translation of global sharing to normal lets++type AstBindings = DEnumMap (AstVarName PrimalSpan)+                            (AstTensor AstMethodLet PrimalSpan)++bindsToLet :: forall s y.+              AstTensor AstMethodLet s y -> AstBindings+           -> AstTensor AstMethodLet s y+{-# INLINE bindsToLet #-}  -- help list fusion+bindsToLet u0 bs = foldl' bindToLet u0 (DMap.toDescList bs)+ where+  bindToLet :: AstTensor AstMethodLet s y+            -> DSum (AstVarName PrimalSpan)+                    (AstTensor AstMethodLet PrimalSpan)+            -> AstTensor AstMethodLet s y+  bindToLet !u (var :=> w) = Ast.AstLet var w u++-- | This replaces 'HordeAd.Core.Ast.AstShare' with 'HordeAd.Core.Ast.AstLet',+-- traversing the term bottom-up.+unshareAstTensor :: AstTensor AstMethodShare PrimalSpan y+                 -> AstTensor AstMethodLet PrimalSpan y+unshareAstTensor tShare =+  let (memoOut, tLet) = unshareAst DMap.empty tShare+  in bindsToLet tLet memoOut++-- This works only because the other code never inserts the same rshare+-- into more than one index element, with the share containing+-- the gather/scatter/build variables corresponding to the index.+unshareAstScoped+  :: forall z s. AstSpan s+  => [IntVarName] -> AstBindings -> AstTensor AstMethodShare s z+  -> (AstBindings, AstTensor AstMethodLet s z)+unshareAstScoped vars0 memo0 v0 =+  let (memo1, v1) = unshareAst memo0 v0+      memoDiff = DMap.difference memo1 memo0+      varsOccur :: [AstVarId] -> AstTensor AstMethodLet PrimalSpan y -> Bool+      varsOccur vs d = any (`varInAst` d) vs+      closeOccurs :: [AstVarId] -> AstBindings -> (AstBindings, AstBindings)+      closeOccurs vars memo =+        let (memoLocal, memoGlobal) = DMap.partition (varsOccur vars) memo+        in if DMap.null memoLocal+           then (memoLocal, memoGlobal)+           else let vars2 = map (\(Some var) -> varNameToAstVarId var)+                                (DMap.keys memoLocal)+                    (memoLocal2, memoGlobal2) = closeOccurs vars2 memoGlobal+                in (DMap.union memoLocal memoLocal2, memoGlobal2)+      (memoLocal1, memoGlobal1) =+        closeOccurs (map varNameToAstVarId vars0) memoDiff+  in (DMap.union memo0 memoGlobal1, bindsToLet v1 memoLocal1)++-- So far, there are no lets in the resulting term,+-- but we mark it as potentially containing lets, because in the future+-- we may optimize this by inserting some lets not at the top-level.+unshareAst+  :: forall s y. AstSpan s+  => AstBindings -> AstTensor AstMethodShare s y+  -> (AstBindings, AstTensor AstMethodLet s y)+unshareAst memo = \case+  Ast.AstPair t1 t2 ->+    let (memo1, v1) = unshareAst memo t1+        (memo2, v2) = unshareAst memo1 t2+    in (memo2, Ast.AstPair v1 v2)+  Ast.AstProject1 t -> second Ast.AstProject1 (unshareAst memo t)+  Ast.AstProject2 t -> second Ast.AstProject2 (unshareAst memo t)+  Ast.AstFromVector snat stk l ->+    let (memo2, l2) = mapAccumR unshareAst memo l+    in (memo2, Ast.AstFromVector snat stk l2)+  Ast.AstSum snat stk v -> second (Ast.AstSum snat stk) (unshareAst memo v)+  Ast.AstReplicate snat stk v ->+    second (Ast.AstReplicate snat stk) (unshareAst memo v)+  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->+    let (memo1, acc02) = unshareAst memo acc0+        (memo2, es2) = unshareAst memo1 es+    in (memo2, Ast.AstMapAccumRDer k bftk eftk f df rf acc02 es2)+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    let (memo1, acc02) = unshareAst memo acc0+        (memo2, es2) = unshareAst memo1 es+    in (memo2, Ast.AstMapAccumLDer k bftk eftk f df rf acc02 es2)+  Ast.AstApply t ll ->+    let (memo1, t2) = unshareAstHFun memo t+        (memo2, ll2) = unshareAst memo1 ll+    in (memo2, Ast.AstApply t2 ll2)+  Ast.AstVar v -> (memo, Ast.AstVar v)+  Ast.AstCond b a2 a3 ->+    let (memo1, b1) = unshareAstBool memo b+        (memo2, t2) = unshareAst memo1 a2+        (memo3, t3) = unshareAst memo2 a3+    in (memo3, Ast.AstCond b1 t2 t3)+  Ast.AstBuild1 snat stk (var, v) ->+    let (memo1, !v2) = unshareAstScoped [var] memo v+    in (memo1, Ast.AstBuild1 snat stk (var, v2))++  -- We assume v is the same if var is the same.+  Ast.AstShare varRaw a | Just Refl <- sameAstSpan @s @PrimalSpan -> case a of+    AstFromS' @y2 ftkz v ->+      let var = mkAstVarName+                  (ftkAst v) (varNameToBounds varRaw) (varNameToAstVarId varRaw)+          astVar0 = cAstFromS @y2 ftkz $ Ast.AstVar var+      in if var `DMap.member` memo+         then (memo, astVar0)+         else let (memo1, !a2) = unshareAst memo v+                    -- DMap is strict, but let's be paranoid+              in (DMap.insert var a2 memo1, astVar0)+    -- The PrimalSpan check ensures there's no need to match for+    -- Ast.AstFromPrimal (Ast.AstFromS).+    _ -> case varNameToFTK varRaw of+      ftk@(FTKR @_ @x sh' x) ->+        withShsFromShR sh' $ \(sh :: ShS sh) ->+          let var = mkAstVarName (FTKS sh x) (varNameToBounds varRaw)+                                 (varNameToAstVarId varRaw)+              astVar0 = cAstFromS @(TKS2 sh x) ftk $ Ast.AstVar var+          in if var `DMap.member` memo+             then (memo, astVar0)+             else let (memo1, !a2) = unshareAst memo (cAstSFromR @sh sh a)+                  in (DMap.insert var a2 memo1, astVar0)+      ftk@(FTKX @_ @x sh' x) ->+        withShsFromShX sh' $ \(sh :: ShS sh) ->+          let var = mkAstVarName (FTKS sh x) (varNameToBounds varRaw)+                                 (varNameToAstVarId varRaw)+              astVar0 = cAstFromS @(TKS2 sh x) ftk $ Ast.AstVar var+          in if var `DMap.member` memo+             then (memo, astVar0)+             else let (memo1, !a2) = unshareAst memo (cAstSFromX @sh sh a)+                  in (DMap.insert var a2 memo1, astVar0)+      -- it maybe not be worth it to recursively convert product+      -- so let's not do that until profiling shows we need it+      _ -> let var = varRaw+               astVar0 = Ast.AstVar var+           in if var `DMap.member` memo+              then (memo, astVar0)+              else let (memo1, !a2) = unshareAst memo a+                   in (DMap.insert var a2 memo1, astVar0)+  Ast.AstShare{} -> error "unshareAst: AstShare not in PrimalSpan"+  Ast.AstToShare v -> (memo, v)  -- nothing to unshare in this subtree++  Ast.AstPrimalPart a -> second Ast.AstPrimalPart $ unshareAst memo a+  Ast.AstDualPart a -> second Ast.AstDualPart $ unshareAst memo a+  Ast.AstFromPrimal a -> second Ast.AstFromPrimal $ unshareAst memo a+  Ast.AstFromDual a -> second Ast.AstFromDual $ unshareAst memo a++  Ast.AstPlusK u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstPlusK u2 v3)+  Ast.AstTimesK u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstTimesK u2 v3)+  Ast.AstN1K opCode u ->+    let (memo2, u2) = unshareAst memo u+    in (memo2, Ast.AstN1K opCode u2)+  Ast.AstR1K opCode u ->+    let (memo2, u2) = unshareAst memo u+    in (memo2, Ast.AstR1K opCode u2)+  Ast.AstR2K opCode u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstR2K opCode u2 v3)+  Ast.AstI2K opCode u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstI2K opCode u2 v3)+  Ast.AstConcreteK k -> (memo, Ast.AstConcreteK k)+  Ast.AstFloorK a -> second Ast.AstFloorK $ unshareAst memo a+  Ast.AstFromIntegralK v -> second Ast.AstFromIntegralK $ unshareAst memo v+  Ast.AstCastK v -> second Ast.AstCastK $ unshareAst memo v++  Ast.AstPlusS u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstPlusS u2 v3)+  Ast.AstTimesS u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstTimesS u2 v3)+  Ast.AstN1S opCode u ->+    let (memo2, u2) = unshareAst memo u+    in (memo2, Ast.AstN1S opCode u2)+  Ast.AstR1S opCode u ->+    let (memo2, u2) = unshareAst memo u+    in (memo2, Ast.AstR1S opCode u2)+  Ast.AstR2S opCode u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstR2S opCode u2 v3)+  Ast.AstI2S opCode u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstI2S opCode u2 v3)+  Ast.AstConcreteS a -> (memo, Ast.AstConcreteS a)+  Ast.AstFloorS a -> second Ast.AstFloorS $ unshareAst memo a+  Ast.AstFromIntegralS v -> second Ast.AstFromIntegralS $ unshareAst memo v+  Ast.AstCastS v -> second Ast.AstCastS $ unshareAst memo v++  Ast.AstIndexS @sh1 shn v ix ->+    let (memo1, v2) = unshareAst memo v+        (memo2, ix2) = mapAccumR unshareAst memo1 (Foldable.toList ix)+    in withKnownShS (shsFromIxS ix)+       (memo2, Ast.AstIndexS @sh1 shn v2 (fromList ix2))+  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->+    let (memo1, ix2) =+          mapAccumR (unshareAstScoped $ listsToList vars)+                    memo (Foldable.toList ix)+        (memo2, v2) = unshareAst memo1 v+        !ix3 = withKnownShS (shsFromIxS ix) $ fromList ix2+    in (memo2, Ast.AstScatterS @shm @shn @shp shn v2 (vars, ix3))+  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->+    let (memo1, ix2) =+          mapAccumR (unshareAstScoped $ listsToList vars)+                    memo (Foldable.toList ix)+        (memo2, v2) = unshareAst memo1 v+        !ix3 = withKnownShS (shsFromIxS ix) $ fromList ix2+    in (memo2, Ast.AstGatherS @shm @shn @shp shn v2 (vars, ix3))+  Ast.AstMinIndexS a -> second Ast.AstMinIndexS $ unshareAst memo a+  Ast.AstMaxIndexS a -> second Ast.AstMaxIndexS $ unshareAst memo a+  Ast.AstIotaS snat -> (memo, Ast.AstIotaS snat)+  Ast.AstAppendS x y ->+    let (memo1, t1) = unshareAst memo x+        (memo2, t2) = unshareAst memo1 y+    in (memo2, Ast.AstAppendS t1 t2)+  Ast.AstSliceS i n k v -> second (Ast.AstSliceS i n k) (unshareAst memo v)+  Ast.AstReverseS v -> second Ast.AstReverseS (unshareAst memo v)+  Ast.AstTransposeS perm v ->+    second (Ast.AstTransposeS perm) $ unshareAst memo v+  Ast.AstReshapeS sh v -> second (Ast.AstReshapeS sh) (unshareAst memo v)++  Ast.AstConvert c v -> second (Ast.AstConvert c) $ unshareAst memo v++  Ast.AstSum0S v -> second Ast.AstSum0S (unshareAst memo v)+  Ast.AstDot0S u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstDot0S u2 v3)+  Ast.AstDot1InS m n u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstDot1InS m n u2 v3)+  Ast.AstMatmul2S m n p u v ->+    let (memo2, u2) = unshareAst memo u+        (memo3, v3) = unshareAst memo2 v+    in (memo3, Ast.AstMatmul2S m n p u2 v3)++unshareAstHFun+  :: AstBindings -> AstHFun s s2 x y -> (AstBindings, AstHFun s s2 x y)+unshareAstHFun memo v0 = case v0 of+  Ast.AstLambda{} ->+    -- No other free variables in l, so no outside lets can reach there,+    -- so we don't need to pass the information from v upwards+    -- nor remove the Share constructors.+    (memo, v0)++unshareAstBool :: AstBindings -> AstBool AstMethodShare+               -> (AstBindings, AstBool AstMethodLet)+unshareAstBool memo = \case+  Ast.AstBoolConst t -> (memo, Ast.AstBoolConst t)+  Ast.AstBoolNot arg ->+    let (memo2, arg2) = unshareAstBool memo arg+    in (memo2, Ast.AstBoolNot arg2)+  Ast.AstBoolAnd arg1 arg2 ->+    let (memo1, b1) = unshareAstBool memo arg1+        (memo2, b2) = unshareAstBool memo1 arg2+    in (memo2, Ast.AstBoolAnd b1 b2)+  Ast.AstLeqK arg1 arg2 ->+    let (memo1, r1) = unshareAst memo arg1+        (memo2, r2) = unshareAst memo1 arg2+    in (memo2, Ast.AstLeqK r1 r2)+  Ast.AstLeqS arg1 arg2 ->+    let (memo1, r1) = unshareAst memo arg1+        (memo2, r2) = unshareAst memo1 arg2+    in (memo2, Ast.AstLeqS r1 r2)
+ src/HordeAd/Core/AstInterpret.hs view
@@ -0,0 +1,451 @@+{-# LANGUAGE CPP #-}+-- | Interpretation of AST terms in an arbitrary tensor operations+-- class instance. With the exception of the the interpretation+-- of the sharing mechanisms and any other performance tweaks,+-- the interpretation is the unique homorphism determined by the instance.+-- The sharing mechanisms are translated so as to preserve sharing in case+-- the instance is a term algebra as well.+module HordeAd.Core.AstInterpret+  ( interpretAstFull, interpretAstPrimal, interpretAstDual, interpretAst+    -- * Exported only to specialize elsewhere+  , interpretAstBool+  ) where++import Prelude++import Data.Coerce (coerce)+import Data.Dependent.EnumMap.Strict qualified as DMap+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import Type.Reflection (typeRep)++import Data.Array.Nested.Shaped.Shape++import HordeAd.Core.Ast+import HordeAd.Core.AstEnv+import HordeAd.Core.AstTools+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++#ifdef WITH_EXPENSIVE_ASSERTIONS+import Control.Exception.Assert.Sugar+#endif++interpretAstFull+  :: forall target y. ADReady target+  => AstEnv target -> AstTensor AstMethodLet FullSpan y+  -> target y+{-# INLINE interpretAstFull #-}+interpretAstFull = interpretAst++-- Strict environment and strict ADVal and Delta make this hard to optimize.+-- Either the environment has to be traversed to remove the dual parts or+-- the dual part needs to be potentially needlessly computed.+-- However, with correct sharing and large tensors, the overall cost+-- is negligible, so we optimize only minimally.+-- It helps that usually the dual part is either trivially computed+-- to be zero or is used elsewhere. It's rarely really lost and forgotten.+interpretAstPrimal+  :: forall target y. ADReady target+  => AstEnv target -> AstTensor AstMethodLet PrimalSpan y+  -> PrimalOf target y+interpretAstPrimal !env v1 = case v1 of+  -- This prevents computing the complex dual parts for mapAccum in ADVal.+  AstMapAccumRDer k bftk eftk f0 df0 rf0 acc0 es ->+    let f = interpretAstHFunPrimal env f0+        df = interpretAstHFunPrimal env df0+        rf = interpretAstHFunPrimal env rf0+        acc02 = interpretAstPrimal env acc0+        es2 = interpretAstPrimal env es+    in tmapAccumRDer (Proxy @(PrimalOf target))+                     k (ftkAst acc0) bftk eftk f df rf acc02 es2+  AstMapAccumLDer k bftk eftk f0 df0 rf0 acc0 es ->+    let f = interpretAstHFunPrimal env f0+        df = interpretAstHFunPrimal env df0+        rf = interpretAstHFunPrimal env rf0+        acc02 = interpretAstPrimal env acc0+        es2 = interpretAstPrimal env es+    in tmapAccumLDer (Proxy @(PrimalOf target))+                     k (ftkAst acc0) bftk eftk f df rf acc02 es2+  -- This prevents multiple ifH expansions in ADVal.+  AstCond b a1 a2 ->+    let c = interpretAstBool env b+    in tcond (ftkToSTK $ ftkAst a1) c+             (interpretAstPrimal env a1) (interpretAstPrimal env a2)+  _ -> tprimalPart (interpretAst env v1)++interpretAstDual+  :: forall target y. ADReady target+  => AstEnv target -> AstTensor AstMethodLet DualSpan y+  -> DualOf target y+{-# INLINE interpretAstDual #-}+interpretAstDual !env v1 =+  tdualPart (ftkToSTK $ ftkAst v1) (interpretAst env v1)++-- A more precise type signature would result in @PrimalOf target@+-- whenever @s@ is @PrimalSpan@, but this would complicate things,+-- e.g., we'd need an extra type family+--+-- type family SpanTarget s target :: Target where+--   SpanTarget PrimalSpan target = PrimalOf target+--   SpanTarget DualSpan target = DualOf target+--   SpanTarget FullSpan target = target+--+-- to be used in AstEnv and the codomain of interpretAst and a lot of other+-- code would need to be changed. So instead we promote results to @target@+-- similarly as in AstEnv and simiarly as we omit @PrimalOf@ in the signatures+-- of most "Ops" methods.+--+-- | Interpret a term in an environment.+--+-- Note that for 'PrimalSpan' term, the results of this function+-- land in @target y@ and not in @PrimalOf target y@.+-- To make it sound nevertheless, we maintain an invariant that a value+-- of interpretation of a term with 'PrimalSpan' has zero dual part+-- and of a term with 'DualSpan' has zero primal part.+-- The invariants holds by the properties of instances of @Ops@+-- (see especially the ADVal instance, which zeroes dual part of many ops)+-- and structural induction on Ast, inspecting spans of constructors.+-- This promotion from @PrimalOf target y@ to @target y@ coincides+-- with how most operations that in Ast have 'PrimalSpan',+-- don't have 'PrimalOf' (but have full target instead)+-- in their method signatures in @Ops@, for user convenience.+-- See, e.g., 'AstConcreteS' vs 'tsconcrete' and 'AstFloorS' vs 'tsfloor'.+interpretAst+  :: forall target s y. (ADReady target, AstSpan s)+  => AstEnv target -> AstTensor AstMethodLet s y+  -> target y+interpretAst !env = \case+  AstPair t1 t2 -> tpair (interpretAst env t1) (interpretAst env t2)+  AstProject1 t -> tproject1 (interpretAst env t)+  AstProject2 t -> tproject2 (interpretAst env t)+  AstFromVector snat stk l ->+    let l2 = V.map (interpretAst env) l+    in tfromVector snat stk l2+  AstSum snat stk v -> tsum snat stk $ interpretAst env v+  AstReplicate snat stk v ->+    treplicate snat stk (interpretAst env v)+  AstMapAccumRDer k bftk eftk f0 df0 rf0 acc0 es ->+    let f = interpretAstHFun env f0+        df = interpretAstHFun env df0+        rf = interpretAstHFun env rf0+        acc02 = interpretAst env acc0+        es2 = interpretAst env es+    in tmapAccumRDer (Proxy @target) k (ftkAst acc0) bftk eftk f df rf acc02 es2+  AstMapAccumLDer k bftk eftk f0 df0 rf0 acc0 es ->+    let f = interpretAstHFun env f0+        df = interpretAstHFun env df0+        rf = interpretAstHFun env rf0+        acc02 = interpretAst env acc0+        es2 = interpretAst env es+    in tmapAccumLDer (Proxy @target) k (ftkAst acc0) bftk eftk f df rf acc02 es2+  AstApply t ll ->+    let t2 = interpretAstHFun env t+        ll2 = interpretAst env ll+    in tApply t2 ll2+  AstVar var ->+    let var2 :: AstVarName FullSpan y+        var2 = coerce var  -- only FullSpan variables permitted in env+    in case DMap.lookup var2 env of+      Just t ->+#ifdef WITH_EXPENSIVE_ASSERTIONS+        withKnownSTK (ftkToSTK $ varNameToFTK var) $+        -- We can't assert anything about bounds, because values can be+        -- symbolic and so not directly comparable to bounds.+        assert (tftk (ftkToSTK $ varNameToFTK var) t == varNameToFTK var+                `blame` ( tftk (ftkToSTK $ varNameToFTK var) t+                        , varNameToFTK var, var, t ))+#endif+        t+      _ -> error $ "interpretAst: unknown AstVar " ++ show var+                   -- ++ " in environment " ++ showsPrecAstEnv 0 env ""+  AstCond b a1 a2 ->+    let c = interpretAstBool env b+    in tcond (ftkToSTK (ftkAst a1)) c+             (interpretAst env a1) (interpretAst env a2)+  AstBuild1 snat stk (var, v) ->+    let f i = interpretAst (extendEnvI var i env) v+    in tbuild1 snat stk f++  -- We assume there are no nested lets with the same variable.+  --+  -- Note that without the second sameAstSpan check, AstLet with both PrimalSpan+  -- would get translated to a composition of ttletPrimal and tfromPrimal,+  -- which doesn't make a difference in a translation from PrimalSpan+  -- terms to PrimalSpan terms, but does in a translation from PrimalSpan+  -- terms to FullSpan terms, causing a loss of a dual part.+  --+  -- However, right now this whole code fragment is disabled, because+  -- it increases the allocation in testsuites by ~3% and slows down the VTO1+  -- benchmark 5 times. To be re-evaluated when rewriting is changed+  -- and also more examples are available.+  AstLet {-@_ @_ @s1-} var u v -> {- case ( sameAstSpan @s1 @PrimalSpan+                                          , sameAstSpan @s @FullSpan ) of+    (Just Refl, Just Refl) ->+      let t = interpretAstPrimal env u+          stk = ftkToSTK (ftkAst u)+          env2 wPrimal = extendEnv var (tfromPrimal stk wPrimal) env+      in ttletPrimal t (\wPrimal -> interpretAst (env2 wPrimal) v)+        -- @ttletPrimal@ can be more frugal in some targets, though we pay+        -- for it with @ftkAst@+    _ -> -}+      let t = interpretAst env u+          env2 w = extendEnv var w env+      in ttlet t (\w -> interpretAst (env2 w) v)++  AstPrimalPart a ->+    tfromPrimal (ftkToSTK (ftkAst a)) (tprimalPart $ interpretAstFull env a)+  AstDualPart a ->+    tfromDual (tdualPart (ftkToSTK (ftkAst a)) $ interpretAstFull env a)+  AstFromPrimal a ->+    -- By the invariant, interpretation of @a@ has zero dual part,+    -- so we don't have to do the following to remove the dual part,+    -- but we still do, because there's almost no rewriting of delta+    -- expressions, so even though they are semantically zero, they'd build+    -- up considerably if not wiped out regularly. By constrast, operations+    -- on AstConstant are rewritten eagerly to AstConstant, so for AstFromDual+    -- we really don't need to do anything.+    tfromPrimal (ftkToSTK (ftkAst a)) (interpretAstPrimal env a)+  AstFromDual a -> interpretAst env a+    -- By the invariant, interpretation of @a@ has zero primal part,+    -- so we don't have to do the following to remove the primal part:+    --   tfromDual (interpretAstDual env a)++  AstPlusK u v -> interpretAst env u + interpretAst env v+  AstTimesK u v -> interpretAst env u * interpretAst env v+  AstN1K opCode u ->+    let u2 = interpretAst env u+    in interpretAstN1 opCode u2+  AstR1K opCode u ->+    let u2 = interpretAst env u+    in interpretAstR1 opCode u2+  AstR2K opCode u v ->+    let u2 = interpretAst env u+        v2 = interpretAst env v+    in interpretAstR2 opCode u2 v2+  AstI2K opCode u v ->+    let u2 = interpretAst env u+        v2 = interpretAst env v+    in interpretAstI2 opCode u2 v2+  AstConcreteK k ->+    tkconcrete @target k+      -- this is equal to the following+      -- (and similarly for tsconcretet and tsiota below):+      -- tfromPrimal @target STKScalar $ tkconcrete @(PrimalOf target) k+  AstFloorK v ->+    -- By the invariant v has zero dual part, so the following suffices:+    tkfloor $ interpretAst env v+  AstFromIntegralK v ->+    -- By the invariant v has zero dual part, so the following suffices:+    tkfromIntegral $ interpretAst env v+  AstCastK v -> tkcast $ interpretAst env v++  AstPlusS u v -> interpretAst env u + interpretAst env v+  AstTimesS u v -> interpretAst env u * interpretAst env v+  AstN1S opCode u -> interpretAstN1 opCode (interpretAst env u)+  AstR1S opCode u -> interpretAstR1 opCode (interpretAst env u)+  AstR2S opCode u v ->+    interpretAstR2 opCode (interpretAst env u) (interpretAst env v)+  AstI2S opCode u v ->+    interpretAstI2 opCode (interpretAst env u) (interpretAst env v)+  AstConcreteS a -> tsconcrete a+  AstFloorS v ->+    -- By the invariant v has zero dual part, so the following suffices:+    tsfloor $ interpretAst env v+  AstFromIntegralS v ->+    -- By the invariant v has zero dual part, so the following suffices:+    tsfromIntegral $ interpretAst env v+  AstCastS @r1 @r2 v ->+    -- Specializing for the cases covered by rules in GHC.Internal.Float.+    case testEquality (typeRep @r1) (typeRep @Double) of+      Just Refl -> case testEquality (typeRep @r2) (typeRep @Float) of+        Just Refl -> tscast @_ @Double @Float $ interpretAst env v+        _ -> tscast @_ @Double $ interpretAst env v+      _ -> case testEquality (typeRep @r1) (typeRep @Float) of+        Just Refl -> case testEquality (typeRep @r2) (typeRep @Double) of+          Just Refl -> tscast @_ @Float @Double $ interpretAst env v+          _ -> tscast @_ @Float $ interpretAst env v+        _ -> tscast $ interpretAst env v++  AstIndexS @sh1 sh2 v ix -> case ftkToSTK (ftkAst v) of+    STKS _ x ->+      withKnownShS (shsFromIxS ix) $+      withKnownShS sh2 $+      withKnownSTK x $+      let v2 = interpretAst env v+          ix3 = interpretAstPrimal env <$> ix+      in tsindex @target @sh1 v2 ix3+  {- TODO: this breaks specialization:+  AstScatterS shn v (ZS, ix) -> case ftkToSTK (ftkAst v) of+    STKS _ x ->+      withKnownShS shn $+      withKnownShS (shsFromIxS ix) $+      withKnownSTK x $+      tsoneHot (interpretAst env v) (interpretAstPrimal env <$> ix) -}+  AstScatterS @shm @shn @shp+              shn v (vars, ix) -> case ftkToSTK (ftkAst v) of+    STKS _ x ->+      withKnownShS (shsFromListS vars) $+      withKnownShS shn $+      withKnownShS (shsFromIxS ix) $+      withKnownSTK x $+      let t1 = interpretAst env v+          f2 :: IxSOf target shm -> IxSOf target shp+          f2 !ix2 = interpretAstPrimal (extendEnvVarsS vars ix2 env) <$> ix+      in tsscatter @_ @shm @shn @shp t1 f2+  AstGatherS shn v (ZS, ix) -> case ftkToSTK (ftkAst v) of+    STKS _ x ->+      withKnownShS shn $+      withKnownShS (shsFromIxS ix) $+      withKnownSTK x $+      tsindex (interpretAst env v) (interpretAstPrimal env <$> ix)+  AstGatherS @shm @shn @shp+             shn v (vars, ix) -> case ftkToSTK (ftkAst v) of+    STKS _ x ->+      withKnownShS (shsFromListS vars) $+      withKnownShS shn $+      withKnownShS (shsFromIxS ix) $+      withKnownSTK x $+      let t1 = interpretAst env v+          f2 :: IxSOf target shm -> IxSOf target shp+          f2 !ix2 = interpretAstPrimal (extendEnvVarsS vars ix2 env) <$> ix+      in tsgather @_ @shm @shn @shp t1 f2+  AstMinIndexS v ->+    -- By the invariant v has zero dual part, so the following suffices:+    tsminIndex $ interpretAst env v+  AstMaxIndexS v ->+    -- By the invariant v has zero dual part, so the following suffices:+    tsmaxIndex $ interpretAst env v+  AstIotaS SNat -> tsiota+  AstAppendS a b -> case ftkToSTK (ftkAst a) of+    STKS _ x ->+      withKnownSTK x $+      let t1 = interpretAst env a+          t2 = interpretAst env b+      in tsappend t1 t2+  AstSliceS i n k v -> case ftkToSTK (ftkAst v) of+    STKS _ x ->+      withKnownSTK x $+      tsslice i n k $ interpretAst env v+  AstReverseS v -> case ftkToSTK (ftkAst v) of+    STKS _ x ->+      withKnownSTK x $+      tsreverse (interpretAst env v)+  AstTransposeS perm v -> case ftkToSTK (ftkAst v) of+    STKS _ x ->+      withKnownSTK x $+      tstranspose perm $ interpretAst env v+  AstReshapeS sh2 v -> case ftkToSTK (ftkAst v) of+    STKS _ x ->+      withKnownSTK x $+      tsreshape sh2 (interpretAst env v)++  AstConvert c a ->+    tconvert c (ftkToSTK (ftkAst a)) (interpretAst env a)++  AstSum0S v -> case ftkToSTK (ftkAst v) of+    STKS sh x ->+      withKnownShS sh $+      withKnownSTK x $+      tssum0 (interpretAst env v)+  AstDot0S u v -> case ftkAst u of+    FTKS sh _ ->+      withKnownShS sh $+      tsdot0 (interpretAst env u) (interpretAst env v)+  AstDot1InS @sh @n sh SNat u v ->+    withKnownShS sh $+    tsdot1In @_ @sh (SNat @n) (interpretAst env u) (interpretAst env v)+  AstMatmul2S SNat SNat SNat u v ->+    tsmatmul2 (interpretAst env u) (interpretAst env v)++interpretAstHFun+  :: forall target x y s s2. (AstSpan s2, BaseTensor target)+  => AstEnv target -> AstHFun s s2 x y+  -> HFunOf target x y+{-# INLINE interpretAstHFun #-}+interpretAstHFun _env (AstLambda var l) =+  tlambda @target (varNameToFTK var)+  $ HFun $ \ws -> interpretAst (extendEnv var ws emptyEnv) l+      -- Interpretation in empty environment makes sense here, because+      -- there are no free variables except for the one declared.++interpretAstHFunPrimal+  :: forall target x y. ADReady target+  => AstEnv target -> AstHFun PrimalSpan PrimalSpan x y+  -> HFunOf (PrimalOf target) x y+{-# INLINE interpretAstHFunPrimal #-}+interpretAstHFunPrimal _env (AstLambda var l) =+  tlambda @(PrimalOf target) (varNameToFTK var)+  $ HFun $ \ws -> interpretAst (extendEnv var ws emptyEnv) l+      -- This is probably optimized as much as possible, because+      -- thanks to the invariant, we get zero dual part from this+      -- PrimalSpan term and so interpretAstPrimal and tfromPrimal+      -- is not needed (and would not be possible, because we lack+      -- FullShapeTK y). From the other end, due to (PrimalOf target),+      -- there won't be any dual part coming from an argument.++interpretAstBool :: ADReady target+                 => AstEnv target -> AstBool AstMethodLet+                 -> BoolOf target+interpretAstBool !env = \case+  AstBoolConst a -> if a then true else false+  AstBoolNot arg -> notB $ interpretAstBool env arg+  AstBoolAnd arg1 arg2 ->+    let b1 = interpretAstBool env arg1+        b2 = interpretAstBool env arg2+    in b1 &&* b2+  AstLeqK arg1 arg2 ->+    let r1 = interpretAstPrimal env arg1+        r2 = interpretAstPrimal env arg2+    in r1 <=. r2+  AstLeqS arg1 arg2 ->+    let r1 = interpretAstPrimal env arg1+        r2 = interpretAstPrimal env arg2+    in r1 <=. r2+++-- * Interpretation of arithmetic, boolean and relation operations++interpretAstN1 :: Num a+               => OpCodeNum1 -> a -> a+{-# INLINE interpretAstN1 #-}+interpretAstN1 NegateOp u = negate u+interpretAstN1 AbsOp u = abs u+interpretAstN1 SignumOp u = signum u++interpretAstR1 :: Floating a+               => OpCode1 -> a -> a+{-# INLINE interpretAstR1 #-}+interpretAstR1 RecipOp u = recip u+interpretAstR1 ExpOp u = exp u+interpretAstR1 LogOp u = log u+interpretAstR1 SqrtOp u = sqrt u+interpretAstR1 SinOp u = sin u+interpretAstR1 CosOp u = cos u+interpretAstR1 TanOp u = tan u+interpretAstR1 AsinOp u = asin u+interpretAstR1 AcosOp u = acos u+interpretAstR1 AtanOp u = atan u+interpretAstR1 SinhOp u = sinh u+interpretAstR1 CoshOp u = cosh u+interpretAstR1 TanhOp u = tanh u+interpretAstR1 AsinhOp u = asinh u+interpretAstR1 AcoshOp u = acosh u+interpretAstR1 AtanhOp u = atanh u++interpretAstR2 :: RealFloatH a+               => OpCode2 -> a -> a -> a+{-# INLINE interpretAstR2 #-}+interpretAstR2 DivideOp u v = u / v+interpretAstR2 PowerOp u v = u ** v+interpretAstR2 LogBaseOp u v = logBase u v+interpretAstR2 Atan2Op u v = atan2H u v++interpretAstI2 :: IntegralH a+               => OpCodeIntegral2 -> a -> a -> a+{-# INLINE interpretAstI2 #-}+interpretAstI2 QuotOp u v = quotH u v+interpretAstI2 RemOp u v = remH u v
+ src/HordeAd/Core/AstPrettyPrint.hs view
@@ -0,0 +1,584 @@+-- | Pretty-printing of the AST. Some of the variants of pretty-printing+-- almost roundtrip, while others are more readable but less faithful.+module HordeAd.Core.AstPrettyPrint+  ( PrintConfig(..), defaulPrintConfig+  , printAstVar, printAst+  ) where++import Prelude++import Data.Foldable qualified as Foldable+import Data.IntMap.Strict (IntMap)+import Data.IntMap.Strict qualified as IM+import Data.List (intersperse)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality ((:~:) (Refl))+import Data.Vector.Generic qualified as V++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Permutation (Perm (..), permToList)+import Data.Array.Nested.Shaped.Shape++import HordeAd.Core.Ast+import HordeAd.Core.AstTools+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * Pretty-printing config++-- Modeled after https://github.com/VMatthijs/CHAD/blob/755fc47e1f8d1c3d91455f123338f44a353fc265/src/TargetLanguage.hs#L335.++-- TODO: ensure that terms roundtrip if neither loseRoudtrip+-- nor ignoreNestedLambdas is set and that explicit sharing is then preserved+-- as opposed to displaying sharing as Haskell lets.+-- Note that disabling ignoreNestedLambdas causes derivatives to be computed,+-- so pretty-printing in this way can be very expensive.+data PrintConfig = PrintConfig+  { loseRoudtrip        :: Bool+  , ignoreNestedLambdas :: Bool+  , varRenames          :: IntMap String+  }++defaulPrintConfig :: PrintConfig+defaulPrintConfig = PrintConfig+  { loseRoudtrip        = True+  , ignoreNestedLambdas = True+  , varRenames          = IM.empty+  }+++-- * Pretty-printing of variables++printAstVarId :: String -> PrintConfig -> AstVarId -> ShowS+printAstVarId prefix cfg var =+  let n = fromEnum var - 100000000+  in showString $ case IM.lookup n (varRenames cfg) of+    Just name | name /= "" -> name+    _ -> prefix ++ show n++printAstVar :: forall s y. AstSpan s+            => PrintConfig -> AstVarName s y -> ShowS+printAstVar cfg var = case isTensorInt (Proxy @s) (varNameToFTK var) of+  Just Refl -> printAstIntVar cfg var+  _ -> let prefix = case lengthSTK (ftkToSTK $ varNameToFTK var) of+             0 -> "x"+             1 -> "v"+             2 -> "m"+             3 -> "t"+             4 -> "u"+             _ -> "w"+       in printAstVarId prefix cfg (varNameToAstVarId var)++printAstIntVar :: PrintConfig -> IntVarName -> ShowS+printAstIntVar cfg var = printAstVarId "i" cfg (varNameToAstVarId var)+++-- * Pretty-printing of AST terms++-- Precedences used are as in Haskell.+printAst :: forall s y ms. AstSpan s+         => PrintConfig -> Int -> AstTensor ms s y -> ShowS+printAst cfg d = \case+  AstPair t1 t2 ->+    showParen (d > 10)+    $ showString "tpair "+      . printAst cfg 11 t1+      . showString " "+      . printAst cfg 11 t2+  AstProject1 t -> printPrefixOp printAst cfg d "tproject1" [t]+  AstProject2 t -> printPrefixOp printAst cfg d "tproject2" [t]+  AstFromVector snat stk l ->+   if loseRoudtrip cfg+   then case stk of+    STKR{} ->+      showParen (d > 10)+      $ showString "rfromVector "+        . (showParen True+           $ showString "fromList "+             . showListWith (printAst cfg 0) (V.toList l))+    STKS{} ->+      showParen (d > 10)+      $ showString "sfromVector "+        . (showParen True+           $ showString "fromList "+             . showListWith (printAst cfg 0) (V.toList l))+    STKX{} ->+      showParen (d > 10)+      $ showString "xfromVector "+        . (showParen True+           $ showString "fromList "+             . showListWith (printAst cfg 0) (V.toList l))+    _ ->  -- scalar and product+      showParen (d > 10)+      $ showString "tfromVector "+        . (showParen True+           $ showString "fromList "+             . showListWith (printAst cfg 0) (V.toList l))+   else showParen (d > 10)+        $ showString ("tfromVector (" ++ show snat ++ ") (" ++ show stk ++ ") ")+          . (showParen True+             $ showString "fromList "+               . showListWith (printAst cfg 0) (V.toList l))+  -- This is too common to be verbose even in no loseRoudtrip mode.+  AstSum snat stk v -> case stk of+    STKR{} -> printPrefixOp printAst cfg d "rsum" [v]+    STKS{} -> printPrefixOp printAst cfg d+                            ("ssum @" ++ show (sNatValue snat)) [v]+    STKX{} -> printPrefixOp printAst cfg d+                            ("xsum @" ++ show (sNatValue snat)) [v]+    _ ->  -- scalar and product+      printPrefixOp printAst cfg d+                    ("tsum (" ++ show snat ++ ") (" ++ show stk ++ ")") [v]+  -- This is too common to be verbose even in no loseRoudtrip mode.+  AstReplicate snat stk v -> case stk of+    STKR{} -> printPrefixOp printAst cfg d+                            ("rreplicate " ++ show (sNatValue snat)) [v]+    STKS{} -> printPrefixOp printAst cfg d+                            ("sreplicate @" ++ show (sNatValue snat)) [v]+    STKX{} -> printPrefixOp printAst cfg d+                            ("xreplicate @" ++ show (sNatValue snat)) [v]+    _ ->  -- scalar and product+      printPrefixOp+        printAst cfg d+        ("treplicate (" ++ show snat ++ ") (" ++ show stk ++ ")") [v]+  AstMapAccumRDer k bftk eftk f df rf acc0 es ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString "tmapAccumRDer "+      . showParen True (shows k)+      . showString " "+      . printAstHFun cfg 10 f+      . showString " "+      . printAstHFun cfg 10 df+      . showString " "+      . printAstHFun cfg 10 rf+      . showString " "+      . printAst cfg 11 acc0+      . showString " "+      . printAst cfg 11 es+   else+    showParen (d > 10)+    $ showString "tmapAccumRDer "+      . showParen True (shows k)+      . showString " "+      . showParen True (shows bftk)+      . showString " "+      . showParen True (shows eftk)+      . showString " "+      . printAstHFun cfg 10 f+      . showString " "+      . printAstHFun cfg 10 df+      . showString " "+      . printAstHFun cfg 10 rf+      . showString " "+      . printAst cfg 11 acc0+      . showString " "+      . printAst cfg 11 es+  AstMapAccumLDer k bftk eftk f df rf acc0 es ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString "tmapAccumLDer "+      . showParen True (shows k)+      . showString " "+      . printAstHFun cfg 10 f+      . showString " "+      . printAstHFun cfg 10 df+      . showString " "+      . printAstHFun cfg 10 rf+      . showString " "+      . printAst cfg 11 acc0+      . showString " "+      . printAst cfg 11 es+   else+    showParen (d > 10)+    $ showString "tmapAccumLDer "+      . showParen True (shows k)+      . showString " "+      . showParen True (shows bftk)+      . showString " "+      . showParen True (shows eftk)+      . showString " "+      . printAstHFun cfg 10 f+      . showString " "+      . printAstHFun cfg 10 df+      . showString " "+      . printAstHFun cfg 10 rf+      . showString " "+      . printAst cfg 11 acc0+      . showString " "+      . printAst cfg 11 es+  AstApply t ll -> showParen (d > 10)+                   $ showString "tApply "+                     . printAstHFunOneUnignore cfg 10 t+                         -- this is a lambda, but not nested, so always printed+                     . showString " "+                     . printAst cfg 11 ll+  AstVar var -> printAstVar cfg var+  AstCond b a1 a2 ->+    showParen (d > 10)+    $ showString "ifH "+      . printAstBool cfg 11 b+      . showString " "+      . printAst cfg 11 a1+      . showString " "+      . printAst cfg 11 a2+  AstBuild1 k stk (var, v) ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString "tbuild1 ("+      . shows k+      . showString ") "+      . (showParen True+         $ showString "\\"+           . printAstIntVar cfg var+           . showString " -> "+           . printAst cfg 0 v)+   else+    showParen (d > 10)+    $ showString "tbuild1 ("+      . shows k+      . showString ") "+      . showParen True (shows stk)+      . showString " "+      . (showParen True+         $ showString "\\"+           . printAstIntVar cfg var+           . showString " -> "+           . printAst cfg 0 v)++  t@(AstLet @_ @_ @s1 var0 u0 v0) ->+    if loseRoudtrip cfg+    then let collect :: AstTensor AstMethodLet s y -> ([(ShowS, ShowS)], ShowS)+             collect (AstLet var u v) =+               let name = printAstVar cfg var+                   uPP = printAst cfg 0 u+                   (rest, corePP) = collect v+               in ((name, uPP) : rest, corePP)+             collect v = ([], printAst cfg 0 v)+             (pairs, core) = collect t+         in showParen (d > 0)+            $ showString "let "+              . foldr (.) id (intersperse (showString " ; ")+                  [name . showString " = " . uPP | (name, uPP) <- pairs])+              . showString " in "+              . core+    else let keyword = case ( sameAstSpan @s1 @PrimalSpan+                            , sameAstSpan @s @FullSpan ) of+               (Just Refl, Just Refl) -> "ttletPrimal "+               _ -> "tlet "+         in showParen (d > 10)+            $ showString keyword+              . printAst cfg 11 u0+              . showString " "+              . (showParen True+                 $ showString "\\"+                   . printAstVar cfg var0+                   . showString " -> "+                   . printAst cfg 0 v0)+  AstShare _var v -> printPrefixOp printAst cfg d "tshare" [v]+  AstToShare v -> printPrefixOp printAst cfg d "toShare" [v]++  AstPrimalPart a ->+    if loseRoudtrip cfg+    then case ftkAst a of+      FTKR{} -> printPrefixOp printAst cfg d "rprimalPart" [a]+      FTKS{} -> printPrefixOp printAst cfg d "sprimalPart" [a]+      FTKX{} -> printPrefixOp printAst cfg d "xprimalPart" [a]+      _      -> printPrefixOp printAst cfg d "tprimalPart" [a]+    else printPrefixOp printAst cfg d "tprimalPart" [a]+  AstDualPart a ->+    if loseRoudtrip cfg+    then case ftkAst a of+      FTKR{} -> printPrefixOp printAst cfg d "rdualPart" [a]+      FTKS{} -> printPrefixOp printAst cfg d "sdualPart" [a]+      FTKX{} -> printPrefixOp printAst cfg d "xdualPart" [a]+      _      -> printPrefixOp printAst cfg d "tdualPart" [a]+    else printPrefixOp printAst cfg d+                       ("tdualPart (" ++ show (ftkToSTK (ftkAst a)) ++ ")") [a]+  AstFromPrimal a ->+    if loseRoudtrip cfg+    then printAst cfg d a+    else printPrefixOp+           printAst cfg d+           ("tfromPrimal (" ++ show (ftkToSTK (ftkAst a)) ++ ")") [a]+  AstFromDual a ->+    if loseRoudtrip cfg+    then printAst cfg d a+    else printPrefixOp printAst cfg d "tfromDual" [a]++  AstPlusK u v -> printBinaryOp printAst cfg d u (6, "+") v+  AstTimesK u v -> printBinaryOp printAst cfg d u (7, "*") v+  AstN1K opCode u -> printAstN1R printAst cfg d opCode u+  AstR1K opCode u -> printAstR1R printAst cfg d opCode u+  AstR2K opCode u v -> printAstR2R printAst cfg d opCode u v+  AstI2K opCode u v -> printAstI2R printAst cfg d opCode u v+  AstConcreteK k -> showNumber k+  AstFloorK v ->+    printPrefixOp printAst cfg d "kfloor" [v]+  AstFromIntegralK v ->+    printPrefixOp printAst cfg d "kfromIntegral" [v]+  AstCastK v ->+    printPrefixOp printAst cfg d "kcast" [v]++  AstPlusS u v -> printBinaryOp printAst cfg d u (6, "+") v+  AstTimesS u v -> printBinaryOp printAst cfg d u (7, "*") v+  AstN1S opCode u -> printAstN1R printAst cfg d opCode u+  AstR1S opCode u -> printAstR1R printAst cfg d opCode u+  AstR2S opCode u v -> printAstR2R printAst cfg d opCode u v+  AstI2S opCode u v -> printAstI2R printAst cfg d opCode u v+  AstConcreteS a -> case Nested.sshape a of+    ZSS -> showParen (d > 10)+           $ showString "sscalar "+             . showNumber (Nested.sunScalar a)+    _ -> showParen (d > 10)+         $ showString "sconcrete "+           . (showParen True+              $ shows a)+  AstFloorS a ->+    printPrefixOp printAst cfg d "sfloor" [a]+  AstFromIntegralS a ->+    printPrefixOp printAst cfg d "sfromIntegral" [a]+  AstCastS a ->+    printPrefixOp printAst cfg d "scast" [a]++  AstIndexS _ v ix ->+    showParen (d > 9)+    $ printAst cfg 10 v+      . showString " !$ "+      . showListWith (printAst cfg 0) (Foldable.toList ix)+  AstScatterS _ v (ZS, ix) ->+    showParen (d > 9)+    $ showString "soneHot "+      . printAst cfg 11 v+      . showString " "+      . showListWith (printAst cfg 0) (Foldable.toList ix)+  AstScatterS sh v (vars, ix) ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString "sscatter "+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . showListWith (printAstIntVar cfg)+                          (listsToList vars)+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+   else+    showParen (d > 10)+    $ showString ("sscatter " ++ show sh ++ " ")+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . showListWith (printAstIntVar cfg)+                          (listsToList vars)+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+  {- Let's re-enable this when/if we remove AstIndexS altogether+     or at least stop rewriting this to AstIndexS but instead optimize+     the instances for this case:+  AstGatherS _ v (ZS, ix) ->+    showParen (d > 9)+    $ printAst cfg 10 v+      . showString " !$ "+      . showListWith (printAst cfg 0) (Foldable.toList ix) -}+  AstGatherS sh v (vars, ix) ->+   if loseRoudtrip cfg+   then+    showParen (d > 10)+    $ showString "sgather "+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . showListWith (printAstIntVar cfg)+                          (listsToList vars)+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+   else+    showParen (d > 10)+    $ showString ("sgather " ++ show sh ++ " ")+      . printAst cfg 11 v+      . showString " "+      . (showParen True+         $ showString "\\"+           . showListWith (printAstIntVar cfg)+                          (listsToList vars)+           . showString " -> "+           . showListWith (printAst cfg 0) (Foldable.toList ix))+  AstMinIndexS a -> printPrefixOp printAst cfg d "sminIndex" [a]+  AstMaxIndexS a -> printPrefixOp printAst cfg d "smaxIndex" [a]+  AstIotaS snat ->+    showParen (d > 10)+    $ showString ("siota (" ++ show snat ++ ")")+  AstAppendS t1 t2 ->+    showParen (d > 10)+    $ showString "sappend "+      . printAst cfg 11 t1+      . showString " "+      . printAst cfg 11 t2+  AstSliceS i n _k v ->+    printPrefixOp printAst cfg d+                  ("sslice (" ++ show i ++ ") (" ++ show n ++ ")") [v]+  AstReverseS v -> printPrefixOp printAst cfg d "sreverse" [v]+  AstTransposeS (SNat' @1 `PCons` SNat' @0 `PCons` PNil) v ->+    printPrefixOp printAst cfg d "str" [v]+  AstTransposeS perm v ->+   if loseRoudtrip cfg+   then printPrefixOp printAst cfg d+                      ("stranspose @" ++ show (permToList perm)) [v]+   else printPrefixOp+          printAst cfg d+          ("ttranspose (makePerm @" ++ show (permToList perm) ++ ")") [v]+  AstReshapeS sh2 v ->+    printPrefixOp printAst cfg d ("sreshape @" ++ show (shsToList sh2)) [v]++  -- TODO: pretty-print correctly szip, sunzip, snestS, sunNestS+  -- or at least make sure they get printed as tconvert, not as the others+  AstConvert c t -> case (ftkToSTK (ftkAst t), convertFTK c (ftkAst t)) of+    (STKScalar, FTKS{}) -> printPrefixOp printAst cfg d "sfromK" [t]+    (STKR{}, FTKS{}) -> printPrefixOp printAst cfg d "sfromR" [t]+    (STKX{}, FTKS{}) -> printPrefixOp printAst cfg d "sfromX" [t]+    (STKS{}, FTKScalar) -> printPrefixOp printAst cfg d "kfromS" [t]+    (STKS{}, FTKR{}) -> printPrefixOp printAst cfg d "rfromS" [t]+    (STKS{}, FTKX{}) -> printPrefixOp printAst cfg d "xfromS" [t]+    (ystk, _) -> let s = "tconvert (" ++ show c ++ ") (" ++ show ystk ++ ")"+                 in printPrefixOp printAst cfg d s [t]++  AstSum0S v ->+    printPrefixOp printAst cfg d "ssum0" [v]+  AstDot0S u v ->+    printPrefixOp printAst cfg d "sdot0" [u, v]+  AstDot1InS _ _ u v ->+    printPrefixOp printAst cfg d "sdot1In" [u, v]+  AstMatmul2S SNat SNat SNat u v ->+    showParen (d > 10)+    $ showString "smatmul2 "+      . printAst cfg 11 u+      . showString " "+      . printAst cfg 11 v++showNumber :: (Ord a, Num a, Show a) => a -> ShowS+{-# INLINE showNumber #-}+showNumber a = showParen (a < 0) $ shows a++-- Differs from standard only in the space after comma.+showListWith :: (a -> ShowS) -> [a] -> ShowS+{-# INLINE showListWith #-}+showListWith = showCollectionWith "[" ", " "]"++showCollectionWith :: String -> String -> String -> (a -> ShowS) -> [a] -> ShowS+{-# INLINE showCollectionWith #-}+showCollectionWith start _   end _     []     s = start ++ end ++ s+showCollectionWith start sep end showx (x:xs) s = start ++ showx x (showl xs)+ where+  showl []     = end ++ s+  showl (y:ys) = sep ++ showx y (showl ys)++printAstHFun :: (AstSpan s, AstSpan s2)+             => PrintConfig -> Int -> AstHFun s s2 x y -> ShowS+printAstHFun cfg d = \case+  AstLambda var l ->+    if loseRoudtrip cfg+    then if ignoreNestedLambdas cfg+         then showString "<lambda>"+         else showParen (d > 0)+              $ showString "\\"+                . printAstVar cfg var+                . showString " -> "+                . printAst cfg 0 l+    else showParen (d > 0)+         $ showString "tlambda $ \\"+           . printAstVar cfg var+           . showString " -> "+           . printAst cfg 0 l++printAstHFunOneUnignore :: (AstSpan s, AstSpan s2)+                        => PrintConfig -> Int -> AstHFun s s2 x y -> ShowS+printAstHFunOneUnignore cfg d = \case+  AstLambda var l ->+    if loseRoudtrip cfg+    then showParen (d > 0)+         $ showString "\\"+           . printAstVar cfg var+           . showString " -> "+           . printAst cfg 0 l+    else showParen (d > 0)+         $ showString "tlambda $ \\"+           . printAstVar cfg var+           . showString " -> "+           . printAst cfg 0 l++printAstBool :: PrintConfig -> Int -> AstBool ms -> ShowS+printAstBool cfg d = \case+  AstBoolConst b -> showString $ if b then "true" else "false"+  AstBoolNot u -> printPrefixOp printAstBool cfg d "notB" [u]+  AstBoolAnd u v -> printBinaryOp printAstBool cfg d u (3, "&&*") v+  AstLeqK u v -> printBinaryOp printAst cfg d u (4, "<=.") v+  AstLeqS u v -> printBinaryOp printAst cfg d u (4, "<=.") v++printAstN1R :: (PrintConfig -> Int -> a -> ShowS)+           -> PrintConfig -> Int -> OpCodeNum1 -> a -> ShowS+printAstN1R pr cfg d opCode u = case opCode of+  NegateOp -> printPrefixOp pr cfg d "negate" [u]+  AbsOp -> printPrefixOp pr cfg d "abs" [u]+  SignumOp -> printPrefixOp pr cfg d "signum" [u]++printAstR1R :: (PrintConfig -> Int -> a -> ShowS)+           -> PrintConfig -> Int -> OpCode1 -> a -> ShowS+printAstR1R pr cfg d opCode u = case opCode of+  RecipOp -> printPrefixOp pr cfg d "recip" [u]+  ExpOp -> printPrefixOp pr cfg d "exp" [u]+  LogOp -> printPrefixOp pr cfg d "log" [u]+  SqrtOp -> printPrefixOp pr cfg d "sqrt" [u]+  SinOp -> printPrefixOp pr cfg d "sin" [u]+  CosOp -> printPrefixOp pr cfg d "cos" [u]+  TanOp -> printPrefixOp pr cfg d "tan" [u]+  AsinOp -> printPrefixOp pr cfg d "asin" [u]+  AcosOp -> printPrefixOp pr cfg d "acos" [u]+  AtanOp -> printPrefixOp pr cfg d "atan" [u]+  SinhOp -> printPrefixOp pr cfg d "sinh" [u]+  CoshOp -> printPrefixOp pr cfg d "cosh" [u]+  TanhOp -> printPrefixOp pr cfg d "tanh" [u]+  AsinhOp -> printPrefixOp pr cfg d "asinh" [u]+  AcoshOp -> printPrefixOp pr cfg d "acosh" [u]+  AtanhOp -> printPrefixOp pr cfg d "atanh" [u]++printAstR2R :: (PrintConfig -> Int -> a -> ShowS)+           -> PrintConfig -> Int -> OpCode2 -> a -> a -> ShowS+printAstR2R pr cfg d opCode u v = case opCode of+  DivideOp -> printBinaryOp pr cfg d u (7, "/") v+  PowerOp -> printBinaryOp pr cfg d u (8, "**") v+  LogBaseOp -> printPrefixOp pr cfg d "logBase" [u, v]+  Atan2Op -> printPrefixOp pr cfg d "atan2H" [u, v]++printAstI2R :: (PrintConfig -> Int -> a -> ShowS)+           -> PrintConfig -> Int -> OpCodeIntegral2 -> a -> a -> ShowS+printAstI2R pr cfg d opCode u v = case opCode of+  QuotOp -> printPrefixOp pr cfg d "quotH" [u, v]+  RemOp -> printPrefixOp pr cfg d "remH" [u, v]++printPrefixOp :: (PrintConfig -> Int -> a -> ShowS)+              -> PrintConfig -> Int -> String -> [a]+              -> ShowS+{-# INLINE printPrefixOp #-}+printPrefixOp pr cfg d funcname args =+  let rs = map (\arg -> showString " " . pr cfg 11 arg) args+  in showParen (d > 10)+     $ showString funcname+       . foldr (.) id rs++printBinaryOp :: (PrintConfig -> Int -> a -> ShowS)+              -> PrintConfig -> Int -> a -> (Int, String) -> a+              -> ShowS+{-# INLINE printBinaryOp #-}+printBinaryOp pr cfg d left (prec, opstr) right =+  showParen (d > prec)+  $ pr cfg (prec + 1) left+    . showString (" " ++ opstr ++ " ")+    . pr cfg (prec + 1) right
+ src/HordeAd/Core/AstSimplify.hs view
@@ -0,0 +1,3881 @@+{-# LANGUAGE AllowAmbiguousTypes, ViewPatterns #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- {-# OPTIONS_GHC -fmax-pmcheck-models=10000 #-}+-- {-# OPTIONS_GHC -freduction-depth=10000 #-}+-- {-# OPTIONS_GHC -fconstraint-solver-iterations=10000 #-}+-- | This module holds smart constructors for AST, that is,+-- term-simplifying combinators corresponding to the Ast constructors.+-- The combinators simplify only on the basis of inspecting the roots of their+-- argument term trees. If the arguments get modified,+-- the modified forms are again inspected and potentially simplified again.+--+-- The limited simplification via combinators is enough to uncover redexes+-- for the factorization rules to fire and to undo some of the complication+-- introduced by vectorization. The intention is to leave intact as much+-- as possible of the original terms provided by the user while making+-- sure subterms introduced by vectorization are maximally simplified.+module HordeAd.Core.AstSimplify+  ( RewritePhase(..), SimplifyKnobs (..), defaultKnobs+  , -- * The simplifying combinators, one for almost each AST constructor+    astPair, astProject1, astProject2, astFromVector, astSum, astReplicate+  , astMapAccumRDer, astMapAccumLDer, astApply, astCond+  , astConcrete, astConcreteK, astConcreteS++  , astLet++  , astPrimalPart, astDualPart++  , astFloorK, astFromIntegralK, astCastK++  , astFloorS, astFromIntegralS, astCastS++  , astIndexS, astIndexKnobsS, astScatterS, astGatherS, astGatherKnobsS+  , astAppendS, astSliceS, astReverseS, astTransposeS, astReshapeS++  , astConvert+  , astFromS', astKFromS', astRFromS', astXFromS'+  , astSFromK', astSFromR', astSFromX'+  , astSum0S, astDot0S, astDot1InS, astMatmul2S++    -- * Helper combinators+  , astLetFun+    -- * Substitution operations+  , substituteAst, substituteAstIxS, substituteAstBool+  ) where++import Prelude++import Control.Exception.Assert.Sugar+import Control.Monad (mapAndUnzipM, mplus)+import Data.Foldable qualified as Foldable+import Data.Functor.Const+import Data.Functor.Product qualified as Fun+import Data.GADT.Compare+import Data.Int (Int64)+import Data.List (findIndex)+import Data.Maybe (catMaybes, fromMaybe, isJust)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Type.Ord (Compare)+import Data.Vector.Generic qualified as V+import Data.Vector.Strict qualified as Data.Vector+import GHC.Exts (IsList (..))+import GHC.TypeLits+  ( Nat+  , OrderingI (..)+  , cmpNat+  , fromSNat+  , sameNat+  , type (+)+  , type (-)+  , type (<=)+  , type (<=?)+  )+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)+import Unsafe.Coerce (unsafeCoerce)++import Data.Array.Nested (Replicate, type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert+  (shrFromShS, shxFromShS, withShsFromShR, withShsFromShX)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation (DropLen, Perm (..), TakeLen, permInverse)+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types+  (Head, Init, Last, Tail, snatPlus, unsafeCoerceRefl)++import HordeAd.Core.Ast+  ( AstBool (AstBoolConst)+  , AstTensor (AstConcreteK, AstConcreteS, AstPlusK, AstPlusS, AstTimesK, AstTimesS)+  )+import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))+import HordeAd.Core.Ast qualified as Ast+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersAst (sunReplicate1, sunReplicateN, sunReplicateScal)+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind++data RewritePhase =+    PhaseUnspecified+  | PhaseVectorization+  | PhaseSimplification+  | PhaseExpansion+  | PhaseContraction+ deriving (Show, Eq)++data SimplifyKnobs = SimplifyKnobs+  { knobPhase :: RewritePhase+  }+ deriving Show++defaultKnobs :: SimplifyKnobs+defaultKnobs = SimplifyKnobs PhaseUnspecified++-- @PhaseVectorization@ should only affect the topmost redex.+deVect :: SimplifyKnobs -> SimplifyKnobs+deVect (SimplifyKnobs PhaseVectorization) = SimplifyKnobs PhaseUnspecified+deVect knobs = knobs+++-- * Expressing operations as Gather; introduces new variable names++-- | We keep AstTranspose terms for as long as possible, because+-- they are small and fuse nicely in many cases. For some forms of indexing+-- and nesting with reshape and gather they don't fuse, which is when+-- this function is invoked.+astTransposeAsGatherS+  :: forall perm sh s r. AstSpan s+  => SimplifyKnobs -> Permutation.Perm perm+  -> AstTensor AstMethodLet s (TKS2 sh r)+  -> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)+{-# NOINLINE astTransposeAsGatherS #-}+astTransposeAsGatherS knobs perm v =+  let FTKS shn _ = ftkAst v+      shnPermuted = shsPermute perm (shsTakeLen perm shn)+  in funToVarsIxS @_ @AstMethodLet shnPermuted $ \ (!vars, !ix) ->+    -- See astGatherCase.AstTransposeS for similar code with more comments.+    gcastWith (lemRankMapJust $ shsTakeLen perm shn) $+    gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $+    permInverse perm $ \(invperm :: Nested.Perm invperm) proof ->+      case proof (ssxFromShX $ shxFromShS $ shsTakeLen perm shn) of+        Refl ->+          gcastWith (unsafeCoerceRefl+                     :: DropLen invperm+                          (Permutation.Permute perm (TakeLen perm sh))+                        :~: '[]) $+          gcastWith (lemAppNil+                       @(Permutation.Permute invperm+                           (Permutation.Permute perm (TakeLen perm sh)))) $+          -- Seriously? This should be deduced from the above:+          gcastWith (unsafeCoerceRefl+                     :: Permutation.PermutePrefix invperm+                          (Permutation.Permute perm (TakeLen perm sh))+                        :~: Permutation.Permute invperm+                          (Permutation.Permute perm (TakeLen perm sh))) $+          -- This should follow from @proof@, if not for MapJust:+          gcastWith (unsafeCoerceRefl+                     :: Permutation.Permute invperm+                          (Permutation.Permute perm (TakeLen perm sh))+                        :~: TakeLen perm sh) $+          let asts :: AstIxS AstMethodLet (TakeLen perm sh)+              asts = ixsPermutePrefix invperm ix+          in gcastWith (unsafeCoerceRefl+                        :: TakeLen perm sh ++ DropLen perm sh :~: sh) $+             astGatherKnobsS @(Permutation.Permute perm (TakeLen perm sh))+                             @(DropLen perm sh)+                             @(TakeLen perm sh)+                             knobs (shsDropLen perm shn) v (vars, asts)++-- TODO: To make this less disastrous, we need to add an extra constructor+-- to AstIndex with the semantics "this index reshaped from shIn to shOut"+-- that fuses perfectly with itself and absorbs normal indexes+-- by substitution. Or perhaps make this the only constructor, with normal+-- indexes represented as "this index reshaped from sh to sh".+-- Or only extend AstGather and possibly also AstIndex with the extra+-- shIn and shOut arguments. This complicates any code related to+-- AstGather and AstIndex, but often prevents nested reshapes from affecting+-- term size in any way. But we'd need to be careful to avoid breaking such+-- an index into components, because that forces index normalization,+-- e.g., index(gather) can no longer simplify recursively by one index+-- component at a time (probably possible only if the index is shorter+-- that the list of variables fo the gather). There are probably bad cases+-- where term size blowup can't be avoided, because the index has to be+-- normalized between each reshape.+--+-- | This generates big terms that don't simplify well,+-- so we keep the AstReshape form until simplification gets stuck.+-- In fact, to simplify the terms we'd need advanced solving of equations+-- in integer arithmetic modulo.+astReshapeAsGatherS+  :: forall sh sh2 r s. AstSpan s+  => SimplifyKnobs -> ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh r)+  -> AstTensor AstMethodLet s (TKS2 sh2 r)+{-# NOINLINE astReshapeAsGatherS #-}+astReshapeAsGatherS knobs shOut v | Refl <- lemAppNil @sh2+                                  , Refl <- lemAppNil @sh+                                  , FTKS shIn _ <- ftkAst v =+  funToVarsIxS shOut $ \ (!vars, !ix) ->+    let fromInt :: Int -> AstInt AstMethodLet+        fromInt i = AstConcreteK (fromIntegral i)+        iUnshared :: AstInt AstMethodLet+        iUnshared = toLinearIdxS @sh2 @'[] fromInt shOut ix+        asts :: AstInt AstMethodLet -> AstIxS AstMethodLet sh+        asts i = fromLinearIdxS fromInt shIn i+    in gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh) $+       gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[]) $+-- This can't be done, because i depends on vars:+--     astLetFunB iUnshared $ \i ->+       let i = iUnshared  -- sharing broken+       in astGatherKnobsS @sh2 @'[] @sh knobs ZSS v (vars, asts i)+++-- * The simplifying combinators, one for almost each AST constructor++astPair :: AstSpan s+        => AstTensor AstMethodLet s x -> AstTensor AstMethodLet s y+        -> AstTensor AstMethodLet s (TKProduct x y)+astPair (Ast.AstFromPrimal v1) (Ast.AstFromPrimal v2) =+  Ast.AstFromPrimal $ astPair v1 v2+astPair (Ast.AstFromDual v1) (Ast.AstFromDual v2) =+  Ast.AstFromDual $ astPair v1 v2+astPair (Ast.AstConvert c1 v1) (Ast.AstConvert c2 v2)+  | checkAstFromS c1 v1 && checkAstFromS c2 v2 =+    astConvert (ConvT2 c1 c2) $ astPair v1 v2+astPair (Ast.AstConvert c1 v1) v2+  | checkAstFromS c1 v1 =+    astConvert (ConvT2 c1 ConvId) $ astPair v1 v2+astPair v1 (Ast.AstConvert c2 v2)+  | checkAstFromS c2 v2 =+    astConvert (ConvT2 ConvId c2) $ astPair v1 v2+astPair v1 v2 = Ast.AstPair v1 v2++astProject1+  :: forall x z s. AstSpan s+  => AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s x+astProject1 u = case u of+  Ast.AstPair x _z -> x+  Ast.AstCond b v1 v2 -> astCond b (astProject1 v1) (astProject1 v2)+  Ast.AstLet var t v -> astLet var t (astProject1 v)+  Ast.AstFromPrimal u1 -> Ast.AstFromPrimal $ astProject1 u1+  Ast.AstFromDual u1 -> Ast.AstFromDual $ astProject1 u1+  Ast.AstConvert c t | FTKProduct yftk _ <- ftkAst t+                     , FTKProduct zftk _ <- convertFTK c (ftkAst t)+                     , Just Refl <- matchingFTK yftk zftk -> astProject1 t+  -- TODO: generalize this somehow to arbitrary Conversions of the right type.+  -- At worst, just generate the canonical (?) c1 for the types at hand.+  Ast.AstConvert c@(ConvT2 c1 _c2) t+    | checkAstFromS c t ->+      astConvert c1 $ astProject1 t+  _ -> Ast.AstProject1 u++astProject2+  :: forall x z s. AstSpan s+  => AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s z+astProject2 u = case u of+  Ast.AstPair _x z -> z+  Ast.AstCond b v1 v2 -> astCond b (astProject2 v1) (astProject2 v2)+  Ast.AstLet var t v -> astLet var t (astProject2 v)+  Ast.AstFromPrimal u1 -> Ast.AstFromPrimal $ astProject2 u1+  Ast.AstFromDual u1 -> Ast.AstFromDual $ astProject2 u1+  Ast.AstConvert c t | FTKProduct _ yftk <- ftkAst t+                     , FTKProduct _ zftk <- convertFTK c (ftkAst t)+                     , Just Refl <- matchingFTK yftk zftk -> astProject2 t+  Ast.AstConvert c@(ConvT2 _c1 c2) t+    | checkAstFromS c t ->+      astConvert c2 $ astProject2 t+  _ -> Ast.AstProject2 u++astFromVector :: forall y k s. AstSpan s+              => SNat k -> SingletonTK y+              -> Data.Vector.Vector (AstTensor AstMethodLet s y)+              -> AstTensor AstMethodLet s (BuildTensorKind k y)+astFromVector (SNat' @1) stk v = astReplicate (SNat @1) stk (v V.! 0)+astFromVector snat@SNat stk l = fromMaybe (Ast.AstFromVector snat stk l) $+  -- This disables some rules, e.g., indexing or summing of fromVector+  -- of concrete arrays, but allocating an extra array of the same size+  -- as the fromVector is not a big deal and early rules are better+  -- then the same rules in contraction phase.+  (case (sameAstSpan @s @PrimalSpan, stk) of+     (Just Refl, STKScalar) ->+       let unConc :: AstTensor AstMethodLet PrimalSpan y+                  -> Maybe (Concrete y)+           unConc (AstConcreteK a) = Just $ Concrete a+           unConc _ = Nothing+       in case V.mapM unConc l of+         Just l4 | V.null l4 -> error "astFromVector: empty vector"+         Just l4 -> Just $ astConcreteS (tfromVector snat stk l4)+         Nothing -> Nothing+     (Just Refl, STKS _ STKScalar) ->+       let unConc :: AstTensor AstMethodLet PrimalSpan y+                  -> Maybe (Concrete y)+           unConc (AstConcreteS a) = Just $ Concrete a+           unConc _ = Nothing+       in case V.mapM unConc l of+         Just l4 | V.null l4 -> error "astFromVector: empty vector"+         Just l4 -> Just $ astConcreteS (tfromVector snat stk l4)+         Nothing -> Nothing+     _ -> Nothing)+  `mplus`+  (case sameAstSpan @s @FullSpan of+     Just Refl ->+       let unFromPrimal :: AstTensor AstMethodLet FullSpan y+                        -> Maybe (AstTensor AstMethodLet PrimalSpan y)+           unFromPrimal (Ast.AstFromPrimal t) = Just t+           unFromPrimal _ = Nothing+       in case V.mapM unFromPrimal l of+         Just l2 | V.null l2 -> error "astFromVector: empty vector"+         Just l2 -> Just $ Ast.AstFromPrimal $ astFromVector snat stk l2+         Nothing -> Nothing+     _ -> Nothing)+  `mplus`+  (case sameAstSpan @s @FullSpan of+     Just Refl ->+       let unFromDual :: AstTensor AstMethodLet FullSpan y+                        -> Maybe (AstTensor AstMethodLet DualSpan y)+           unFromDual (Ast.AstFromDual t) = Just t+           unFromDual _ = Nothing+       in case V.mapM unFromDual l of+         Just l2 | V.null l2 -> error "astFromVector: empty vector"+         Just l2 -> Just $ Ast.AstFromDual $ astFromVector snat stk l2+         Nothing -> Nothing+     _ -> Nothing)+  `mplus`+  (let unFrom :: FullShapeTK x+              -> AstTensor AstMethodLet s y+              -> Maybe (AstTensor AstMethodLet s x)+       unFrom xftk (AstFromS' _ t) =+         case matchingFTK (ftkAst t) xftk of+           Just Refl -> Just t+           Nothing -> error "astFromVector: impossible shape"+       unFrom _ _ = Nothing+   in case V.uncons l of+     Just (Ast.AstConvert c t, _) ->+       let xftk = ftkAst t+       in case V.mapM (unFrom xftk) l of+         Just l2 ->+           -- Here we heavily depend on c being semantically determined+           -- by the domain and codomain. We choose one such c,+           -- not necessarily the most efficient of them all.+           Just $ astConvert (buildTKConversion snat xftk c)+                $ astFromVector snat (ftkToSTK xftk) l2+         Nothing -> Nothing+     Just{} -> Nothing+     Nothing -> error "astFromVector: empty vector")++astSum :: forall y k s. AstSpan s+       => SNat k -> SingletonTK y+       -> AstTensor AstMethodLet s (BuildTensorKind k y)+       -> AstTensor AstMethodLet s y+astSum snat@SNat stk t0 = case t0 of+  _ | Just Refl <- testEquality snat (SNat @0) ->+    let ftk = razeFTK snat stk (ftkAst t0)+    in fromPrimal $ astConcrete ftk (tdefTarget ftk)+  AstConcreteS @_ @sh2 t -> case stk of+    STKS @sh _ STKScalar ->+      gcastWith (unsafeCoerceRefl :: k ': sh :~: sh2) $+      astConcreteS (tsum snat stk $ Concrete t)+    STKScalar ->+      gcastWith (unsafeCoerceRefl :: '[k] :~: sh2) $+      astConcreteK (tsum snat stk $ Concrete t)+  Ast.AstIotaS @_ @r (SNat @n) ->+    let i :: r+        i = fromInteger $ valueOf @n * (valueOf @n - 1) `div` 2+    in case stk of+      STKScalar -> AstConcreteK i+      STKS ZSS STKScalar -> AstConcreteS $ Nested.sscalar i+  Ast.AstReverseS v -> astSum snat stk v+  _ | Just Refl <- testEquality snat (SNat @1)+    , STKScalar <- stk ->+      astKFromS' $ astIndexS ZSS t0 (0 :.$ ZIS)+  _ | Just Refl <- testEquality snat (SNat @1)+    , STKS sh _  <- stk ->  -- other cases too rare+      astIndexS sh t0 (0 :.$ ZIS)  -- astReshape slows down the CNNO test+  Ast.AstFromVector @y2 _ _ l ->+    gcastWith (unsafeCoerceRefl :: y2 :~: y) $+    case stk of+      STKScalar -> foldr1 (+) l+      STKR _ STKScalar -> foldr1 (+) l+      STKS _ STKScalar -> foldr1 (+) l+      STKX _ STKScalar -> foldr1 (+) l+      _ -> Ast.AstSum snat stk t0+  -- See the analogous astSliceS rule.+  Ast.AstTransposeS (SNat' @1 `PCons` SNat' @0 `PCons` PNil) t+    | FTKS (_ :$$ _ :$$ _) _ <- ftkAst t+    , STKS (snat1 :$$ sh3) x <- stk+    , Just u <- unRepl1 t ->+      astReplicate snat1 (STKS sh3 x)+      $ astSum snat (STKS sh3 x) u+  Ast.AstReplicate _ STKScalar v | STKScalar <- stk ->+    v * fromPrimal (AstConcreteK $ fromInteger $ fromSNat snat)+  Ast.AstReplicate _ _ v | STKR _ (STKScalar @r) <- stk ->+    case ftkAst v of+      ftk@(FTKR sh' FTKScalar) ->+        withShsFromShR sh' $ \(sh :: ShS sh) ->+          v * astFromS'+                ftk (fromPrimal $ AstConcreteS @r+                     $ Nested.sreplicateScal sh $ fromInteger $ fromSNat snat)+  Ast.AstReplicate _ _ v | STKX _ (STKScalar @r) <- stk ->+    case ftkAst v of+      ftk@(FTKX sh' FTKScalar) ->+        withShsFromShX sh' $ \(sh :: ShS sh) ->+          v * astFromS'+                ftk (fromPrimal $ AstConcreteS @r+                     $ Nested.sreplicateScal sh $ fromInteger $ fromSNat snat)+  Ast.AstReplicate _ STKS{} v | STKS sh (STKScalar @r) <- stk ->+    case ftkAst v of+      ftk ->+          v * astFromS'+                ftk (fromPrimal $ AstConcreteS @r+                     $ Nested.sreplicateScal sh $ fromInteger $ fromSNat snat)+  -- This keeps tensors alive for longer, but it enables new simplifications,+  -- while hiding a sum inside let not often prevents other simplifications,+  -- because there are few redexes with sum but not at the top.+  Ast.AstLet var u v -> astLet var u (astSum snat stk v)+  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astSum snat stk v+  Ast.AstFromDual v -> Ast.AstFromDual $ astSum snat stk v+  Ast.AstScatterS @shm @shn @shp shn v (vars, (:.$) @k2 i1 rest)+    | STKS{} <- stk ->+      -- This boolean term may have free variables that act as universally+      -- quantified.+      case 0 <=. i1 &&* i1 <=. valueOf @k2 - 1 of+        AstBoolConst True ->+          astScatterS @shm @shn @(Tail shp) shn v (vars, rest)+        _ -> Ast.AstSum snat stk t0+  Ast.AstConvert c t | checkAstFromS c t -> case ftkAst t of+    FTKS (snat2 :$$ rest) x ->+      -- Here we'd need to change the types inside c, so instead we construct+      -- a new conversion based on the domain and codomain.+      astFromS' (razeFTK snat stk (ftkAst t0))+      $ astSum snat2 (STKS rest (ftkToSTK x)) t+    _ -> Ast.AstSum snat stk t0  -- products probably not worth the effort+  _ -> Ast.AstSum snat stk t0++astReplicate :: forall y k s. AstSpan s+             => SNat k -> SingletonTK y+             -> AstTensor AstMethodLet s y+             -> AstTensor AstMethodLet s (BuildTensorKind k y)+astReplicate snat stk t0 = case t0 of+  Ast.AstPair t1 t2 | STKProduct stk1 stk2 <- stk ->+    astPair (astReplicate snat stk1 t1) (astReplicate snat stk2 t2)+  -- This doesn't prevent indexing of replicate, because indexing goes inside+  -- the conditional, but it prevents the sum(replicate(cond)) simplification,+  -- because sum can't go inside, because it's costly and cond is eager.+  -- Ast.AstCond b v1 v2 ->+  --  astCond b (astReplicate snat stk v1) (astReplicate snat stk v2)+  -- TODO: This rules is, in principle, very good, because it permits many other+  -- rules to fire. However, one of these other rules is indexing of transpose+  -- that in some cases complicates terms and causes OOM in CNNI tests.+  -- We need to restrict the indexing rule more effectively first.+  -- Ast.AstLet var t v -> astLet var t (astReplicate snat stk v)+  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astReplicate snat stk v+  Ast.AstFromDual v -> Ast.AstFromDual $ astReplicate snat stk v+  AstConcreteK t -> astConcreteS $ treplicate snat stk $ Concrete t+  AstConcreteS t -> astConcreteS $ treplicate snat stk $ Concrete t+  Ast.AstConvert c t | checkAstFromS c t ->+    let xftk = ftkAst t+    in astConvert (buildTKConversion snat xftk c)+                  (astReplicate snat (ftkToSTK xftk) t)+  _ -> Ast.AstReplicate snat stk t0+  -- TODO: maybe add a rule and then generalize:+  -- replicate n1 (str (replicate n2 u))+  -- ~> transpose [0, 2, 1] (replicate n1 (replicate n2 u))+  -- but the reverse rule is already in astTransposeS++-- TODO: also pull up AstFromPrimal, etc.+astMapAccumRDer+  :: forall accy by ey k s. AstSpan s+  => SNat k+  -> FullShapeTK by+  -> FullShapeTK ey+  -> AstHFun s s+             (TKProduct accy ey) (TKProduct accy by)+  -> AstHFun s s+             (TKProduct (ADTensorKind (TKProduct accy ey))+                        (TKProduct accy ey))+             (ADTensorKind (TKProduct accy by))+  -> AstHFun s s+             (TKProduct (ADTensorKind (TKProduct accy by))+                        (TKProduct accy ey))+             (ADTensorKind (TKProduct accy ey))+  -> AstTensor AstMethodLet s accy+  -> AstTensor AstMethodLet s (BuildTensorKind k ey)+  -> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))+astMapAccumRDer k bftk eftk (AstLambda varf vf)+                            (AstLambda vard vd)+                            (AstLambda varr vr)+                (AstFromS' @accyFrom accftk acc0From) es =+  let accftkFrom = ftkAst acc0From+      accFromSTK = ftkToSTK accftkFrom+      ftkf2 = FTKProduct accftkFrom eftk+      varf2 = mkAstVarName ftkf2 (varNameToBounds varf) (varNameToAstVarId varf)+      astf2 = astVar varf2+      vf2 =+        let subbed =+              substituteAst+                (astPair (astFromS' @accyFrom accftk (astProject1 astf2))+                         (astProject2 astf2))+                varf vf+        in astSFrom' @(TKProduct accy by)+                    (STKProduct accFromSTK (ftkToSTK bftk))+                    subbed+      ftkd2 = FTKProduct+                (adFTK $ FTKProduct accftkFrom eftk)+                (FTKProduct accftkFrom eftk)+      vard2 = mkAstVarName ftkd2 (varNameToBounds vard) (varNameToAstVarId vard)+      astd2 = astVar vard2+      vd2 =+        let subbed =+              substituteAst+                (astPair (astPair (astFromS' @(ADTensorKind accyFrom)+                                     (adFTK accftk)+                                     (astProject1 (astProject1 astd2)))+                                  (astProject2 (astProject1 astd2)))+                         (astPair (astFromS' @accyFrom accftk+                                     (astProject1 (astProject2 astd2)))+                                  (astProject2 (astProject2 astd2))))+                vard vd+        in astSFrom' @(ADTensorKind (TKProduct accy by))+                    (adSTK $ STKProduct accFromSTK (ftkToSTK bftk))+                    subbed+      ftkr2 = FTKProduct+                (adFTK $ FTKProduct accftkFrom bftk)+                (FTKProduct accftkFrom eftk)+      varr2 = mkAstVarName ftkr2 (varNameToBounds varr) (varNameToAstVarId varr)+      astr2 = astVar varr2+      vr2 =+        let subbed =+              substituteAst+                (astPair (astPair (astFromS' @(ADTensorKind accyFrom)+                                     (adFTK accftk)+                                     (astProject1 (astProject1 astr2)))+                                  (astProject2 (astProject1 astr2)))+                         (astPair (astFromS' @accyFrom accftk+                                     (astProject1 (astProject2 astr2)))+                                  (astProject2 (astProject2 astr2))))+                varr vr+        in astSFrom' @(ADTensorKind (TKProduct accy ey))+                    (adSTK $ STKProduct accFromSTK (ftkToSTK eftk))+                    subbed+  in astFromS' @(TKProduct accyFrom (BuildTensorKind k by))+               (FTKProduct accftk (buildFTK k bftk))+     $ astMapAccumRDer k bftk eftk (AstLambda varf2 vf2)+                                   (AstLambda vard2 vd2)+                                   (AstLambda varr2 vr2)+                                   acc0From es+astMapAccumRDer k bftk eftk (AstLambda varf vf)+                            (AstLambda vard vd)+                            (AstLambda varr vr)+                acc0 (AstFromS' @esShsFrom _esShsFTK esFrom) =+  let accftk = ftkAst acc0+      accstk = ftkToSTK accftk+      esShsFrom = ftkAst esFrom+      esShsFromSTK = ftkToSTK esShsFrom+  in case razeSTK esShsFromSTK of+    (eftkFromSTK :: SingletonTK eyFrom) ->+      gcastWith (unsafeCoerceRefl+                 :: BuildTensorKind k eyFrom :~: esShsFrom) $+      let eftkFrom = razeFTK k eftkFromSTK esShsFrom+          ftkf2 = FTKProduct accftk eftkFrom+          varf2 =+            mkAstVarName ftkf2 (varNameToBounds varf) (varNameToAstVarId varf)+          astf2 = astVar varf2+          vf2 =+            let subbed =+                  substituteAst+                    (astPair (astProject1 astf2)+                             (astFromS' @eyFrom eftk (astProject2 astf2)))+                    varf vf+            in subbed+          ftkd2 = FTKProduct+                    (adFTK $ FTKProduct accftk eftkFrom)+                    (FTKProduct accftk eftkFrom)+          vard2 =+            mkAstVarName ftkd2 (varNameToBounds vard) (varNameToAstVarId vard)+          astd2 = astVar vard2+          vd2 =+            let subbed =+                  substituteAst+                    (astPair (astPair (astProject1 (astProject1 astd2))+                                      (astFromS' @(ADTensorKind eyFrom)+                                         (adFTK eftk)+                                         (astProject2 (astProject1 astd2))))+                             (astPair (astProject1 (astProject2 astd2))+                                      (astFromS' @eyFrom eftk+                                         (astProject2 (astProject2 astd2)))))+                    vard vd+            in subbed+          ftkr2 = FTKProduct+                    (adFTK $ FTKProduct accftk bftk)+                    (FTKProduct accftk eftkFrom)+          varr2 =+            mkAstVarName ftkr2 (varNameToBounds varr) (varNameToAstVarId varr)+          astr2 = astVar varr2+          vr2 =+            let subbed =+                  substituteAst+                    (astPair (astProject1 astr2)+                             (astPair (astProject1 (astProject2 astr2))+                                      (astFromS' @eyFrom eftk+                                         (astProject2 (astProject2 astr2)))))+                    varr vr+            in astSFrom' @(ADTensorKind (TKProduct accy ey))+                        (adSTK $ STKProduct accstk eftkFromSTK)+                        subbed+      in astMapAccumRDer k bftk eftkFrom (AstLambda varf2 vf2)+                                         (AstLambda vard2 vd2)+                                         (AstLambda varr2 vr2)+                                         acc0 esFrom+astMapAccumRDer k bftk eftk f df rf acc0 es =+  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es++astMapAccumLDer+  :: forall accy by ey k s. AstSpan s+  => SNat k+  -> FullShapeTK by+  -> FullShapeTK ey+  -> AstHFun s s+             (TKProduct accy ey) (TKProduct accy by)+  -> AstHFun s s+             (TKProduct (ADTensorKind (TKProduct accy ey))+                        (TKProduct accy ey))+             (ADTensorKind (TKProduct accy by))+  -> AstHFun s s+             (TKProduct (ADTensorKind (TKProduct accy by))+                        (TKProduct accy ey))+             (ADTensorKind (TKProduct accy ey))+  -> AstTensor AstMethodLet s accy+  -> AstTensor AstMethodLet s (BuildTensorKind k ey)+  -> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))+astMapAccumLDer k bftk eftk (AstLambda varf vf)+                            (AstLambda vard vd)+                            (AstLambda varr vr)+                (AstFromS' @accyFrom accftk acc0From) es =+  let accftkFrom = ftkAst acc0From+      accFromSTK = ftkToSTK accftkFrom+      ftkf2 = FTKProduct accftkFrom eftk+      varf2 = mkAstVarName ftkf2 (varNameToBounds varf) (varNameToAstVarId varf)+      astf2 = astVar varf2+      vf2 =+        let subbed =+              substituteAst+                (astPair (astFromS' @accyFrom accftk (astProject1 astf2))+                         (astProject2 astf2))+                varf vf+        in astSFrom' @(TKProduct accy by)+                    (STKProduct accFromSTK (ftkToSTK bftk))+                    subbed+      ftkd2 = FTKProduct+                (adFTK $ FTKProduct accftkFrom eftk)+                (FTKProduct accftkFrom eftk)+      vard2 = mkAstVarName ftkd2 (varNameToBounds vard) (varNameToAstVarId vard)+      astd2 = astVar vard2+      vd2 =+        let subbed =+              substituteAst+                (astPair (astPair (astFromS' @(ADTensorKind accyFrom)+                                     (adFTK accftk)+                                     (astProject1 (astProject1 astd2)))+                                  (astProject2 (astProject1 astd2)))+                         (astPair (astFromS' @accyFrom accftk+                                     (astProject1 (astProject2 astd2)))+                                  (astProject2 (astProject2 astd2))))+                vard vd+        in astSFrom' @(ADTensorKind (TKProduct accy by))+                    (adSTK $ STKProduct accFromSTK (ftkToSTK bftk))+                    subbed+      ftkr2 = FTKProduct+                (adFTK $ FTKProduct accftkFrom bftk)+                (FTKProduct accftkFrom eftk)+      varr2 = mkAstVarName ftkr2 (varNameToBounds varr) (varNameToAstVarId varr)+      astr2 = astVar varr2+      vr2 =+        let subbed =+              substituteAst+                (astPair (astPair (astFromS' @(ADTensorKind accyFrom)+                                     (adFTK accftk)+                                     (astProject1 (astProject1 astr2)))+                                  (astProject2 (astProject1 astr2)))+                         (astPair (astFromS' @accyFrom accftk+                                     (astProject1 (astProject2 astr2)))+                                  (astProject2 (astProject2 astr2))))+                varr vr+        in astSFrom' @(ADTensorKind (TKProduct accy ey))+                    (adSTK $ STKProduct accFromSTK (ftkToSTK eftk))+                    subbed+  in astFromS' @(TKProduct accyFrom (BuildTensorKind k by))+               (FTKProduct accftk (buildFTK k bftk))+     $ astMapAccumLDer k bftk eftk (AstLambda varf2 vf2)+                                   (AstLambda vard2 vd2)+                                   (AstLambda varr2 vr2)+                                   acc0From es+astMapAccumLDer k bftk eftk (AstLambda varf vf)+                            (AstLambda vard vd)+                            (AstLambda varr vr)+                acc0 (AstFromS' @esShsFrom _esShsFTK esFrom) =+  let accftk = ftkAst acc0+      accstk = ftkToSTK accftk+      esShsFrom = ftkAst esFrom+      esShsFromSTK = ftkToSTK esShsFrom+  in case razeSTK esShsFromSTK of+    (eftkFromSTK :: SingletonTK eyFrom) ->+      gcastWith (unsafeCoerceRefl+                 :: BuildTensorKind k eyFrom :~: esShsFrom) $+      let eftkFrom = razeFTK k eftkFromSTK esShsFrom+          ftkf2 = FTKProduct accftk eftkFrom+          varf2 =+            mkAstVarName ftkf2 (varNameToBounds varf) (varNameToAstVarId varf)+          astf2 = astVar varf2+          vf2 =+            let subbed =+                  substituteAst+                    (astPair (astProject1 astf2)+                             (astFromS' @eyFrom eftk (astProject2 astf2)))+                    varf vf+            in subbed+          ftkd2 = FTKProduct+                    (adFTK $ FTKProduct accftk eftkFrom)+                    (FTKProduct accftk eftkFrom)+          vard2 =+            mkAstVarName ftkd2 (varNameToBounds vard) (varNameToAstVarId vard)+          astd2 = astVar vard2+          vd2 =+            let subbed =+                  substituteAst+                    (astPair (astPair (astProject1 (astProject1 astd2))+                                      (astFromS' @(ADTensorKind eyFrom)+                                         (adFTK eftk)+                                         (astProject2 (astProject1 astd2))))+                             (astPair (astProject1 (astProject2 astd2))+                                      (astFromS' @eyFrom eftk+                                         (astProject2 (astProject2 astd2)))))+                    vard vd+            in subbed+          ftkr2 = FTKProduct+                    (adFTK $ FTKProduct accftk bftk)+                    (FTKProduct accftk eftkFrom)+          varr2 =+            mkAstVarName ftkr2 (varNameToBounds varr) (varNameToAstVarId varr)+          astr2 = astVar varr2+          vr2 =+            let subbed =+                  substituteAst+                    (astPair (astProject1 astr2)+                             (astPair (astProject1 (astProject2 astr2))+                                      (astFromS' @eyFrom eftk+                                         (astProject2 (astProject2 astr2)))))+                    varr vr+            in astSFrom' @(ADTensorKind (TKProduct accy ey))+                        (adSTK $ STKProduct accstk eftkFromSTK)+                        subbed+      in astMapAccumLDer k bftk eftkFrom (AstLambda varf2 vf2)+                                         (AstLambda vard2 vd2)+                                         (AstLambda varr2 vr2)+                                         acc0 esFrom+astMapAccumLDer k bftk eftk f df rf acc0 es =+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es++astApply :: forall x z s1 s. (AstSpan s1, AstSpan s)+         => AstHFun s1 s x z -> AstTensor AstMethodLet s1 x+         -> AstTensor AstMethodLet s z+astApply (AstLambda !var !v) u = astLet var u v++astCond :: AstSpan s+        => AstBool AstMethodLet+        -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y+        -> AstTensor AstMethodLet s y+astCond (AstBoolConst b) v w = if b then v else w+astCond (Ast.AstBoolNot b) v w = astCond b w v+astCond b (Ast.AstFromPrimal v) (Ast.AstFromPrimal w) =+  Ast.AstFromPrimal $ astCond b v w+astCond b (Ast.AstFromDual v) (Ast.AstFromDual w) =+  Ast.AstFromDual $ astCond b v w+astCond b v@(AstFromS' FTKScalar _) w = Ast.AstCond b v w+-- We rely here on c and the other conversion being semantically equal.+astCond b (Ast.AstConvert c v) (AstFromS' _ w) =+  case matchingFTK (ftkAst v) (ftkAst w) of+    Just Refl -> astConvert c $ astCond b v w+    Nothing -> error "astCond: shapes don't match"+astCond b v w = Ast.AstCond b v w++-- Invariant: if the variable has bounds, the expression can only have+-- values within the bounds (regardless of what the `bounds` call would say).+astLet :: forall y z s s2. (AstSpan s, AstSpan s2)+       => AstVarName s y -> AstTensor AstMethodLet s y+       -> AstTensor AstMethodLet s2 z+       -> AstTensor AstMethodLet s2 z+astLet _var _u v@Ast.AstConcreteK{} = v+astLet _var _u v@Ast.AstConcreteS{} = v+astLet _var _u v@Ast.AstIotaS{} = v+astLet var u v@(Ast.AstVar var2) =+  if varNameToAstVarId var2 == varNameToAstVarId var+  then case sameAstSpan @s @s2 of+    Just Refl -> case testEquality var var2 of+      Just Refl -> u+      _ -> error "astLet: wrong variable types at AstVar"+    _ -> error "astLet: wrong span at AstVar"+  else v+astLet var u v@(Ast.AstPrimalPart (Ast.AstVar var2)) =  -- a common noop+  if varNameToAstVarId var2 == varNameToAstVarId var+  then case sameAstSpan @s @FullSpan of+    Just Refl -> case testEquality var var2 of+      Just Refl -> astPrimalPart u+      _ -> error "astLet: wrong variable types at AstPrimalPart"+    _ -> error "astLet: wrong span at AstPrimalPart"+  else v+astLet var u v@(Ast.AstDualPart (Ast.AstVar var2)) =  -- a noop+  if varNameToAstVarId var2 == varNameToAstVarId var+  then case sameAstSpan @s @FullSpan of+    Just Refl -> case testEquality var var2 of+      Just Refl -> astDualPart u+      _ -> error "astLet: wrong variable types at AstDualPart"+    _ -> error "astLet: wrong span at AstDualPart"+  else v+astLet var u v | astIsSmall True u =+  substituteAst u var v+astLet var (Ast.AstPair u1 u2) v =+  astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->+    substituteAst (Ast.AstPair ast1 ast2) var v+astLet var (Ast.AstFromPrimal (Ast.AstPair u1 u2)) v =+  astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->+    substituteAst (Ast.AstFromPrimal (Ast.AstPair ast1 ast2)) var v+astLet var (Ast.AstFromDual (Ast.AstPair u1 u2)) v =+  astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->+    substituteAst (Ast.AstFromDual (Ast.AstPair ast1 ast2)) var v+astLet var (Ast.AstLet varN uN (Ast.AstPair u1 u2)) v =+  astLet varN uN+  $ astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->+      substituteAst (Ast.AstPair ast1 ast2) var v+astLet var (Ast.AstFromPrimal (Ast.AstLet varN uN (Ast.AstPair u1 u2))) v =+  astLet varN uN+  $ astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->+      substituteAst (Ast.AstFromPrimal (Ast.AstPair ast1 ast2)) var v+astLet var (Ast.AstFromDual (Ast.AstLet varN uN (Ast.AstPair u1 u2))) v =+  astLet varN uN+  $ astLetFun u1 $ \ !ast1 -> astLetFun u2 $ \ !ast2 ->+      substituteAst (Ast.AstFromDual (Ast.AstPair ast1 ast2)) var v+-- This is a common case, e.g., from representing conditionals.+astLet var (Ast.AstFromVector snat stk u) v | V.length u == 2 =+  astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->+    substituteAst (Ast.AstFromVector snat stk $ fromList [ast1, ast2]) var v+astLet var (Ast.AstFromPrimal+              (Ast.AstFromVector snat stk u)) v | V.length u == 2 =+  astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->+    substituteAst (Ast.AstFromPrimal (Ast.AstFromVector snat stk+                                      $ fromList [ast1, ast2])) var v+astLet var (Ast.AstFromDual+              (Ast.AstFromVector snat stk u)) v | V.length u == 2 =+  astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->+    substituteAst (Ast.AstFromDual (Ast.AstFromVector snat stk+                                    $ fromList [ast1, ast2])) var v+astLet var (Ast.AstLet varN uN+              (Ast.AstFromVector snat stk u)) v | V.length u == 2 =+  astLet varN uN+  $ astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->+      substituteAst (Ast.AstFromVector snat stk $ fromList [ast1, ast2]) var v+astLet var (Ast.AstFromPrimal+              (Ast.AstLet varN uN+                 (Ast.AstFromVector snat stk u))) v | V.length u == 2 =+  astLet varN uN+  $ astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->+      substituteAst (Ast.AstFromPrimal (Ast.AstFromVector snat stk+                                        $ fromList [ast1, ast2])) var v+astLet var (Ast.AstFromDual+              (Ast.AstLet varN uN+                 (Ast.AstFromVector snat stk u))) v | V.length u == 2 =+  astLet varN uN+  $ astLetFun (u V.! 0) $ \ !ast1 -> astLetFun (u V.! 1) $ \ !ast2 ->+      substituteAst (Ast.AstFromDual (Ast.AstFromVector snat stk+                                      $ fromList [ast1, ast2])) var v+astLet var (Ast.AstReplicate snat stk a) v =+  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)+      ast = Ast.AstReplicate snat stk $ astVar var2+  in astLet var2 a (substituteAst ast var v)+astLet var (Ast.AstFromPrimal (Ast.AstReplicate snat stk a)) v =+  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)+      ast = Ast.AstFromPrimal (Ast.AstReplicate snat stk $ astVar var2)+  in astLet var2 a (substituteAst ast var v)+astLet var (Ast.AstFromDual (Ast.AstReplicate snat stk a)) v =+  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)+      ast = Ast.AstFromDual (Ast.AstReplicate snat stk $ astVar var2)+  in astLet var2 a (substituteAst ast var v)+astLet var (Ast.AstTransposeS perm a) v =+  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)+      ast = Ast.AstTransposeS perm $ astVar var2+  in astLet var2 a (substituteAst ast var v)+astLet var (Ast.AstFromPrimal (Ast.AstTransposeS perm a)) v =+  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)+      ast = Ast.AstFromPrimal (Ast.AstTransposeS perm $ astVar var2)+  in astLet var2 a (substituteAst ast var v)+astLet var (Ast.AstFromDual (Ast.AstTransposeS perm a)) v =+  let var2 = mkAstVarName (ftkAst a) Nothing (varNameToAstVarId var)+      ast = Ast.AstFromDual (Ast.AstTransposeS perm $ astVar var2)+  in astLet var2 a (substituteAst ast var v)+astLet var u@(AstFromS' FTKScalar _) v = Ast.AstLet var u v+astLet var (Ast.AstConvert c a) v | checkAstFromS c a =+  let var2 =+        mkAstVarName (ftkAst a) (varNameToBounds var) (varNameToAstVarId var)+      ast = astConvert c $ astVar var2+  in astLet var2 a (substituteAst ast var v)+astLet var u (Ast.AstFromPrimal v0) = Ast.AstFromPrimal $ astLet var u v0+astLet var u (Ast.AstFromDual v0) = Ast.AstFromDual $ astLet var u v0+astLet var u v@(AstFromS' FTKScalar _) = Ast.AstLet var u v+astLet var u (Ast.AstConvert c v) | checkAstFromS c v =+  astConvert c $ astLet var u v+astLet var u v = Ast.AstLet var u v++astPrimalPart :: AstTensor AstMethodLet FullSpan y+              -> AstTensor AstMethodLet PrimalSpan y+astPrimalPart t = case t of+  Ast.AstPair t1 t2 -> astPair (astPrimalPart t1) (astPrimalPart t2)+  Ast.AstProject1 v -> astProject1 (astPrimalPart v)+  Ast.AstProject2 v -> astProject2 (astPrimalPart v)+  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map astPrimalPart l)+  Ast.AstSum snat stk v ->+    astSum snat stk (astPrimalPart v)+  Ast.AstReplicate snat stk v ->+    astReplicate snat stk (astPrimalPart v)+  {- Ast.AstMapAccumRDer k bftk eftk (AstLambda varf vf)+                                  (AstLambda vard vd)+                                  (AstLambda varr vr) acc0 es ->+    astMapAccumRDer k bftk eftk f df rf+                    (astPrimalPart acc0) (astPrimalPart es)+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    astMapAccumLDer k bftk eftk f df rf+                    (astPrimalPart acc0) (astPrimalPart es) -}+  Ast.AstMapAccumRDer{} -> Ast.AstPrimalPart t  -- TODO+  Ast.AstMapAccumLDer{} -> Ast.AstPrimalPart t+  Ast.AstApply (AstLambda !var !v) ll ->+    astApply (AstLambda var (astPrimalPart v)) ll+  Ast.AstVar{} -> Ast.AstPrimalPart t  -- the only normal form+  Ast.AstCond b a2 a3 -> astCond b (astPrimalPart a2) (astPrimalPart a3)+  Ast.AstBuild1 k stk (var, v) ->+    let !v2 = astPrimalPart v+    in Ast.AstBuild1 k stk (var, v2)++  Ast.AstLet var u v -> astLet var u (astPrimalPart v)++  Ast.AstFromPrimal v -> v+  Ast.AstFromDual v ->+    let ftk = ftkAst v+    in astConcrete ftk (tdefTarget ftk)++  AstPlusK u v -> astPrimalPart u + astPrimalPart v+  AstTimesK u v -> astPrimalPart u * astPrimalPart v+  Ast.AstN1K NegateOp u -> negate (astPrimalPart u)+  Ast.AstN1K AbsOp u -> abs (astPrimalPart u)+  Ast.AstN1K SignumOp u -> signum (astPrimalPart u)+  Ast.AstR1K opCode u -> Ast.AstR1K opCode (astPrimalPart u)+  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (astPrimalPart u) (astPrimalPart v)+  Ast.AstI2K QuotOp u v -> quotH (astPrimalPart u) (astPrimalPart v)+  Ast.AstI2K RemOp u v -> remH (astPrimalPart u) (astPrimalPart v)+  Ast.AstCastK v -> astCastK $ astPrimalPart v++  AstPlusS u v -> astPrimalPart u + astPrimalPart v+  AstTimesS u v -> astPrimalPart u * astPrimalPart v+  Ast.AstN1S NegateOp u -> negate (astPrimalPart u)+  Ast.AstN1S AbsOp u -> abs (astPrimalPart u)+  Ast.AstN1S SignumOp u -> signum (astPrimalPart u)+  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astPrimalPart u)+  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (astPrimalPart u)+                                             (astPrimalPart v)+  Ast.AstI2S QuotOp u v -> quotH (astPrimalPart u) (astPrimalPart v)+  Ast.AstI2S RemOp u v -> remH (astPrimalPart u) (astPrimalPart v)+  Ast.AstCastS v -> astCastS $ astPrimalPart v++  Ast.AstIndexS shn v ix ->+    astIndexS shn (astPrimalPart v) ix+  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->+    astScatterS @shm @shn @shp shn (astPrimalPart v) (vars, ix)+  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->+    astGatherS @shm @shn @shp shn (astPrimalPart v) (vars, ix)+  Ast.AstAppendS x y -> astAppendS (astPrimalPart x) (astPrimalPart y)+  Ast.AstSliceS i n k v -> astSliceS i n k (astPrimalPart v)+  Ast.AstReverseS v -> astReverseS (astPrimalPart v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astPrimalPart v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astPrimalPart v)++  -- All conversions need to stay down here to cancel out.+  Ast.AstConvert{} -> Ast.AstPrimalPart t++  -- These should not appear in this context unless via wacky tests.+  Ast.AstSum0S{} -> Ast.AstPrimalPart t+  Ast.AstDot0S{} -> Ast.AstPrimalPart t+  Ast.AstDot1InS{} -> Ast.AstPrimalPart t+  Ast.AstMatmul2S{} -> Ast.AstPrimalPart t++-- Note how this can't be pushed down into, say, multiplication, because it+-- multiplies the dual part by the primal part. Addition is fine, though.+astDualPart :: AstTensor AstMethodLet FullSpan y+            -> AstTensor AstMethodLet DualSpan y+astDualPart t = case t of+  Ast.AstPair t1 t2 -> astPair (astDualPart t1) (astDualPart t2)+  Ast.AstProject1 v -> astProject1 (astDualPart v)+  Ast.AstProject2 v -> astProject2 (astDualPart v)+  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map astDualPart l)+  Ast.AstSum snat stk v ->+    astSum snat stk (astDualPart v)+  Ast.AstReplicate snat stk v ->+    astReplicate snat stk (astDualPart v)+  {- Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->+    astMapAccumRDer k bftk eftk f df rf+                    (astDualPart acc0) (astDualPart es)+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    astMapAccumLDer k bftk eftk f df rf+                    (astDualPart acc0) (astDualPart es) -}+  Ast.AstMapAccumRDer{} -> Ast.AstDualPart t  -- TODO+  Ast.AstMapAccumLDer{} -> Ast.AstDualPart t+  Ast.AstApply (AstLambda !var !v) ll ->+    astApply (AstLambda var (astDualPart v)) ll+  Ast.AstVar{} -> Ast.AstDualPart t+  Ast.AstCond b a2 a3 -> astCond b (astDualPart a2) (astDualPart a3)+  Ast.AstBuild1 k stk (var, v) ->+    let !v2 = astDualPart v+    in Ast.AstBuild1 k stk (var, v2)++  Ast.AstLet var u v -> astLet var u (astDualPart v)++  Ast.AstFromPrimal v ->+    let ftk = ftkAst v+    in Ast.AstDualPart $ Ast.AstFromPrimal+       $ astConcrete ftk (tdefTarget ftk)+           -- let's hope this is smaller than v+  Ast.AstFromDual v -> v++  AstPlusK u v -> astDualPart u + astDualPart v+  -- This one is mathematically wrong, dual numbers don't mult like that:+  -- AstTimesK u v -> astDualPart u * astDualPart v+  Ast.AstN1K NegateOp u -> negate (astDualPart u)+  {- Some of these are wrong, so let's be conservative:+  Ast.AstN1K AbsOp u -> abs (astDualPart u)+  Ast.AstN1K SignumOp u -> signum (astDualPart u)+  Ast.AstR1K opCode u -> Ast.AstR1K opCode (astDualPart u)+  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (astDualPart u) (astDualPart v)+  Ast.AstI2K QuotOp u v -> quotH (astDualPart u) (astDualPart v)+  Ast.AstI2K RemOp u v -> remH (astDualPart u) (astDualPart v)+  -}+  Ast.AstCastK v -> astCastK $ astDualPart v++  AstPlusS u v -> astDualPart u + astDualPart v+  -- This one is mathematically wrong, dual numbers don't mult like that:+  -- AstTimesS u v -> astDualPart u * astDualPart v+  Ast.AstN1S NegateOp u -> negate (astDualPart u)+  {- Some of these are wrong, so let's be conservative:+  Ast.AstN1S AbsOp u -> abs (astDualPart u)+  Ast.AstN1S SignumOp u -> signum (astDualPart u)+  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astDualPart u)+  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (astDualPart u)+                                             (astDualPart v)+  Ast.AstI2S opCode u v -> Ast.AstI2S opCode (astDualPart u)+                                             (astDualPart v)+  -}+  Ast.AstCastS v -> astCastS $ astDualPart v++  Ast.AstIndexS shn v ix ->+    astIndexS shn (astDualPart v) ix+  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->+    astScatterS @shm @shn @shp shn (astDualPart v) (vars, ix)+  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->+    astGatherS @shm @shn @shp shn (astDualPart v) (vars, ix)+  Ast.AstAppendS x y -> astAppendS (astDualPart x) (astDualPart y)+  Ast.AstSliceS i n k v -> astSliceS i n k (astDualPart v)+  Ast.AstReverseS v -> astReverseS (astDualPart v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astDualPart v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astDualPart v)++  -- All conversions need to stay down here to cancel out.+  Ast.AstConvert{} -> Ast.AstDualPart t++  -- These should not appear in this context unless via wacky tests.+  Ast.AstSum0S{} -> Ast.AstDualPart t+  Ast.AstDot0S{} -> Ast.AstDualPart t+  Ast.AstDot1InS{} -> Ast.AstDualPart t+  Ast.AstMatmul2S{} -> Ast.AstDualPart t++  _ -> Ast.AstDualPart t++astConcreteK :: GoodScalar r+             => Concrete (TKScalar r)+             -> AstTensor AstMethodLet PrimalSpan (TKScalar r)+astConcreteK = AstConcreteK . unConcrete++astFloorK :: (GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2)+          => AstTensor AstMethodLet PrimalSpan (TKScalar r1)+          -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)+astFloorK t = case t of+  Ast.AstLet var u v -> astLet var u (astFloorK v)+  -- This increases work and term size, because conditional is eager.+  -- Ast.AstCond b a2 a3 -> Ast.AstCond b (astFloorK a2) (astFloorK a3)+  -- These values are small, so we can simplify them ASAP.+  AstConcreteK k -> astConcreteK (tkfloor $ Concrete k)+  Ast.AstFloorK v -> astFloorK v+  Ast.AstFromIntegralK v -> astFromIntegralK v+  Ast.AstCastK v -> astFloorK v+  _ -> Ast.AstFloorK t++-- Beware that increasing the number of calls to this constructor+-- sometimes increases runtime, because not enough copies cancel out.+-- Hence the commented out rules below.+astFromIntegralK :: forall r1 r2. (GoodScalar r1, GoodScalar r2, Integral r1)+                 => AstTensor AstMethodLet PrimalSpan (TKScalar r1)+                 -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)+astFromIntegralK t = case t of+  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t+  Ast.AstLet var u v -> astLet var u (astFromIntegralK v)+  AstConcreteK k -> astConcreteK (tkfromIntegral $ Concrete k)+  Ast.AstN1K NegateOp u -> negate (astFromIntegralK u)+  Ast.AstN1K AbsOp u -> abs (astFromIntegralK u)+  Ast.AstN1K SignumOp u -> signum (astFromIntegralK u)+  Ast.AstFromIntegralK v -> astFromIntegralK v+  _ -> Ast.AstFromIntegralK t++astCastK :: forall r1 r2 s.+            (GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2, AstSpan s)+         => AstTensor AstMethodLet s (TKScalar r1)+         -> AstTensor AstMethodLet s (TKScalar r2)+astCastK t = case t of+  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t+  Ast.AstLet var u v -> astLet var u (astCastK v)+  AstConcreteK k -> astConcreteK (tkcast $ Concrete k)+  -- TODO: which should go deeper, casts or fromPrimal? Or maybe alternate+  -- in different phases to make sure both can cancel out?+  -- Rethink. For now, astFromPrimalis not called, to avoid loops.+  -- The same with many others+  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astCastK v+  Ast.AstFromDual v -> Ast.AstFromDual $ astCastK v+  Ast.AstN1K NegateOp u -> negate (astCastK u)+  Ast.AstN1K AbsOp u -> abs (astCastK u)+  Ast.AstN1K SignumOp u -> signum (astCastK u)+--  Ast.AstR1K opCode u -> Ast.AstR1K opCode (astCastK u)+  Ast.AstFromIntegralK v -> astFromIntegralK v+  Ast.AstCastK v -> astCastK v+  _ -> Ast.AstCastK t++astConcreteS :: GoodScalar r+             => Concrete (TKS sh r)+             -> AstTensor AstMethodLet PrimalSpan (TKS sh r)+astConcreteS = AstConcreteS . unConcrete++astFloorS :: forall r1 r2 sh.+             (GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2)+          => AstTensor AstMethodLet PrimalSpan (TKS sh r1)+          -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)+astFloorS t = case t of+  _ | FTKS (snat :$$ sh2) _ <- ftkAst t+    , Just u <- unRepl1 t ->+      astReplicate snat (STKS sh2 STKScalar) (astFloorS u)+  Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->+    Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, astFloorS v)+  Ast.AstBuild1 snat STKScalar (var, v) ->+    Ast.AstBuild1 snat STKScalar (var, astFloorK v)+  Ast.AstLet var u v -> astLet var u (astFloorS v)+  Ast.AstScatterS shn v (vars, ix) ->+    astScatterS shn (astFloorS v) (vars, ix)+  Ast.AstGatherS shn v (vars, ix) ->+    astGatherS shn (astFloorS v) (vars, ix)+  Ast.AstIotaS snat -> Ast.AstIotaS snat+  Ast.AstReverseS v -> astReverseS (astFloorS v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astFloorS v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astFloorS v)+  AstSFromK' a -> astSFromK' (astFloorK a)+  Ast.AstFloorS v -> astFloorS v+  Ast.AstFromIntegralS v -> astFromIntegralS v+  Ast.AstCastS v -> astFloorS v+  _ -> Ast.AstFloorS t++astFromIntegralS :: forall r1 r2 sh. (GoodScalar r1, GoodScalar r2, Integral r1)+                 => AstTensor AstMethodLet PrimalSpan (TKS sh r1)+                 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)+astFromIntegralS t = case t of+  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t+  _ | FTKS (snat :$$ sh2) _ <- ftkAst t+    , Just u <- unRepl1 t ->+      astReplicate snat (STKS sh2 STKScalar) (astFromIntegralS u)+--  Ast.AstFromVector snat (STKS sh STKScalar) l ->+--   astFromVector snat (STKS sh STKScalar) (V.map astFromIntegralS l)+--  Ast.AstFromVector snat STKScalar l ->+--   astFromVector snat STKScalar (V.map astFromIntegralK l)+  Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->+    Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, astFromIntegralS v)+  Ast.AstBuild1 snat STKScalar (var, v) ->+    Ast.AstBuild1 snat STKScalar (var, astFromIntegralK v)+  Ast.AstLet var u v -> astLet var u (astFromIntegralS v)+  Ast.AstN1S NegateOp u -> negate (astFromIntegralS u)+  Ast.AstN1S AbsOp u -> abs (astFromIntegralS u)+  Ast.AstN1S SignumOp u -> signum (astFromIntegralS u)+  Ast.AstFromIntegralS v -> astFromIntegralS v+--  Ast.AstIndexS shn v ix -> astIndexS shn (astFromIntegralS v) ix+    -- increases work; also index goes into fromIntegral, so we'd loop+  Ast.AstScatterS shn v (vars, ix) ->+    astScatterS shn (astFromIntegralS v) (vars, ix)+  Ast.AstGatherS shn v (vars, ix) ->+    astGatherS shn (astFromIntegralS v) (vars, ix)+  Ast.AstIotaS snat -> Ast.AstIotaS snat+--  Ast.AstSliceS i n k v -> astSliceS i n k (astFromIntegralS v)+  Ast.AstReverseS v -> astReverseS (astFromIntegralS v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astFromIntegralS v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astFromIntegralS v)+  AstSFromK' a -> astSFromK' (astFromIntegralK a)+  _ -> Ast.AstFromIntegralS t++astCastS :: forall r1 r2 s sh.+            (GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2, AstSpan s)+         => AstTensor AstMethodLet s (TKS sh r1)+         -> AstTensor AstMethodLet s (TKS sh r2)+astCastS t = case t of+  _ | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) -> t+  _ | FTKS (snat :$$ sh2) _ <- ftkAst t+    , Just u <- unRepl1 t ->+      astReplicate snat (STKS sh2 STKScalar) (astCastS u)+--  Ast.AstFromVector snat (STKS sh STKScalar) l ->+--   astFromVector snat (STKS sh STKScalar) (V.map astCastS l)+--  Ast.AstFromVector snat STKScalar l ->+--   astFromVector snat STKScalar (V.map astCastK l)+  {- This (and other similar rules) is bad, because it changes semantics+     and also impacts performance negatively (a is larger than sum a):+  Ast.AstSum snat (STKS sh STKScalar) a ->+    astSum snat (STKS sh STKScalar) (astCastS a) -}+  Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, v) ->+    Ast.AstBuild1 snat (STKS sh2 STKScalar) (var, astCastS v)+  Ast.AstBuild1 snat STKScalar (var, v) ->+    Ast.AstBuild1 snat STKScalar (var, astCastK v)+  Ast.AstLet var u v -> astLet var u (astCastS v)+  Ast.AstPrimalPart a -> Ast.AstPrimalPart $ astCastS a+  Ast.AstDualPart a -> Ast.AstDualPart $ astCastS a+  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astCastS v+  Ast.AstFromDual v -> Ast.AstFromDual $ astCastS v+  Ast.AstN1S NegateOp u -> negate (astCastS u)+  Ast.AstN1S AbsOp u -> abs (astCastS u)+  Ast.AstN1S SignumOp u -> signum (astCastS u)+  Ast.AstFromIntegralS v -> astFromIntegralS v+  Ast.AstCastS v -> astCastS v+--  Ast.AstIndexS shn v ix -> astIndexS shn (astCastS v) ix+    -- increases work; also index goes into fromIntegral, so we'd loop+  Ast.AstScatterS shn v (vars, ix) -> astScatterS shn (astCastS v) (vars, ix)+  Ast.AstGatherS shn v (vars, ix) -> astGatherS shn (astCastS v) (vars, ix)+--  Ast.AstMinIndexS v -> Ast.AstMinIndexS (astCastS v)+  Ast.AstIotaS snat -> Ast.AstIotaS snat+--  Ast.AstSliceS i n k v -> astSliceS i n k (astCastS v)+  Ast.AstReverseS v -> astReverseS (astCastS v)+  Ast.AstTransposeS perm v -> astTransposeS perm (astCastS v)+  Ast.AstReshapeS sh v -> astReshapeS sh (astCastS v)+  AstSFromK' a -> astSFromK' (astCastK a)+  _ -> Ast.AstCastS t++astIndexS+  :: forall shm shn s r. AstSpan s+  => ShS shn+  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r) -> AstIxS AstMethodLet shm+  -> AstTensor AstMethodLet s (TKS2 shn r)+astIndexS = astIndexKnobsS defaultKnobs++astIndexKnobsS+  :: forall shm shn s r. AstSpan s+  => SimplifyKnobs+  -> ShS shn+  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)+  -> AstIxS AstMethodLet shm+  -> AstTensor AstMethodLet s (TKS2 shn r)+astIndexKnobsS _ _ v0 ZIS = v0+astIndexKnobsS _ shn v0 (i1 :.$ _)+  | let (lb, ub) = bounds i1+-- this doesn't work in GHC 9.10:+--      FTKS (snat :$$ _) x = ftkAst v0+  , FTKS (snat :$$ _) x <- ftkAst v0+  , ub < 0 || lb >= fromInteger (fromSNat snat) =+    let ftk = FTKS shn x+    in fromPrimal $ astConcrete ftk (tdefTarget ftk)+astIndexKnobsS knobs shn v0 (Ast.AstCond b i1 i2 :.$ rest0)+  | knobPhase knobs `notElem` [PhaseUnspecified, PhaseVectorization] =+      -- don't undo vectorization tweaks+    astLetFun v0 $ \v ->+    shareIx rest0 $ \rest ->+      astCond b (astIndexKnobsS knobs shn v (i1 :.$ rest))+                (astIndexKnobsS knobs shn v (i2 :.$ rest))+astIndexKnobsS knobs shn v0 ix@((:.$) @in1 @shm1 i1 rest1) =+ let FTKS _ x = ftkAst v0+     astIndex+       :: forall shm' shn' s' r'. AstSpan s'+       => ShS shn'+       -> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') r')+       -> AstIxS AstMethodLet shm'+       -> AstTensor AstMethodLet s' (TKS2 shn' r')+     astIndex shn' v2 ix2 = astIndexKnobsS (deVect knobs) shn' v2 ix2+     astGather+       :: forall shm' shn' shp'.+          ShS shn'+       -> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') r)+       -> (AstVarListS shm', AstIxS AstMethodLet shp')+       -> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') r)+     astGather shn' v2 (vars2, ix2) =+       astGatherKnobsS @shm' @shn' @shp' (deVect knobs) shn' v2 (vars2, ix2)+ in case v0 of+  Ast.AstProject1{} -> Ast.AstIndexS shn v0 ix+  Ast.AstProject2{} -> Ast.AstIndexS shn v0 ix+  Ast.AstFromVector _ STKS{} l | AstConcreteK it <- i1 ->+    let i = fromIntegral it+    in astIndex shn (l V.! i) rest1+  Ast.AstFromVector _ STKScalar l | AstConcreteK it <- i1, ZIS <- rest1 ->+    let i = fromIntegral it+    in astSFromK' (l V.! i)+  Ast.AstFromVector{} | ZIS <- rest1 ->  -- normal form+    Ast.AstIndexS shn v0 ix+  Ast.AstFromVector snat STKS{} l ->+    shareIx rest1 $ \ !rest2 ->+      Ast.AstIndexS @'[in1] @shn shn (astFromVector snat (STKS shn (ftkToSTK x))+                                      $ V.map (\a -> astIndex shn a rest2) l)+                    (i1 :.$ ZIS)+  Ast.AstSum snat@(SNat @n1) STKS{} v ->+    let perm3 = backpermCycle $ shsLength (shsFromIxS ix) + 1+    in Permutation.permFromList perm3 $ \(perm :: Permutation.Perm perm3P) ->+         gcastWith (unsafeCoerceRefl+                    :: Compare (Rank perm3P) (Rank (n1 : shm ++ shn))+                       :~: LT) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix perm3P (n1 : (shm ++ shn))+                       :~: shm ++ (n1 : shn)) $+         fromMaybe (error "astIndexKnobsS: impossible non-permutation")+         $ Permutation.permCheckPermutation perm+         $ astSum snat (STKS shn (ftkToSTK x))+         $ astIndex @shm @(n1 : shn) (snat :$$ shn)+                    (astTransposeS @perm3P @(n1 : shm ++ shn) perm v)+                    ix+  Ast.AstReplicate _ STKS{} v ->+    let ftk = FTKS shn x+        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+    in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of+      AstBoolConst b -> if b then astIndex shn v rest1 else defArr+      _ -> Ast.AstIndexS shn v0 ix+  Ast.AstReplicate _ STKScalar v | ZIS <- rest1 ->+    let ftk = FTKS shn x+        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+    in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of+      AstBoolConst b -> if b then astSFromK' v else defArr+      _ -> Ast.AstIndexS shn v0 ix+  Ast.AstApply{} -> Ast.AstIndexS shn v0 ix+  Ast.AstVar{} -> Ast.AstIndexS shn v0 ix+  Ast.AstCond b v w ->+    shareIx ix $ \ !ix2 ->+      astCond b (astIndexKnobsS knobs shn v ix2)+                (astIndexKnobsS knobs shn w ix2)+{- This is wrong: in a counterfactual case, astLet assigns OOB i to var2,+   violating the invariant about variables bounds:+  Ast.AstBuild1 (SNat @k) STKS{} (var2, v) ->+    let ftk = FTKS shn x+        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+    in astLetFunB i1 $ \i ->+    astCond (0 <=. i &&* i <=. valueOf @k - 1)+            (astIndex shn (astLet var2 i v) rest1)+            defArr -}+  Ast.AstBuild1 (SNat @k) STKS{} (var2, v) ->+    let ftk = FTKS shn x+        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+    in case 0 <=. i1 &&* i1 <=. valueOf @k - 1 of+      AstBoolConst b ->+        if b then astIndex shn (astLet var2 i1 v) rest1 else defArr+      _ -> Ast.AstIndexS shn v0 ix+  Ast.AstBuild1 (SNat @k) STKScalar (var2, v) | ZIS <- rest1 ->+    let ftk = FTKS shn x+        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+    in case 0 <=. i1 &&* i1 <=. valueOf @k - 1 of+      AstBoolConst b ->+        if b then astSFromK' $ astLet var2 i1 v else defArr+      _ -> Ast.AstIndexS shn v0 ix++  Ast.AstLet var u v -> astLet var u (astIndexKnobsS knobs shn v ix)++  Ast.AstPrimalPart{} -> Ast.AstIndexS shn v0 ix  -- must be a NF+  Ast.AstDualPart{} -> Ast.AstIndexS shn v0 ix+  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astIndexKnobsS knobs shn v ix+  Ast.AstFromDual v -> Ast.AstFromDual $ astIndexKnobsS knobs shn v ix++  AstPlusS u v ->+    shareIx ix $ \ !ix2 ->+    astIndexKnobsS knobs shn u ix2 + astIndexKnobsS knobs shn v ix2+  AstTimesS u v ->+    shareIx ix $ \ !ix2 ->+    astIndexKnobsS knobs shn u ix2 * astIndexKnobsS knobs shn v ix2+  Ast.AstN1S NegateOp u -> negate (astIndexKnobsS knobs shn u ix)+  Ast.AstN1S AbsOp u -> abs (astIndexKnobsS knobs shn u ix)+  Ast.AstN1S SignumOp u -> signum (astIndexKnobsS knobs shn u ix)+  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astIndexKnobsS knobs shn u ix)+  Ast.AstR2S opCode u v ->+    shareIx ix+    $ \ !ix2 -> Ast.AstR2S opCode (astIndexKnobsS knobs shn u ix2)+                                  (astIndexKnobsS knobs shn v ix2)+  Ast.AstI2S opCode u v ->+    shareIx ix+    $ \ !ix2 -> Ast.AstI2S opCode (astIndexKnobsS knobs shn u ix2)+                                  (astIndexKnobsS knobs shn v ix2)+  AstConcreteS a | AstConcreteK i <- i1 ->+    let u = withKnownShS (shsFromIxS rest1 `shsAppend` shn) $+            tsindex (Concrete a) (Concrete i :.$ ZIS)+    in astIndex shn (astConcreteS u) rest1+  AstConcreteS{} -> case unRepl1 v0 of+    Just u ->+      let ftk = FTKS shn x+          defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+      in case 0 <=. i1 &&* i1 <=. valueOf @in1 - 1 of+        AstBoolConst b -> if b then astIndex shn u rest1 else defArr+        _ -> Ast.AstIndexS shn v0 ix+    _ -> Ast.AstIndexS shn v0 ix+  Ast.AstFloorS v -> astFloorS $ astIndexKnobsS knobs shn v ix+  Ast.AstFromIntegralS v -> astFromIntegralS $ astIndexKnobsS knobs shn v ix+  Ast.AstCastS t -> astCastS $ astIndexKnobsS knobs shn t ix++  Ast.AstIndexS _ v (ix2 :: AstIxS AstMethodLet sh4)+    | Refl <- lemAppAssoc (Proxy @sh4) (Proxy @shm) (Proxy @shn) ->+      astIndexKnobsS knobs shn v (ix2 `ixsAppend` ix)+  Ast.AstScatterS @shm7 @shn7 @shp7 shn7 v (vars, AstIntVar var5 :.$ ix2)+    | AstIntVar var6 <- i1, var6 == var5 ->+        astIndex shn (astScatterS @shm7 @shn7 @(Tail shp7)+                                  shn7 v (vars, ix2)) rest1+  Ast.AstScatterS @shm7 @shn7 @shp7 shn7+                  v (vars, AstConcreteK i5 :.$ ix2)+    | AstConcreteK i6 <- i1 ->+        if i6 == i5+        then astIndex shn (astScatterS @shm7 @shn7 @(Tail shp7)+                                       shn7 v (vars, ix2)) rest1+        else let ftk = FTKS shn x+             in fromPrimal $ astConcrete ftk (tdefTarget ftk)+  -- AstScatter sh v (vars2, ZIR) ->+  --   AstScatter sh (astIndex (astTranspose perm3 v) ix) (vars2, ZIR)+  Ast.AstScatterS{} ->  -- normal form+    Ast.AstIndexS shn v0 ix+  -- This is not a possible normal form, but pattern needs to be exhaustive.+  Ast.AstGatherS @_ @_ @shp' _shn' v (ZS, ix2) ->+    gcastWith (unsafeCoerceRefl+              :: shp' ++ (in1 ': shm1) ++ shn+                 :~: shp' ++ (in1 ': shm1 ++ shn)) $+    astIndex @(shp' ++ shm) @shn shn v (ix2 `ixsAppend` ix)+  Ast.AstGatherS @_ @shn' @shp' shn'+                 v ((::$) @m71 @shm71 (Const var2) vars, ix2) ->+    gcastWith (unsafeCoerceRefl :: shm71 ++ shn' :~: shm1 ++ shn) $+    let ftk = FTKS shn x+        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+        w :: AstTensor AstMethodLet s (TKS2 (shm1 ++ shn) r)+        w = astGather @shm71 @shn' @shp' shn' v (vars, ix2)+        u = astLet var2 i1 $ astIndex @shm1 @shn shn w rest1+          -- this let makes it impossible to use astCond when i1 is OOB+    in case 0 <=. i1 &&* i1 <=. valueOf @m71 - 1 of+      AstBoolConst b -> if b then u else defArr+      _ -> Ast.AstIndexS shn v0 ix+  Ast.AstMinIndexS @n1 @shz v -> case ftkAst v of+    FTKS nsh _ -> case shsLast nsh of+     nl@(SNat @nl) ->+      let shnl = shn `shsAppend` (nl :$$ ZSS)+      in gcastWith (unsafeCoerceRefl+                    :: Permutation.Index 0 (shn ++ '[nl])+                       ': Drop 1 (shn ++ '[nl])+                       :~: shn ++ '[nl]) $+         gcastWith (unsafeCoerceRefl+                    :: Init (shn ++ '[nl]) :~: shn) $+         gcastWith (unsafeCoerceRefl+                    :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) $+         Ast.AstMinIndexS @(Head (shn ++ '[nl]))+                          @(Tail (shn ++ '[nl]))+         $ astIndexKnobsS @shm @(shn ++ '[nl]) knobs shnl v ix+  Ast.AstMaxIndexS @n1 @shz v -> case ftkAst v of+    FTKS nsh _ -> case shsLast nsh of+     nl@(SNat @nl) ->+      let shnl = shn `shsAppend` (nl :$$ ZSS)+      in gcastWith (unsafeCoerceRefl+                    :: Permutation.Index 0 (shn ++ '[nl])+                       ': Drop 1 (shn ++ '[nl])+                       :~: shn ++ '[nl]) $+         gcastWith (unsafeCoerceRefl+                    :: Init (shn ++ '[nl]) :~: shn) $+         gcastWith (unsafeCoerceRefl+                    :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) $+         Ast.AstMaxIndexS @(Head (shn ++ '[nl]))+                          @(Tail (shn ++ '[nl]))+         $ astIndexKnobsS @shm @(shn ++ '[nl]) knobs shnl v ix+  Ast.AstIotaS (SNat @k) -> case testEquality shn ZSS of+    Just Refl ->+      let ftk = FTKS ZSS x+          defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+      in astLetFunB i1 $ \i ->+        astCond (0 <=. i &&* i <=. valueOf @k - 1)+                (astFromIntegralS $ astSFromK' i)+                defArr+    _ -> error "astIndexKnobsS: shape not []"+  Ast.AstAppendS u v | FTKS (SNat @m :$$ _) _ <- ftkAst u ->+    astLetFunB i1 $ \i ->+    shareIx rest1 $ \ !rest2 ->+    let ulen = valueOf @m+        ix1 = i :.$ rest2+        ix2 = i - ulen :.$ rest2+    in case ulen <=. i of+      AstBoolConst b -> if b then astIndex shn v ix2 else astIndex shn u ix1+      bExpr ->+        -- This results in a larger term, so we consider this late.+        if knobPhase knobs == PhaseExpansion+        then astCond bExpr (astIndex shn v ix2) (astIndex shn u ix1)+        else Ast.AstIndexS shn v0 ix+  Ast.AstSliceS i@(SNat @i) (SNat @n) k@SNat v ->+    astLetFunB i1 $ \iShared ->+    let ftk = FTKS shn x+        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+        b = (if sNatValue i == 0 then true else 0 <=. iShared )+            &&* (if sNatValue k == 0 then true else iShared <=. valueOf @n - 1)+        ii = valueOf @i + iShared+    in astCond b (astIndex shn v (ii :.$ rest1)) defArr+  Ast.AstReverseS v ->+    let iRev = valueOf @in1 - 1 - i1+    in astIndex shn v (iRev :.$ rest1)+  Ast.AstTransposeS @_ @sh2 perm v+    | gcompare (shsRank (shsFromIxS ix)) (Permutation.permRank perm) /= GLT ->+      -- TODO: remake once there's an S version of permInverse:+      permInverse perm $ \(permR :: Permutation.Perm permR) _ ->+      let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix permR shm)+          ix2 = ixsPermutePrefix permR ix+      in gcastWith (unsafeCoerceRefl+                    :: sh2 :~: Permutation.PermutePrefix permR shm ++ shn) $+         astIndex @(Permutation.PermutePrefix permR shm) shn v ix2+  Ast.AstTransposeS @perm perm v+    | knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion] ->+      astIndex shn (astTransposeAsGatherS @perm (deVect knobs) perm v) ix+  Ast.AstTransposeS{} -> Ast.AstIndexS shn v0 ix+  -- This results in a larger term, so we consider this late.+  Ast.AstReshapeS sh v | knobPhase knobs == PhaseExpansion+                         || shsLength sh <= 1 ->+    astIndex shn (astReshapeAsGatherS (deVect knobs) sh v) ix+  Ast.AstReshapeS{} -> Ast.AstIndexS shn v0 ix++  AstFromS' ftkz v -> case matchingFTK (ftkAst v) ftkz of+    Just Refl -> astIndexKnobsS knobs shn v ix+      -- rare, usually simplifies away earlier+    Nothing -> error "astIndexKnobsS: wrong tensor kinds in AstFromS"+  -- These conversions need to stay down, so this is NF, see vectorization.+  Ast.AstConvert{} -> Ast.AstIndexS shn v0 ix++  -- These should not appear here unless via wacky tests.+  Ast.AstDot1InS{} -> Ast.AstIndexS shn v0 ix+  Ast.AstMatmul2S{} -> Ast.AstIndexS shn v0 ix++-- TODO: compared to tletIx, it adds many lets, not one, but does not+-- create other (and non-simplified!) big terms and also uses astIsSmall,+-- so it's probably more efficient. Use this instead of tletIx+-- or design something even better.+--+-- See https://github.com/Mikolaj/horde-ad/issues/119.+shareIx :: AstSpan s+        => AstIxS AstMethodLet shm+        -> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)+        -> AstTensor AstMethodLet s y+{-# NOINLINE shareIx #-}+shareIx ix f = unsafePerformIO $ do+  let shareI :: AstInt AstMethodLet+             -> IO ( Maybe (IntVarName, AstInt AstMethodLet)+                   , AstInt AstMethodLet )+      shareI i | astIsSmall True i = return (Nothing, i)+      shareI i =+        -- i can be OOB, so we can't use shape to determine its bounds+        let bds = bounds i+        in funToAstIntVarIO (Just bds) $ \ (!varFresh, !astVarFresh) ->+                                           (Just (varFresh, i), astVarFresh)+  (bindings, ix2) <- mapAndUnzipM shareI (Foldable.toList ix)+  return $! foldr (uncurry astLet)+                  (withKnownShS (shsFromIxS ix) $ f $ fromList ix2)+                  (catMaybes bindings)++-- TODO: fuse scatters, scatter and sum, and perhaps more (fromList?)+astScatterS :: forall shm shn shp r s. AstSpan s+            => ShS shn+            -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)+            -> (AstVarListS shm, AstIxS AstMethodLet shp)+            -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)+astScatterS _shn v (ZS, ZIS) = v+astScatterS shn v0 (_vars, ix@((:.$) @k i1 _))+  | let (lb, ub) = bounds i1+        FTKS _ x = ftkAst v0+  , ub < 0 || lb >= valueOf @k =+    let ftk = FTKS (shsFromIxS ix `shsAppend` shn) x+    in fromPrimal $ astConcrete ftk (tdefTarget ftk)+astScatterS shn v (vars, (:.$) @k (AstConcreteK _) rest)+  | Just Refl <- sameNat (SNat @k) (SNat @1)+  , FTKS _ x <- ftkAst v =+    astReplicate (SNat @1) (STKS (shsFromIxS rest `shsAppend` shn) (ftkToSTK x))+    $ astScatterS shn v (vars, rest)+astScatterS shn v (Const var ::$ (vars :: AstVarListS sh3), ix)+  | not $ var `varNameInIxS` ix+  , FTKS _ x <- ftkAst v =+      astScatterS @sh3 @shn @shp shn+        (astSum SNat (STKS (shsFromListS vars `shsAppend` shn) (ftkToSTK x)) v)+        (vars, ix)+-- TODO? astScatterS v (ZR, ix) = update (rzero sh 0) ix v+astScatterS shn (Ast.AstLet var u v) (vars, ix) =+      astLet var u (astScatterS @shm @shn @shp shn v (vars, ix))+astScatterS shn (Ast.AstFromPrimal v) (vars, ix) =+  Ast.AstFromPrimal $ astScatterS @shm @shn @shp shn v (vars, ix)+astScatterS shn (Ast.AstFromDual v) (vars, ix) =+  Ast.AstFromDual $ astScatterS @shm @shn @shp shn v (vars, ix)+astScatterS shn v (vars, ix) = Ast.AstScatterS @shm @shn @shp shn v (vars, ix)++flipCompare :: forall (a :: Nat) b. Compare a b ~ GT => Compare b a :~: LT+flipCompare = unsafeCoerceRefl++astGatherS+  :: forall shm shn shp r s. AstSpan s+  => ShS shn+  -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)+  -> (AstVarListS shm, AstIxS AstMethodLet shp)+  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)+astGatherS = astGatherKnobsS @shm @shn @shp defaultKnobs++-- Assumption: vars0 don't not occur in v0. The assumption only holds+-- when newly generated variables are fresh, which is the case as long+-- as resetVarCounter is not used. The assumption makes it easier to spot+-- bugs or corruption, hence we assert it in the code below.+astGatherKnobsS+  :: forall shm shn shp r s. AstSpan s+  => SimplifyKnobs+  -> ShS shn+  -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)+  -> (AstVarListS shm, AstIxS AstMethodLet shp)+  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)+astGatherKnobsS _ _ v0 (!vars0, !_ix0)+  | any (`varNameInAst` v0) $ listsToList vars0 =+    error $ "astGatherKnobsS: gather vars in v0: " ++ show (vars0, v0)+astGatherKnobsS knobs shn v0 (ZS, ix0) = astIndexKnobsS knobs shn v0 ix0+astGatherKnobsS _ _ v0 (vars0, ZIS) =+  astReplicateNS @shm @shn (shsFromListS vars0) v0+astGatherKnobsS _ shn v0 (vars, i1 :.$ _)+  | let (lb, ub) = bounds i1+-- this doesn't work in GHC 9.10:+--      FTKS (snat :$$ _) x = ftkAst v0+  , FTKS (snat :$$ _) x <- ftkAst v0+  , ub < 0 || lb >= fromInteger (fromSNat snat) =+    let ftk = FTKS (shsFromListS vars `shsAppend` shn) x+    in fromPrimal $ astConcrete ftk (tdefTarget ftk)+astGatherKnobsS knobs shn v0 (vars0@(var1 ::$ vars1), ix0)+  | not (getConst var1 `varNameInIxS` ix0) =+    let k :$$ sh' = shsFromListS vars0+        FTKS _ x = ftkAst v0+    in astReplicate k (STKS (sh' `shsAppend` shn) (ftkToSTK x))+                    (astGatherKnobsS @(Tail shm) @shn @shp+                                     knobs shn v0 (vars1, ix0))+astGatherKnobsS knobs shn v0 (vars0@(_ ::$ _), ix0@(_ :.$ _))+  | let ixInit = ixsInit ix0+        varInit = listsInit vars0+        varLast = getConst $ listsLast vars0+  , AstIntVar ixvarLast <- ixsLast ix0+  , ixvarLast == varLast+  , not (varLast `varNameInIxS` ixInit)+  , kLast@SNat <- shsLast (shsFromListS vars0)+  , Just Refl <- testEquality kLast (shsLast (shsFromIxS ix0)) =+    gcastWith (unsafeCoerceRefl+               :: Init shp ++ (Last shm ': shn) :~: shp ++ shn) $+    gcastWith (unsafeCoerceRefl+               :: Init shm ++ (Last shm ': shn) :~: shm ++ shn) $+    astGatherKnobsS @(Init shm) @(Last shm ': shn) @(Init shp)+                    knobs (kLast :$$ shn) v0 (varInit, ixInit)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstCond+           (Ast.AstBoolAnd a@(AstLeqInt (AstConcreteK j) AstIntVar{}) b)+           v w :.$ prest )+    | j <= 0 || j >= valueOf @m || ixIsSmall prest =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstLet varN uN+           (Ast.AstCond+              (Ast.AstBoolAnd a@(AstLeqInt (AstConcreteK j) AstIntVar{}) b)+              v w) :.$ prest )+    | j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstCond+           (Ast.AstBoolAnd a@(AstLeqInt (AstConcreteK j)+                                        (Ast.AstN1K NegateOp AstIntVar{})) b)+           v w :.$ prest )+    | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstLet varN uN+           (Ast.AstCond+              (Ast.AstBoolAnd a@(AstLeqInt (AstConcreteK j)+                                           (Ast.AstN1K NegateOp AstIntVar{})) b)+              v w) :.$ prest )+    | - j + 1 <= 0 || - j + 1 >= valueOf @m+      || ixIsSmall prest && astIsSmall True uN =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)++astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstCond+           (Ast.AstBoolAnd+              a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK j) AstIntVar{})) b)+           v w :.$ prest )+    | j <= 0 || j >= valueOf @m || ixIsSmall prest =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstLet varN uN+           (Ast.AstCond+              (Ast.AstBoolAnd+                 a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK j) AstIntVar{})) b)+              v w) :.$ prest )+    | j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstCond+           (Ast.AstBoolAnd+              a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK j)+                                           (Ast.AstN1K NegateOp+                                                       AstIntVar{}))) b)+           v w :.$ prest )+    | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstLet varN uN+           (Ast.AstCond+              (Ast.AstBoolAnd+                 a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK j)+                                              (Ast.AstN1K NegateOp+                                                          AstIntVar{}))) b)+              v w) :.$ prest )+    | - j + 1 <= 0 || - j + 1 >= valueOf @m+      || ixIsSmall prest && astIsSmall True uN =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)++astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstCond+      (Ast.AstBoolAnd+         a@(Ast.AstBoolNot+              (Ast.AstBoolAnd (AstLeqInt (AstConcreteK j) AstIntVar{}) _)) b)+      v w :.$ prest )+    | j <= 0 || j >= valueOf @m || ixIsSmall prest =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstLet varN uN+           (Ast.AstCond+              (Ast.AstBoolAnd+                 a@(Ast.AstBoolNot+                      (Ast.AstBoolAnd (AstLeqInt (AstConcreteK j)+                                                 AstIntVar{}) _)) b)+              v w) :.$ prest )+    | j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstCond+           (Ast.AstBoolAnd+              a@(Ast.AstBoolNot+                   (Ast.AstBoolAnd (AstLeqInt (AstConcreteK j)+                                              (Ast.AstN1K NegateOp+                                                          AstIntVar{})) _)) b)+           v w :.$ prest )+    | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstLet varN uN+           (Ast.AstCond+              (Ast.AstBoolAnd+                 a@(Ast.AstBoolNot+                      (Ast.AstBoolAnd+                         (AstLeqInt (AstConcreteK j)+                                    (Ast.AstN1K NegateOp+                                                AstIntVar{})) _)) b)+              v w) :.$ prest )+    | - j + 1 <= 0 || - j + 1 >= valueOf @m+      || ixIsSmall prest && astIsSmall True uN =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)++astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstCond+           (Ast.AstBoolAnd+              a@(Ast.AstBoolNot+                   (Ast.AstBoolAnd+                      (Ast.AstBoolNot+                         (AstLeqInt (AstConcreteK j) AstIntVar{})) _)) b)+           v w :.$ prest )+    | j <= 0 || j >= valueOf @m || ixIsSmall prest =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstLet varN uN+           (Ast.AstCond+              (Ast.AstBoolAnd+                 a@(Ast.AstBoolNot+                      (Ast.AstBoolAnd+                         (Ast.AstBoolNot (AstLeqInt (AstConcreteK j)+                                                    AstIntVar{})) _)) b)+              v w) :.$ prest )+    | j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstCond+           (Ast.AstBoolAnd+              a@(Ast.AstBoolNot+                   (Ast.AstBoolAnd+                      (Ast.AstBoolNot+                         (AstLeqInt (AstConcreteK j)+                                   (Ast.AstN1K NegateOp AstIntVar{}))) _)) b)+           v w :.$ prest )+    | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, i :.$ prest)+astGatherKnobsS+  knobs shn v0+  ( vars@((::$) @m _ _)+  , Ast.AstLet varN uN+           (Ast.AstCond+              (Ast.AstBoolAnd+                 a@(Ast.AstBoolNot+                      (Ast.AstBoolAnd+                         (Ast.AstBoolNot+                            (AstLeqInt (AstConcreteK j)+                                       (Ast.AstN1K NegateOp+                                                   AstIntVar{}))) _)) b)+              v w) :.$ prest )+    | - j + 1 <= 0 || - j + 1 >= valueOf @m+      || ixIsSmall prest && astIsSmall True uN =+  let i = astLetFunB w $ \wShared -> astCond a (astCond b v wShared) wShared+  in astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i :.$ prest)+-- Rules with AstConcreteK on the right hand side of AstPlusK are+-- not needed, thanks to the normal form of AstPlusK rewriting.+astGatherKnobsS knobs shn v0+  (vars, AstPlusK (AstConcreteK i64) i1 :.$ prest)+  | let (lb, ub) = bounds i1+  , lb >= 0  -- if not, we may need to apply astReverse first+  , FTKS (SNat @p :$$ _) x <- ftkAst v0 =+    if i64 >= 0 then+      withSNat (fromIntegral i64) $ \(SNat @i) ->+      withSNat (fromIntegral $ min (valueOf @p - i64) (ub + 1)) $ \(SNat @k) ->+        gcastWith (unsafeCoerceRefl :: (i + k <=? p) :~: True) $+        let v2 = astSliceS (SNat @i) (SNat @k) (SNat @(p - (i + k))) v0+        in astGatherKnobsS knobs shn v2 (vars, i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+    else+      withSNat (- fromIntegral i64) $ \(SNat @i) ->+        let ftk = FTKS (SNat @i :$$ shsFromIxS prest `shsAppend` shn) x+            v2 = fromPrimal (astConcrete ftk (tdefTarget ftk))+                 `astAppendS`+                 v0+        in astGatherKnobsS knobs shn v2 (vars, i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+astGatherKnobsS knobs shn v0+  (vars, Ast.AstLet varN uN (AstPlusK (AstConcreteK i64) i1) :.$ prest)+  | let (lb, ub) = bounds i1+  , lb >= 0  -- if not, we may need to apply astReverse first+  , FTKS (SNat @p :$$ _) x <- ftkAst v0 =+    if i64 >= 0 then+      withSNat (fromIntegral i64) $ \(SNat @i) ->+      withSNat (fromIntegral $ min (valueOf @p - i64) (ub + 1)) $ \(SNat @k) ->+        gcastWith (unsafeCoerceRefl :: (i + k <=? p) :~: True) $+        let v2 = astSliceS (SNat @i) (SNat @k) (SNat @(p - (i + k))) v0+        in astGatherKnobsS knobs shn v2 (vars, Ast.AstLet varN uN i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+    else+      withSNat (- fromIntegral i64) $ \(SNat @i) ->+        let ftk = FTKS (SNat @i :$$ shsFromIxS prest `shsAppend` shn) x+            v2 = fromPrimal (astConcrete ftk (tdefTarget ftk))+                 `astAppendS`+                 v0+        in astGatherKnobsS knobs shn v2 (vars, Ast.AstLet varN uN i1 :.$ prest)+             -- this gather may still index out of bounds, which is fine+astGatherKnobsS knobs shn v0+  ( vars@((::$) @m (Const varm) mrest)+  , Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar varp)) i1 i2+    :.$ prest )+  | varNameToAstVarId varm == varNameToAstVarId varp+  , j <= 0 || j >= valueOf @m || ixIsSmall prest =+    if | j <= 0 ->+         astGatherKnobsS knobs shn v0 (vars, i1 :.$ prest)+       | j >= valueOf @m ->+         astGatherKnobsS knobs shn v0 (vars, i2 :.$ prest)+       | otherwise ->+         withSNat (fromIntegral j) $ \(SNat @j) ->+         gcastWith (unsafeCoerceRefl :: (j <=? m) :~: True) $+         astLetFun v0 $ \v ->+         let varm2 = mkAstVarName (varNameToFTK varm)+                                  (Just (0, j - 1))+                                  (varNameToAstVarId varm)+             varm3 = mkAstVarName (varNameToFTK varm)+                                  (Just (0, valueOf @m - j - 1))+                                  (varNameToAstVarId varm)+         in astGatherKnobsS knobs shn v+              ( (::$) @j (Const varm2) mrest+              , substituteAstIxS (astVar varm2)+                                 varm (i2 :.$ prest) )+            `astAppendS`+            astGatherKnobsS knobs shn v+              ( (::$) @(m - j) (Const varm3) mrest+              , substituteAstIxS (AstConcreteK j + astVar varm3)+                                 varm (i1 :.$ prest) )+astGatherKnobsS knobs shn v0+  ( vars@((::$) @m (Const varm) mrest)+  , Ast.AstCond (AstLeqInt (AstConcreteK j)+                           (Ast.AstN1K NegateOp (AstIntVar varp))) i1 i2+    :.$ prest )+  | varNameToAstVarId varm == varNameToAstVarId varp+  , - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest =+    if | - j + 1 <= 0 ->+         astGatherKnobsS knobs shn v0 (vars, i2 :.$ prest)+       | - j + 1 >= valueOf @m ->+         astGatherKnobsS knobs shn v0 (vars, i1 :.$ prest)+       | otherwise ->+         withSNat (- fromIntegral j + 1) $ \(SNat @mj) ->+         gcastWith (unsafeCoerceRefl :: (mj <=? m) :~: True) $+         astLetFun v0 $ \v ->+         let varm2 = mkAstVarName (varNameToFTK varm)+                                  (Just (0, valueOf @mj - 1))+                                  (varNameToAstVarId varm)+             varm3 = mkAstVarName (varNameToFTK varm)+                                  (Just (0, valueOf @m - valueOf @mj - 1))+                                  (varNameToAstVarId varm)+         in astGatherKnobsS knobs shn v+              ( (::$) @mj (Const varm2) mrest+              , substituteAstIxS (astVar varm2)+-- TODO: when I use AstIntVar here, which is wrong, I get phantom errors;+-- make sure this vanished after the fixes in HEAD+                                 varm (i1 :.$ prest) )+            `astAppendS`+            astGatherKnobsS knobs shn v+              ( (::$) @(m - mj) (Const varm3) mrest+              , substituteAstIxS (AstConcreteK (- j + 1) + astVar varm3)+                                 varm (i2 :.$ prest))+astGatherKnobsS knobs shn v0+  ( vars@((::$) @m (Const varm) mrest)+  , Ast.AstLet varN uN+      (Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar varp)) i1 i2)+      :.$ prest )+  | varNameToAstVarId varm == varNameToAstVarId varp+  , j <= 0 || j >= valueOf @m || ixIsSmall prest && astIsSmall True uN =+    if | j <= 0 ->+         astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i1 :.$ prest)+       | j >= valueOf @m ->+         astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i2 :.$ prest)+       | otherwise ->+         withSNat (fromIntegral j) $ \(SNat @j) ->+         gcastWith (unsafeCoerceRefl :: (j <=? m) :~: True) $+         astLetFun v0 $ \v ->+         let varm2 = mkAstVarName (varNameToFTK varm)+                                  (Just (0, j - 1))+                                  (varNameToAstVarId varm)+             varm3 = mkAstVarName (varNameToFTK varm)+                                  (Just (0, valueOf @m - j - 1))+                                  (varNameToAstVarId varm)+         in astGatherKnobsS knobs shn v+              ( (::$) @j (Const varm2) mrest+              , substituteAstIxS (astVar varm2) varm+                                 (Ast.AstLet varN uN i2 :.$ prest) )+            `astAppendS`+            astGatherKnobsS knobs shn v+              ( (::$) @(m - j) (Const varm3) mrest+              , substituteAstIxS (AstConcreteK j + astVar varm3)+                                 varm (Ast.AstLet varN uN i1 :.$ prest) )+astGatherKnobsS knobs shn v0+  ( vars@((::$) @m (Const varm) mrest)+  , Ast.AstLet varN uN+      (Ast.AstCond (AstLeqInt (AstConcreteK j)+                              (Ast.AstN1K NegateOp (AstIntVar varp))) i1 i2)+    :.$ prest )+  | varNameToAstVarId varm == varNameToAstVarId varp+  , - j + 1 <= 0 || - j + 1 >= valueOf @m+    || ixIsSmall prest && astIsSmall True uN =+    if | - j + 1 <= 0 ->+         astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i2 :.$ prest)+       | - j + 1 >= valueOf @m ->+         astGatherKnobsS knobs shn v0 (vars, Ast.AstLet varN uN i1 :.$ prest)+       | otherwise ->+         withSNat (- fromIntegral j + 1) $ \(SNat @mj) ->+         gcastWith (unsafeCoerceRefl :: (mj <=? m) :~: True) $+         astLetFun v0 $ \v ->+         let varm2 = mkAstVarName (varNameToFTK varm)+                                  (Just (0, valueOf @mj - 1))+                                  (varNameToAstVarId varm)+             varm3 = mkAstVarName (varNameToFTK varm)+                                  (Just (0, valueOf @m - valueOf @mj - 1))+                                  (varNameToAstVarId varm)+         in astGatherKnobsS knobs shn v+              ( (::$) @mj (Const varm2) mrest+              , substituteAstIxS (astVar varm2)+                                 varm (Ast.AstLet varN uN i1 :.$ prest) )+            `astAppendS`+            astGatherKnobsS knobs shn v+              ( (::$) @(m - mj) (Const varm3) mrest+              , substituteAstIxS (AstConcreteK (- j + 1) + astVar varm3)+                                 varm (Ast.AstLet varN uN i2 :.$ prest))+astGatherKnobsS knobs shn v0+  ( (::$) @m @shmTail (Const varm) mrest+  , (:.$) @p @shpTail (AstIntVar varp) prest )+  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]+      -- prevent a loop+  , varm == varp+  , not (varm `varNameInIxS` prest)+  , FTKS _ x <- ftkAst v0 =+    withSNat (min (valueOf @p) (valueOf @m)) $ \(SNat @m2) ->+    gcastWith (unsafeCoerceRefl :: (m2 <=? p) :~: True) $+    gcastWith (unsafeCoerceRefl :: (m2 <=? m) :~: True) $+    Permutation.permFromList (permCycle $ shsLength (shsFromListS mrest) + 1)+    $ \(permVars :: Permutation.Perm permVars) ->+    Permutation.permFromList (backpermCycle $ shsLength (shsFromIxS prest) + 1)+    $ \(permIx :: Permutation.Perm permIx) ->+       gcastWith (unsafeCoerceRefl+                  :: m2 ': shmTail ++ shn+                     :~: Permutation.PermutePrefix+                           permVars (shmTail ++ (m2 ': shn))) $+       gcastWith (unsafeCoerceRefl+                  :: shpTail ++ (m2 ': shn)+                     :~: Permutation.PermutePrefix+                           permIx (m2 ': shpTail ++ shn)) $+       gcastWith (unsafeCoerceRefl+                  :: (Rank permVars <=? Rank (shmTail ++ (m2 ': shn)))+                     :~: True) $+       gcastWith (unsafeCoerceRefl+                  :: (Rank permIx <=? Rank (m2 ': shpTail ++ shn)) :~: True) $+       fromMaybe (error "astGatherKnobsS: impossible non-permutation")+       $ Permutation.permCheckPermutation permVars+       $ fromMaybe (error "astGatherKnobsS: impossible non-permutation")+       $ Permutation.permCheckPermutation permIx+       $ let v2 = astTransposeS permIx+                  $ astSliceS (SNat @0) (SNat @m2) (SNat @(p - m2)) v0+             u = astGatherKnobsS knobs (SNat @m2 :$$ shn) v2 (mrest, prest)+             ftk =+               FTKS (SNat @(m - m2) :$$ shsFromListS mrest `shsAppend` shn) x+         in astTransposeS permVars u+            `astAppendS`+            fromPrimal (astConcrete ftk (tdefTarget ftk))+astGatherKnobsS knobs shn v7@(Ast.AstFromVector _ (STKS _ x2) l)+                ( ((::$) @m1' @shm4 (Const var4) vrest4)+                , ((:.$) @_ @shp1' i4 rest4) )+  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]+  , let g = case i4 of+          AstIntVar var | var == var4 -> Just id+          AstTimesK (AstConcreteK n) (AstIntVar var)+            | var == var4 -> Just (n *)+          -- TODO: add more or define evaluation+          _ -> Nothing+  , Just h <- g+  , ixIsSmall rest4 =+    let f i =+          let subRest4 = substituteAstIxS (AstConcreteK i) var4 rest4+              j = fromIntegral $ h i+          in if j >= V.length l+             then let FTKS _ x = ftkAst v7+                      ftk = FTKS (shsFromListS vrest4 `shsAppend` shn) x+                  in fromPrimal $ astConcrete ftk (tdefTarget ftk)+             else astGatherKnobsS @shm4 @shn @shp1' knobs shn+                                  (l V.! j) (vrest4, subRest4)+    in astFromVector (SNat @m1')+                     (STKS (shsFromListS vrest4 `shsAppend` shn) x2)+       $ V.fromList $ map f [0 .. valueOf @m1' - 1]+astGatherKnobsS knobs shn v0 (vars0, i1 :.$ rest1)+  | knobPhase knobs `notElem` [PhaseVectorization, PhaseExpansion]+      -- prevent a loop+  , not (any (`varNameInAst` i1) $ listsToList vars0) =+    astGatherKnobsS @shm @shn+      knobs shn+      (astIndexKnobsS knobs (shsFromIxS rest1 `shsAppend` shn) v0 (i1 :.$ ZIS))+      (vars0, rest1)+astGatherKnobsS knobs shn v0+  (vars@((::$) @m _ _), ix@(i1 :.$ prest))+  | let intInteresting = \case+          AstPlusK (AstConcreteK _) i2+            | fst (bounds i2) >= 0 -> True+          Ast.AstLet _ _ (AstPlusK (AstConcreteK _) i2)+            | fst (bounds i2) >= 0 -> True+          Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar var)) _ _+            | j <= 0 || j >= valueOf @m || ixIsSmall prest+            , any ((== varNameToAstVarId var) . varNameToAstVarId)+                  (listsToList vars) -> True+          Ast.AstCond (AstLeqInt (AstConcreteK j)+                                 (Ast.AstN1K NegateOp (AstIntVar var))) _ _+            | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest+            , any ((== varNameToAstVarId var) . varNameToAstVarId)+                  (listsToList vars) -> True+          Ast.AstLet _ uN+            (Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar var)) _ _)+            | j <= 0 || j >= valueOf @m+              || ixIsSmall prest && astIsSmall True uN+            , any ((== varNameToAstVarId var) . varNameToAstVarId)+                  (listsToList vars) -> True+          Ast.AstLet _ uN+            (Ast.AstCond (AstLeqInt (AstConcreteK j)+                                    (Ast.AstN1K NegateOp (AstIntVar var))) _ _)+            | - j + 1 <= 0 || - j + 1 >= valueOf @m+              || ixIsSmall prest && astIsSmall True uN+            , any ((== varNameToAstVarId var) . varNameToAstVarId)+                  (listsToList vars) -> True+          Ast.AstCond+            (Ast.AstBoolAnd+               (AstLeqInt (AstConcreteK j) (AstIntVar var)) _) _ _+            | j <= 0 || j >= valueOf @m || ixIsSmall prest+            , any ((== varNameToAstVarId var) . varNameToAstVarId)+                  (listsToList vars) -> True+          Ast.AstCond+            (Ast.AstBoolAnd+               (AstLeqInt (AstConcreteK j)+                          (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _+            | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest+            , any ((== varNameToAstVarId var) . varNameToAstVarId)+                  (listsToList vars) -> True+          Ast.AstLet _ uN+            (Ast.AstCond+               (Ast.AstBoolAnd+                  (AstLeqInt (AstConcreteK j) (AstIntVar var)) _) _ _)+            | j <= 0 || j >= valueOf @m+              || ixIsSmall prest && astIsSmall True uN+            , any ((== varNameToAstVarId var) . varNameToAstVarId)+                  (listsToList vars) -> True+          Ast.AstLet _ uN+            (Ast.AstCond+               (Ast.AstBoolAnd+                  (AstLeqInt (AstConcreteK j)+                             (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _)+            | - j + 1 <= 0 || - j + 1 >= valueOf @m+              || ixIsSmall prest && astIsSmall True uN+            , any ((== varNameToAstVarId var) . varNameToAstVarId)+                  (listsToList vars) -> True+          AstIntVar var+            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]+            , null $ drop 1 $ filter (var `varNameInAst`) (Foldable.toList ix)+            , any ((== varNameToAstVarId var) . varNameToAstVarId)+                  (listsToList vars) -> True+          ik | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]+             , not (any (`varNameInAst` ik) $ listsToList vars) -> True+          -- We can't reorder ix for the gather(fromVector) rule above,+          -- because it becomes gather(transpose); we can only reorder vars.+          _ -> False+  , not (intInteresting i1)  -- now vars may need to be reordered, too+  , Just i <- findIndex intInteresting+                        (Foldable.toList ix) = assert (i > 0) $+    Permutation.permFromList (backpermCycle $ i + 1)+    $ \(perm :: Permutation.Perm perm) ->+    gcastWith (unsafeCoerceRefl+               :: Permutation.PermutePrefix perm (shp ++ shn)+                  :~: Permutation.PermutePrefix perm shp ++ shn) $+    gcastWith (unsafeCoerceRefl :: (Rank perm <=? Rank (shp ++ shn)) :~: True) $+    fromMaybe (error "astGatherKnobsS: impossible non-permutation")+    $ Permutation.permCheckPermutation perm+    $ let v2 = astTransposeS perm v0+      in astGatherKnobsS knobs shn v2 (vars, ixsPermutePrefix perm ix)+        -- this call is guaranteed to simplify as above, so the tranpose+        -- won't reduce it back to the original and cause a loop+astGatherKnobsS knobs shn v0+  (vars@((::$) @m _ _), ix@(i1 :.$ prest))+  | let varInteresting = \case+          Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar var)) _ _+            | j <= 0 || j >= valueOf @m || ixIsSmall prest ->+              Just var+          Ast.AstCond (AstLeqInt (AstConcreteK j)+                                 (Ast.AstN1K NegateOp (AstIntVar var))) _ _+            | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest ->+              Just var+          Ast.AstLet _ uN+            (Ast.AstCond (AstLeqInt (AstConcreteK j) (AstIntVar var)) _ _)+            | j <= 0 || j >= valueOf @m+              || ixIsSmall prest && astIsSmall True uN ->+              Just var+          Ast.AstLet _ uN+            (Ast.AstCond (AstLeqInt (AstConcreteK j)+                                    (Ast.AstN1K NegateOp+                                                (AstIntVar var))) _ _)+            | - j + 1 <= 0 || - j + 1 >= valueOf @m+              || ixIsSmall prest && astIsSmall True uN ->+              Just var+          Ast.AstCond+            (Ast.AstBoolAnd+               (AstLeqInt (AstConcreteK j) (AstIntVar var)) _) _ _+            | j <= 0 || j >= valueOf @m || ixIsSmall prest ->+              Just var+          Ast.AstCond+            (Ast.AstBoolAnd+               (AstLeqInt (AstConcreteK j)+                          (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _+            | - j + 1 <= 0 || - j + 1 >= valueOf @m || ixIsSmall prest ->+              Just var+          Ast.AstLet _ uN+            (Ast.AstCond+               (Ast.AstBoolAnd+                  (AstLeqInt (AstConcreteK j) (AstIntVar var)) _) _ _)+            | j <= 0 || j >= valueOf @m+              || ixIsSmall prest && astIsSmall True uN ->+              Just var+          Ast.AstLet _ uN+            (Ast.AstCond+               (Ast.AstBoolAnd+                  (AstLeqInt (AstConcreteK j)+                             (Ast.AstN1K NegateOp (AstIntVar var))) _) _ _)+            | - j + 1 <= 0 || - j + 1 >= valueOf @m+              || ixIsSmall prest && astIsSmall True uN ->+              Just var+          AstIntVar var+            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]+            , not (var `varNameInIxS` prest) -> Just var+          i4  -- has to be last, because ix can't be reordered+            | knobPhase knobs `elem` [PhaseSimplification, PhaseContraction]+            , Ast.AstFromVector{} <- v0+            , ixIsSmall prest+            , mvar <- case i4 of+                AstIntVar var -> Just var+                AstTimesK AstConcreteK{} (AstIntVar var) -> Just var+                _ -> Nothing+            , Just{} <- mvar -> mvar+          _ -> Nothing+  , Just varp <- varInteresting i1+  , Just i <- findIndex ((== varNameToAstVarId varp) . varNameToAstVarId)+                        (listsToList vars) = assert (i > 0) $+    Permutation.permFromList (backpermCycle $ i + 1)+    $ \(permWhole :: Permutation.Perm permWhole) ->+    permInverse permWhole $ \(invperm :: Nested.Perm invperm) _ ->+    gcastWith (unsafeCoerceRefl+               :: shm ++ shn+                  :~: Permutation.PermutePrefix permWhole+                        (Permutation.PermutePrefix invperm shm ++ shn)) $+    gcastWith (unsafeCoerceRefl+               :: (Rank permWhole+                   <=? Rank (Permutation.PermutePrefix invperm shm ++ shn))+                  :~: True) $+    fromMaybe (error "astGatherKnobsS: impossible non-permutation")+    $ Permutation.permCheckPermutation permWhole+    $ astTransposeS permWhole+    $ astGatherKnobsS knobs shn v0 (listsPermutePrefix invperm vars, ix)+        -- this call is guaranteed to simplify as above, so the tranpose+        -- won't reduce it back to the original and cause a loop+astGatherKnobsS knobs shn v4 (vars4, ix4@((:.$) @in1 @shp1' i4 rest4))+  | FTKS _ x <- ftkAst v4 = case v4 of+    Ast.AstProject1{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstProject2{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    {- Ast.AstFromVector{} | gatherFromNF (shsFromIxS rest4) vars4 ix4 ->+        -- normal form+      Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4) -}+    {- this rule seems counterproductive in many cases, so disabled until+       we can detect cases where it helps:+    Ast.AstFromVector snat STKS{} l ->+      -- Term rest4 is duplicated without sharing and we can't help it,+      -- because it needs to be in scope of vars4, so we can't use tlet.+      funToVarsIxS @shm (shsFromListS vars4) $ \ (!varsFresh, IxS !ixFresh) ->+        let f v = astGather @shm @shn @shp1' shn v (vars4, rest4)+            -- This subst doesn't currently break sharing because it's a rename.+            subst i =+              foldr (\(i2, var2) v2 -> substituteAst i2 var2 v2)+                    i+                    (listsToList $ zipSizedS ixFresh vars4)+            i5 = subst i4+       in astGather @shm @shn @(p1' ': shm)+                    shn (astFromVector snat (STKS (shsFromListS varsFresh+                                                    `shsAppend` shn)+                                                   (ftkToSTK x))+                          $ V.map f l)+                    (varsFresh, i5 :.$ IxS ixFresh) -}+    Ast.AstFromVector{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    -- This accomplishes fusion if v is a gather or anything+    -- that gather can fuse with, but at the cost of an extra transpose+    -- that doesn't fuse here unless astTransposeAsGatherS is used.+    -- Since the transpose is O(1), let's leave this as is.+    Ast.AstSum snat@(SNat @n1) STKS{} v ->+      let perm3 = backpermCycle $ shsLength (shsFromIxS ix4) + 1+          perm4 = permCycle $ shsLength (shsFromListS vars4) + 1+      in Permutation.permFromList perm3+         $ \(perm3S :: Permutation.Perm perm3P) ->+         gcastWith (unsafeCoerceRefl+                    :: Compare (Rank perm3P) (Rank (n1 : shp ++ shn))+                       :~: LT) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix perm3P (n1 : (shp ++ shn))+                       :~: shp ++ (n1 : shn)) $+         fromMaybe (error "astGatherKnobsS: impossible non-permutation")+         $ Permutation.permCheckPermutation perm3S+         $ Permutation.permFromList perm4+         $ \(perm4S :: Permutation.Perm perm4P) ->+         gcastWith (unsafeCoerceRefl+                    :: Compare (Rank perm4P) (Rank (shm ++ (n1 : shn)))+                       :~: LT) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix perm4P (shm ++ (n1 : shn))+                       :~: n1 : (shm ++ shn)) $+         fromMaybe (error "astGatherKnobsS: impossible non-permutation")+         $ Permutation.permCheckPermutation perm4S+         $ let innerGather =+                 astGather @shm @(n1 : shn) @shp+                           (snat :$$ shn) (astTransposeS perm3S v) (vars4, ix4)+           in astSum snat (STKS (shsFromListS vars4 `shsAppend` shn)+                                (ftkToSTK x))+              $ astTransposeS perm4S innerGather+                {- TODO: disabled until we can reliably fuse back to transpose+                if not (knobExpand knobs)+                then astTransposeS perm4S innerGather+                else astTransposeAsGatherS knobs perm4S innerGather -}+    Ast.AstReplicate _ STKS{} v ->+      let ftk = FTKS (shsFromListS vars4 `shsAppend` shn) x+          defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+      -- This boolean term may have free variables that act as universally+      -- quantified.+      in case 0 <=. i4 &&* i4 <=. valueOf @in1 - 1 of+        AstBoolConst b ->+          if b then astGather @shm @shn @shp1' shn v (vars4, rest4) else defArr+        _ -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstReplicate (SNat @k) STKScalar v | ZIS <- rest4 ->+      let ftk = FTKS (shsFromListS vars4 `shsAppend` shn) x+          defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+      in case 0 <=. i4 &&* i4 <=. valueOf @k - 1 of+        AstBoolConst b ->+          if b then astGather @shm @shn @shp1'+                              shn (astSFromK' v) (vars4, rest4) else defArr+        _ -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstApply{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstVar{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstCond b v w | ixIsSmall ix4 ->+      astCond b (astGather @shm @shn @shp shn v (vars4, ix4))+                (astGather @shm @shn @shp shn w (vars4, ix4))+    Ast.AstCond{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstBuild1{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    AstConcreteS{} -> case unRepl1 v4 of+      Just u ->+        let ftk = FTKS (shsFromListS vars4 `shsAppend` shn) x+            defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)+        in case 0 <=. i4 &&* i4 <=. valueOf @in1 - 1 of+          AstBoolConst b ->+            if b+            then astGather @shm @shn @shp1' shn u (vars4, rest4)+            else defArr+          _ -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+      _ -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+             -- free variables possible in the index, so can't compute the array++    Ast.AstLet var u v ->+      astLet var u (astGather @shm @shn @shp shn v (vars4, ix4))++    Ast.AstPrimalPart{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstDualPart{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstFromPrimal v ->+      Ast.AstFromPrimal $ astGather @shm @shn @shp shn v (vars4, ix4)+    Ast.AstFromDual v ->+      Ast.AstFromDual $ astGather @shm @shn @shp shn v (vars4, ix4)++    -- Going inside a binary ops usually makes a term more expensive+    -- to interpret and inverting that requires comparing two arguments,+    -- so it's not practical.+    AstPlusS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    AstTimesS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstN1S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstR1S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstR2S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstI2S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstFloorS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstFromIntegralS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstCastS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)++    {- is reverted in astGatherKnobsS immediatedly; only do in expansion phase?+    Ast.AstIndexS @shm2 _shn2 v2 (i2 :.$ ZIS) ->+        astGather @shm @shn @(shm2 ++ shp) shn v2 (vars4, i2 :.$ ix4) -}+    Ast.AstIndexS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstScatterS @shm7 @shn7 @shp7 shn7 v (vars, AstIntVar var5 :.$ ix2)+      | AstIntVar var6 <- i4, var6 == var5 ->+          astGather @shm @shn @shp1' shn+                    (astScatterS @shm7 @shn7 @(Tail shp7) shn7+                                 v (vars, ix2))+                    (vars4, rest4)+    Ast.AstScatterS{} ->  -- normal form+      Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstGatherS @shm2 @shn2 @shp2 shn2 v2 (vars2, ix2)+                   | SNat @rank4 <- ixsRank ix4+                   , SNat @rank2 <- listsRank vars2 ->+      let subst :: AstIxS AstMethodLet shm7 -> AstVarListS shm7+                -> AstInt AstMethodLet+                -> AstInt AstMethodLet+          subst (IxS ix) vars t0 =+            foldr (\ (v, i) -> substituteAst i v)+                  t0 (listsFold (\(Fun.Pair (Const v) (Const i)) -> [(v, i)])+                      $ listsZip vars ix)+          inBounds :: AstIxS AstMethodLet shm7 -> AstVarListS shm7 -> Bool+          inBounds (IxS ix) vars =+            let inb (v, i) =+                  let (lbi, ubi) = bounds i+                  in case varNameToBounds v of+                    Nothing -> True+                    Just (lbv, ubv) -> lbv <= lbi && ubi <= ubv+            in all inb (listsFold (\(Fun.Pair (Const v) (Const i)) -> [(v, i)])+                        $ listsZip vars ix)+          IxS list4 = ix4+          composedGather ::  -- rank4 <= rank2+            Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) r))+          composedGather =+            -- we have: shm2 ++ shn2 == shp ++ shn+            -- so from ranks:+            gcastWith (unsafeCoerceRefl :: TakeLen shp shm2 :~: shp) $+            -- and from congruence:+--            gcastWith (unsafeCoerceRefl+--                       :: DropLen shp shm2 ++ shn2 :~: shn) $+            -- from congruence:+            gcastWith (unsafeCoerceRefl+                       :: (shm ++ DropLen shp shm2) ++ shn2+                          :~: shm ++ shn) $+            let vars2p = listsTakeLen list4 vars2+                vars22 = listsDropLen list4 vars2+                ix22 = fmap (subst ix4 vars2p) ix2+                list422 = vars4 `listsAppend` vars22+            in if ixIsSmall ix4 && inBounds ix4 vars2p+               then Just $ astGather shn2 v2 (list422, ix22)+               else Nothing+          assimilatedGather ::  -- rank2 <= rank4+            Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) r))+          assimilatedGather =+            -- we have: shm2 ++ shn2 == shp ++ shn+            -- so from ranks:+            gcastWith (unsafeCoerceRefl :: TakeLen shm2 shp :~: shm2) $+            -- and from congruence:+--            gcastWith (unsafeCoerceRefl+--                       :: DropLen shm2 shp ++ shn :~: shn2) $+            -- from congruence:+            gcastWith (unsafeCoerceRefl+                       :: (shp2 ++ DropLen shm2 shp) ++ shn+                          :~: shp2 ++ shn2) $+            let ix42 = IxS $ listsTakeLen vars2 list4+                ix44 = IxS $ listsDropLen vars2 list4+                ix22 = fmap (subst ix42 vars2) ix2+                ix2244 = ix22 `ixsAppend` ix44+            in if ixIsSmall ix42 && inBounds ix42 vars2+               then Just $ astGather shn v2 (vars4, ix2244)+               else Nothing+      in fromMaybe (Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4))+         $ case cmpNat (Proxy @rank4) (Proxy @rank2) of+        LTI -> composedGather+        EQI -> assimilatedGather+        GTI -> gcastWith (flipCompare @rank4 @rank2) assimilatedGather+    Ast.AstMinIndexS @n @sh v -> case ftkAst v of+     FTKS nsh _ -> case shsLast nsh of+      nl@(SNat @nl) ->+        let shnl = shn `shsAppend` (nl :$$ ZSS)+        in gcastWith (unsafeCoerceRefl+                     :: shp ++ (shn ++ '[nl]) :~: n ': sh) $+           gcastWith (unsafeCoerceRefl+                      :: Head (shm ++ (shn ++ '[nl]))+                         ': Tail (shm ++ (shn ++ '[nl]))+                         :~: shm ++ (shn ++ '[nl])) $+           gcastWith (unsafeCoerceRefl+                      :: Init (shm ++ (shn ++ '[nl]))+                      :~: shm ++ shn) $+           Ast.AstMinIndexS @(Head (shm ++ (shn ++ '[nl])))+                            @(Tail (shm ++ (shn ++ '[nl])))+           $ astGather shnl v (vars4, ix4)+    Ast.AstMaxIndexS @n @sh v -> case ftkAst v of+     FTKS nsh _ -> case shsLast nsh of+      nl@(SNat @nl) ->+        let shnl = shn `shsAppend` (nl :$$ ZSS)+        in gcastWith (unsafeCoerceRefl+                     :: shp ++ (shn ++ '[nl]) :~: n ': sh) $+           gcastWith (unsafeCoerceRefl+                      :: Head (shm ++ (shn ++ '[nl]))+                         ': Tail (shm ++ (shn ++ '[nl]))+                         :~: shm ++ (shn ++ '[nl])) $+           gcastWith (unsafeCoerceRefl+                      :: Init (shm ++ (shn ++ '[nl]))+                      :~: shm ++ shn) $+           Ast.AstMaxIndexS @(Head (shm ++ (shn ++ '[nl])))+                            @(Tail (shm ++ (shn ++ '[nl])))+           $ astGather shnl v (vars4, ix4)+    Ast.AstIotaS{} ->  -- probably nothing can be simplified; a normal form+      Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstAppendS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+      -- fusing would result in gather([gather, gather]), so no gain+    Ast.AstSliceS{}-> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+      -- slicing is O(1) so no point fusing and complicating the expression;+      -- if it did not simplify further with slice, it wouldn't with gather+    Ast.AstReverseS{}-> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+      -- reversing is O(1)+    Ast.AstTransposeS @perm @sh perm v | FTKS sh _ <- ftkAst v ->+      let rankPerm = Permutation.permRank perm+      in case gcompare (ixsRank ix4) rankPerm of+        GLT ->  -- TODO: this does not provide any proof, so use cmpNat :(+          if knobPhase knobs `elem` [PhaseVectorization, PhaseExpansion]+          then astGather @shm @shn @shp+                         shn (astTransposeAsGatherS knobs perm v) (vars4, ix4)+          else Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+        _ ->+          gcastWith (lemRankMapJust $ shsTakeLen perm sh) $+          gcastWith (unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) $+          permInverse perm+          $ \(invperm :: Nested.Perm invperm) proof ->+            case proof (ssxFromShX $ shxFromShS $ shsTakeLen perm sh) of+              Refl ->+                -- from PermutePrefix and ranks:+                gcastWith+                  (unsafeCoerceRefl+                   :: Permutation.PermutePrefix invperm shp ++ shn+                      :~: Permutation.PermutePrefix invperm (shp ++ shn)) $+                -- from AstTransposeS:+--                gcastWith+--                  (unsafeCoerceRefl+--                   :: Permutation.PermutePrefix invperm (shp ++ shn)+--                      :~: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)) $+                -- from PermutePrefix:+--                gcastWith+--                  (unsafeCoerceRefl+--                   :: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)+--                      :~: {-1-} Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))+--                          ++ {-2-} DropLen invperm (Permutation.PermutePrefix perm sh)) $+                -- 1. from PermutePrefix:+--                gcastWith+--                  (unsafeCoerceRefl+--                   :: Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))+--                      :~: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))) $+                -- ranks+                gcastWith+                  (unsafeCoerceRefl+                   :: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))+                      :~: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))) $+                -- from permInverse but MapJust-unwrapped:+                gcastWith+                  (unsafeCoerceRefl+                   :: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))+                      :~: TakeLen perm sh) $+                -- end of 1.+                -- 2. from PermutePrefix+--                gcastWith+--                  (unsafeCoerceRefl+--                   :: DropLen invperm (Permutation.PermutePrefix perm sh)+--                      :~: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)) $+                -- ranks+                gcastWith+                  (unsafeCoerceRefl+                   :: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)+                      :~: DropLen perm sh) $+                -- end of 2.+                -- from TakeLen:+                gcastWith+                  (unsafeCoerceRefl+                   :: TakeLen perm sh ++ DropLen perm sh :~: sh) $+                let invix4 = ixsPermutePrefix invperm ix4+                in astGather shn v (vars4, invix4)+    Ast.AstReshapeS sh v ->+      if shsLength sh <= 1+      then astGather @shm @shn @shp shn+                     (astReshapeAsGatherS knobs sh v) (vars4, ix4)+      else Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)++    AstFromS' ftkz v -> case matchingFTK (ftkAst v) ftkz of+      Just Refl -> astGather @shm @shn @shp shn v (vars4, ix4)+        -- rare, usually simplifies away earlier+      Nothing -> error "astGatherKnobsS: wrong tensor kinds in AstFromS"+    -- These conversions need to stay down.+    Ast.AstConvert{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)++    -- These should not appear here unless via wacky tests.+    Ast.AstDot1InS{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+    Ast.AstMatmul2S{} -> Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4)+ where+  astGather+    :: forall shm' shn' shp' s' r'. AstSpan s'+    => ShS shn'+    -> AstTensor AstMethodLet s' (TKS2 (shp' ++ shn') r')+    -> (AstVarListS shm', AstIxS AstMethodLet shp')+    -> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') r')+  astGather shn' v2 (vars2, ix2) =+    astGatherKnobsS @shm' @shn' @shp' knobs shn' v2 (vars2, ix2)++-- Normal form of chains of appends has the append constructor on the right.+astAppendS :: AstSpan s+           => AstTensor AstMethodLet s (TKS2 (m ': sh) r)+           -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)+           -> AstTensor AstMethodLet s (TKS2 ((m + n) ': sh) r)+astAppendS u v | FTKS (SNat' @0 :$$ _) _ <- ftkAst u = v+astAppendS u v | FTKS (SNat' @0 :$$ _) _ <- ftkAst v = u+astAppendS (Ast.AstFromVector (SNat @k1) stk2@STKS{} l1)+           (Ast.AstFromVector (SNat @k2) STKS{} l2) =+  astFromVector (SNat @(k1 + k2)) stk2 $ l1 V.++ l2+astAppendS (Ast.AstFromVector (SNat @k1) stk2@STKScalar l1)+           (Ast.AstFromVector (SNat @k2) STKScalar l2) =+  astFromVector (SNat @(k1 + k2)) stk2 $ l1 V.++ l2+astAppendS (Ast.AstFromPrimal u) (Ast.AstFromPrimal v) =+  Ast.AstFromPrimal $ astAppendS u v+astAppendS (Ast.AstFromDual u) (Ast.AstFromDual v) =+  Ast.AstFromDual $ astAppendS u v+astAppendS (AstConcreteS u) (AstConcreteS v) =+  astConcreteS (tsappend (Concrete u) (Concrete v))+astAppendS (AstConcreteS u) (Ast.AstAppendS (AstConcreteS v) w) =+  astAppendS (astConcreteS (tsappend (Concrete u) (Concrete v))) w+astAppendS (Ast.AstAppendS v u) w = astAppendS v (astAppendS u w)+astAppendS u v = Ast.AstAppendS u v++astSliceS :: forall i n k sh s r. AstSpan s+          => SNat i -> SNat n -> SNat k+          -> AstTensor AstMethodLet s (TKS2 (i + n + k ': sh) r)+          -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)+astSliceS SNat SNat SNat (Ast.AstFromVector _ stk@STKS{} l) =+  astFromVector (SNat @n) stk $ V.take (valueOf @n) $ V.drop (valueOf @i) l+astSliceS SNat SNat SNat (Ast.AstFromVector _ stk@STKScalar l) =+  astFromVector (SNat @n) stk $ V.take (valueOf @n) $ V.drop (valueOf @i) l+astSliceS SNat SNat SNat (Ast.AstReplicate _ stk@STKS{} v) =+  astReplicate (SNat @n) stk v+astSliceS SNat SNat SNat (Ast.AstReplicate _ stk@STKScalar v) =+  astReplicate (SNat @n) stk v+astSliceS (SNat' @0) SNat (SNat' @0) v = v+astSliceS SNat (SNat' @1) SNat v | FTKS (_ :$$ sh) x <- ftkAst v =+  astReplicate (SNat @1) (STKS sh (ftkToSTK x))+               (astIndexS sh v (valueOf @i :.$ ZIS))+astSliceS SNat SNat SNat (Ast.AstGatherS shn v (Const var ::$ vars, ix)) =+  let varn = mkAstVarName (varNameToFTK var)+                          (Just (0, valueOf @n - 1))+                          (varNameToAstVarId var)+      ivar = valueOf @i + astVar varn+      ix2 = substituteAstIxS ivar var ix  -- cheap subst, because ivar is tiny+  in astGatherS shn v (Const varn ::$ vars, ix2)+astSliceS i n@(SNat @n0) _k (Ast.AstAppendS v1 v2)+  | FTKS (m1@(SNat @m1) :$$ _) _ <- ftkAst v1+  , FTKS (m2@(SNat @m2) :$$ _) _ <- ftkAst v2 =+    let i1 = sNatValue i `min` sNatValue m1+        n1 = sNatValue n `min` (sNatValue m1 - i1)+        k1 = sNatValue m1 - i1 - n1+        i2' = sNatValue i `max` sNatValue m1+        i2 = i2' - sNatValue m1+        n2 = sNatValue n - n1+        k2 = sNatValue m2 - i2 - n2+    in withSNat i1 $ \si1@(SNat @i1) ->+       withSNat n1 $ \sn1@(SNat @n1) ->+       withSNat k1 $ \sk1@(SNat @k1) ->+       withSNat i2 $ \si2@(SNat @i2) ->+       withSNat n2 $ \sn2@(SNat @n2) ->+       withSNat k2 $ \sk2@(SNat @k2) ->+         gcastWith (unsafeCoerceRefl :: n1 + n2 :~: n0) $+         gcastWith (unsafeCoerceRefl :: i1 + n1 + k1 :~: m1) $+         gcastWith (unsafeCoerceRefl :: i2 + n2 + k2 :~: m2) $+         astAppendS (astSliceS si1 sn1 sk1 v1) (astSliceS si2 sn2 sk2 v2)+astSliceS i n k (Ast.AstSliceS i2 _n2 k2 v) =+  astSliceS (snatPlus i i2) n (snatPlus k k2) v+astSliceS i n k (Ast.AstReverseS v) = astReverseS (astSliceS k n i v)+-- This enlarges the term and increases computation, but sometimes+-- it permits eliminating the AstFromVector node altogether, so we risk it+-- for cases that commonly emerge from conditionals.+astSliceS i n@SNat k (Ast.AstTransposeS+                        perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil)+                        (Ast.AstFromVector (SNat' @2) (STKS (_ :$$ sh) x) l)) =+    Ast.AstTransposeS perm+    $ astFromVector (SNat @2) (STKS (n :$$ sh) x) (V.map (astSliceS i n k) l)+-- TODO: generalize (maybe the above, too) using unReplN, but it's hard.+-- TODO: does it really work only for replicate-like things in-between?+astSliceS i n@SNat k+          (Ast.AstTransposeS perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil) t)+  | STKS (snat :$$ _ :$$ sh3) x <- ftkToSTK $ ftkAst t+  , Just u <- unRepl1 t =+    astTransposeS perm+    $ astReplicate snat (STKS (n :$$ sh3) x)+    $ astSliceS i n k u+astSliceS i n k v1 = case v1 of+  Ast.AstCond b a2 a3 ->+    astCond b (astSliceS i n k a2) (astSliceS i n k a3)+  Ast.AstLet var u v -> astLet var u (astSliceS i n k v)+  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astSliceS i n k v+  Ast.AstFromDual v -> Ast.AstFromDual $ astSliceS i n k v+  AstPlusS u v -> astSliceS i n k u + astSliceS i n k v+  AstTimesS u v -> astSliceS i n k u * astSliceS i n k v+  Ast.AstN1S NegateOp u -> negate (astSliceS i n k u)+  Ast.AstN1S AbsOp u -> abs (astSliceS i n k u)+  Ast.AstN1S SignumOp u -> signum (astSliceS i n k u)+  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astSliceS i n k u)+  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (astSliceS i n k u)+                                             (astSliceS i n k v)+  Ast.AstI2S opCode u v -> Ast.AstI2S opCode (astSliceS i n k u)+                                             (astSliceS i n k v)+  Ast.AstFloorS a -> astFloorS $ astSliceS i n k a+  Ast.AstFromIntegralS v -> astFromIntegralS $ astSliceS i n k v+  Ast.AstCastS v -> astCastS $ astSliceS i n k v+  AstConcreteS v -> astConcreteS (tsslice i n k $ Concrete v)+  _ -> Ast.AstSliceS i n k v1++astReverseS :: forall n sh s r. AstSpan s+            => AstTensor AstMethodLet s (TKS2 (n ': sh) r)+            -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)+astReverseS (Ast.AstFromVector snat stk l) =+  astFromVector snat stk $ V.reverse l+astReverseS (Ast.AstReplicate snat stk v) = astReplicate snat stk v+astReverseS (Ast.AstCond b a2 a3) = astCond b (astReverseS a2) (astReverseS a3)+astReverseS (Ast.AstLet var u v) = astLet var u (astReverseS v)+astReverseS (Ast.AstFromPrimal v) = Ast.AstFromPrimal $ astReverseS v+astReverseS (Ast.AstFromDual v) = Ast.AstFromDual $ astReverseS v+astReverseS (Ast.AstGatherS @shm @shn @shp+                            shn v ((::$) @k (Const var) vars, ix)) =+  let ivar = valueOf @k - astVar var+      ix2 = substituteAstIxS ivar var ix  -- cheap subst, because ivar is tiny+  in astGatherS @shm @shn @shp shn v (Const var ::$ vars, ix2)+astReverseS (Ast.AstReverseS v) = v+astReverseS (Ast.AstTransposeS perm@(SNat' @1 `PCons` SNat' @0 `PCons` PNil) t)+  | STKS (snat :$$ sh2@(_ :$$ _)) x <- ftkToSTK $ ftkAst t+  , Just u <- unRepl1 t =+    astTransposeS perm+    $ astReplicate snat (STKS sh2 x)+    $ astReverseS u+astReverseS (AstConcreteS v) = astConcreteS (tsreverse $ Concrete v)+astReverseS v = Ast.AstReverseS v++-- TODO: try to completely cover the AstGatherS case here, which would permit+-- not expanding to astTransposeAsGatherS in astGatherCase+-- | Beware, this does not do full simplification, which often requires+-- the gather form, so astTransposeAsGather needs to be called in addition+-- if full simplification is required.+astTransposeS+  :: forall perm sh s r.+     (Permutation.IsPermutation perm, Rank perm <= Rank sh, AstSpan s)+  => Permutation.Perm perm -> AstTensor AstMethodLet s (TKS2 sh r)+  -> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)+astTransposeS perm t =+  gcastWith (unsafeCoerceRefl+             :: TakeLen perm sh ++ DropLen perm sh :~: sh) $ case perm of+ PNil -> t+ PCons (SNat' @0) PNil ->+   gcastWith (unsafeCoerceRefl :: Permutation.PermutePrefix '[0] sh :~: sh) $+   t+ _ | FTKS sh _ <- ftkAst t+   , Just u2 <- unReplN @_ @_ @(DropLen perm sh) (shsTakeLen perm sh) t ->+   astReplicateNS (shsPermute perm (shsTakeLen perm sh)) u2+ _ -> case t of+  Ast.AstFromVector snat@(SNat @n) (STKS @sh2 sh2 x) l+    | SNat' @0 `PCons` _ <- perm -> case permUnShift1 perm of+      (perm2 :: Permutation.Perm perm2) ->+        fromMaybe (error "astTransposeS: impossible non-permutation")+        $ Permutation.permCheckPermutation perm2+        $ gcastWith (unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)+        -- for GHC 9.10 only:+        $ gcastWith (unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)+        $ gcastWith (unsafeCoerceRefl+                     :: Permutation.PermutePrefix perm (n : sh2)+                        :~: n : Permutation.PermutePrefix perm2 sh2)+        $ astFromVector snat (STKS (shsPermutePrefix perm2 sh2) x)+                        (V.map (astTransposeS perm2) l)+  Ast.AstSum snat@(SNat @n) (STKS sh x) v ->+    let zsuccP :: Permutation.Perm (0 : Permutation.MapSucc perm)+        zsuccP = Permutation.permShift1 perm+    in+      gcastWith (unsafeCoerceRefl :: Rank (0 : Permutation.MapSucc perm)+                                     :~: 1 + Rank perm) $+      gcastWith (unsafeCoerceRefl+                 :: Permutation.PermutePrefix+                      (0 : Permutation.MapSucc perm) (n : sh)+                    :~: n : Permutation.PermutePrefix perm sh) $+      fromMaybe (error "astTransposeS: impossible non-permutation")+      $ Permutation.permCheckPermutation zsuccP+      $ astSum snat (STKS (shsPermutePrefix perm sh) x) $ astTransposeS zsuccP v+  _ | SNat' @0 `PCons` _ <- perm+    , STKS ((:$$) @n @sh2 snat sh2) x <- ftkToSTK $ ftkAst t+    , Just u <- unRepl1 t -> case permUnShift1 perm of+      (perm2 :: Permutation.Perm perm2) ->+        fromMaybe (error "astTransposeS: impossible non-permutation")+        $ Permutation.permCheckPermutation perm2+        $ gcastWith (unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)+        -- for GHC 9.10 only:+        $ gcastWith (unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)+        $ gcastWith (unsafeCoerceRefl+                     :: Permutation.PermutePrefix perm (n : sh2)+                        :~: n : Permutation.PermutePrefix perm2 sh2)+        $ astReplicate snat (STKS (shsPermutePrefix perm2 sh2) x)+                       (astTransposeS perm2 u)+  -- This increases term size and work, so limited to size 2.+  Ast.AstReplicate snat1@SNat _+                   (Ast.AstFromVector snat2@(SNat' @2) stk2@STKScalar l)+    | SNat' @1 `PCons` SNat' @0 `PCons` PNil <- perm ->+      astFromVector snat2 (STKS (snat1 :$$ ZSS) stk2)+                    (V.map (astReplicate snat1 stk2) l)+  -- TODO: generalize+  Ast.AstReplicate snat1@SNat _+                   (Ast.AstFromVector snat2@(SNat' @2) stk2@(STKS sh x) l)+    | SNat' @1 `PCons` SNat' @0 `PCons` PNil <- perm ->+      astFromVector snat2 (STKS (snat1 :$$ sh) x)+                    (V.map (astReplicate snat1 stk2) l)+  Ast.AstCond b u v -> astCond b (astTransposeS perm u) (astTransposeS perm v)++  Ast.AstLet var u v -> astLet var u (astTransposeS perm v)++  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astTransposeS perm v+  Ast.AstFromDual v -> Ast.AstFromDual $ astTransposeS perm v++  AstPlusS u v -> astTransposeS perm u + astTransposeS perm v+  AstTimesS u v -> astTransposeS perm u * astTransposeS perm v+  Ast.AstN1S NegateOp u -> negate (astTransposeS perm u)+  Ast.AstN1S AbsOp u -> abs (astTransposeS perm u)+  Ast.AstN1S SignumOp u -> signum (astTransposeS perm u)+  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astTransposeS perm u)+  Ast.AstR2S opCode u v ->+    Ast.AstR2S opCode (astTransposeS perm u) (astTransposeS perm v)+  Ast.AstI2S opCode u v ->+    Ast.AstI2S opCode (astTransposeS perm u) (astTransposeS perm v)+  AstConcreteS v -> astConcreteS (tstranspose perm $ Concrete v)+--  Ast.AstFloorS v -> astFloorS $ astTransposeS perm v+--  Ast.AstFromIntegralS v -> astFromIntegralS $ astTransposeS perm v+--  Ast.AstCastS v -> astCastS $ astTransposeS perm v++  Ast.AstIndexS @shm shn v ix | SNat @n <- ixsRank ix ->+    Permutation.permFromList+      (Permutation.permToList'+       $ iterate (unsafeCoerce Permutation.permShift1) perm+         !! (valueOf @n))  -- this has a fake type, but that's fine+      $ \ (permn :: Perm permn) ->+        fromMaybe (error "astTransposeS: impossible non-permutation")+        $ Permutation.permCheckPermutation permn+        $ gcastWith (unsafeCoerceRefl+                     :: Permutation.PermutePrefix permn (shm ++ sh)+                        :~: shm ++ Permutation.PermutePrefix perm sh)+-- should suffice, but it doesn't+--      $ gcastWith (unsafeCoerceRefl :: Rank permn :~: n + Rank perm)+        $ gcastWith (unsafeCoerceRefl+                     :: (Rank permn <=? Rank (shm ++ sh)) :~: True)+        $ astIndexS (shsPermutePrefix perm shn) (astTransposeS permn v) ix+  Ast.AstScatterS @shm @shn @shp shn v (vars, ix)+    -- TODO: should the below be backpermute or permute?+    | gcompare (Permutation.permRank perm) (ixsRank ix) /= GGT ->+        let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix perm shp)+            ix2 = ixsPermutePrefix perm ix+        in gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm shp ++ shn+                         :~: Permutation.PermutePrefix perm (shp ++ shn)) $+           astScatterS @shm @shn @(Permutation.PermutePrefix perm shp)+                       shn v (vars, ix2)+  Ast.AstGatherS @shm @shn @shp shn v (vars, ix)+    -- TODO: should the below be backpermute or permute?+    | gcompare (Permutation.permRank perm) (listsRank vars) /= GGT ->+        let vars2 :: AstVarListS (Permutation.PermutePrefix perm shm)+            vars2 = listsPermutePrefix perm vars+        in gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix perm shm ++ shn+                         :~: Permutation.PermutePrefix perm (shm ++ shn)) $+           astGatherS @(Permutation.PermutePrefix perm shm) @shn @shp+                      shn v (vars2, ix)+  Ast.AstAppendS u v | SNat' @0 `PCons` _ <- perm+                     , FTKS ((:$$) @m @sh2 _ _) _ <- ftkAst u+                     , FTKS ((:$$) @n _ _) _ <- ftkAst v ->+    gcastWith (unsafeCoerceRefl+               :: Permutation.PermutePrefix perm ((m + n) : sh2)+                  :~: m + n : Tail (Permutation.PermutePrefix+                                      perm (m : sh2))) $+    gcastWith (unsafeCoerceRefl+               :: Permutation.PermutePrefix perm ((m + n) : sh2)+                  :~: m + n : Tail (Permutation.PermutePrefix+                                      perm (n : sh2))) $+    astAppendS (astTransposeS perm u) (astTransposeS perm v)+  Ast.AstSliceS i n@(SNat @n) k u | SNat' @0 `PCons` _ <- perm+                                  , FTKS ((:$$) @ink @sh2 _ _) _ <- ftkAst u ->+    gcastWith (unsafeCoerceRefl+               :: Permutation.PermutePrefix perm (n : sh2)+                  :~: n : Tail (Permutation.PermutePrefix perm (ink : sh2))) $+    astSliceS i n k (astTransposeS perm u)+  Ast.AstReverseS u | SNat' @0 `PCons` _ <- perm ->+    astReverseS (astTransposeS perm u)+  Ast.AstTransposeS @_ @sh2 perm2 u | FTKS sh2 _ <- ftkAst u ->+    -- TODO: try to perform at type level+    let permV = Permutation.permToList' perm+        perm2V = Permutation.permToList' perm2+        perm2Matched =+          perm2V+          ++ take (length permV - length perm2V) (drop (length perm2V) [0 ..])+        perm3V = normalizePermutationHack+                 $ backpermutePrefixList permV perm2Matched+    in Permutation.permFromList perm3V $ \(perm3 :: Permutation.Perm perm3) ->+      fromMaybe (error "astTransposeS: impossible non-permutation")+      $ Permutation.permCheckPermutation perm3+      $ gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix perm3 sh2+                      :~: Permutation.PermutePrefix perm sh) $+        case compare (length perm3V) (shsLength sh2) of+          LT -> gcastWith (unsafeCoerceRefl+                           :: Compare (Rank perm3) (Rank sh2) :~: LT) $+                astTransposeS perm3 u+          EQ -> gcastWith (unsafeCoerceRefl+                           :: Compare (Rank perm3) (Rank sh2) :~: EQ) $+                astTransposeS perm3 u+          GT -> error "astTransposeS: GT"+  u -> Ast.AstTransposeS perm u++-- TODO: try to cover the AstGatherS case here, which would permit+-- not expanding to astReshapeAsGatherS in astGatherCase+-- | Beware, this does not do full simplification, which often requires+-- the gather form, so astReshapeAsGather needs to be called in addition+-- if full simplification is required.+astReshapeS :: forall sh sh2 x s. (Product sh ~ Product sh2, AstSpan s)+            => ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh x)+            -> AstTensor AstMethodLet s (TKS2 sh2 x)+astReshapeS sh2 t = case t of+  Ast.AstReplicate (SNat' @1) STKS{} x -> astReshapeS sh2 x+  _ | Just u <- unRepl t+    , Refl <- lemAppNil @sh2 -> astReplicateNS sh2 u+  _ | STKS ((:$$) @_ @sh1 k _) x <- ftkToSTK $ ftkAst t+    , Just u <- unRepl1 t+    , (:$$) @_ @rest2 k2 rest2 <- sh2+    , Just Refl <- testEquality k k2 ->+      gcastWith (unsafeCoerceRefl :: Product rest2 :~: Product sh1) $+      astReplicate k (STKS rest2 x) $ astReshapeS rest2 u+  Ast.AstLet var u v -> astLet var u (astReshapeS @_ @sh2 sh2 v)+  Ast.AstFromPrimal v -> Ast.AstFromPrimal $ astReshapeS sh2 v+  Ast.AstFromDual v -> Ast.AstFromDual $ astReshapeS sh2 v+  -- Reshaping can be costly, so we don't touch AstTimesS, etc.+  Ast.AstN1S NegateOp u -> negate (astReshapeS @_ @sh2 sh2 u)+  Ast.AstN1S AbsOp u -> abs (astReshapeS @_ @sh2 sh2 u)+  Ast.AstN1S SignumOp u -> signum (astReshapeS @_ @sh2 sh2 u)+  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astReshapeS @_ @sh2 sh2 u)+  Ast.AstReshapeS _ v -> astReshapeS @_ @sh2 sh2 v+  _ | FTKS sh _ <- ftkAst t -> case testEquality sh sh2 of+    Just Refl -> t+    _ -> Ast.AstReshapeS sh2 t++astConvert+  :: AstSpan s+  => TKConversion y z -> AstTensor AstMethodLet s y+  -> AstTensor AstMethodLet s z+astConvert c a = case (ftkAst a, convertFTK c (ftkAst a)) of+  (yftk, zftk) | Just Refl <- matchingFTK yftk zftk -> a+    -- this covers the ConvId case and more, so not simplifying c at worst+    -- causes c to take more memory but doesn't inhibit rewriting+  _ | Ast.AstConvert c2 t2 <- a -> astConvert (ConvCmp c c2) t2+  (FTKScalar @ry, zftk@(FTKS ZSS (FTKScalar @rz)))+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      astConvertSFromK c zftk a+  (FTKR shr xy, zftk@(FTKS sh xz))+    | Just Refl <- matchingFTK xy xz+    , Just Refl <- testEquality (shrRank shr) (shsRank sh) ->+      astConvertSFromR c zftk a+  (FTKX shx xy, zftk@(FTKS sh xz))+    | Just Refl <- matchingFTK xy xz+    , Just Refl <- testEquality (shxRank shx) (shsRank sh) ->+      astConvertSFromX c zftk a+  (FTKS{}, zftk) -> astConvertFromS c zftk a+  (yftk@FTKProduct{}, zftk) | checkFtkAstFromS yftk zftk ->+    astConvertFromS c zftk a+  (yftk@FTKProduct{}, zftk) | checkFtkAstSFrom yftk zftk ->+    astConvertSFrom c zftk a+  _ -> Ast.AstConvert c a++-- We are pulling conversions from shaped tensors up, generally.+-- z is shaped or a product (presumably with some shaped components,+-- but it's not checked; the other components are supposed to be+-- converted identically, which is not checked in c, either).+astConvertFromS+  :: AstSpan s+  => TKConversion y z -> FullShapeTK z -> AstTensor AstMethodLet s y+  -> AstTensor AstMethodLet s z+astConvertFromS c zftk a = case (zftk, a) of+  (_, Ast.AstConvert c2 a2) -> astConvert (ConvCmp c c2) a2+  (FTKScalar @r1, AstConcreteS @r2 v)+    | ZSS <- Nested.sshape v+    , Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->+      AstConcreteK (Nested.sunScalar v)+  (_, Ast.AstFromPrimal v) ->+    Ast.AstFromPrimal $ astConvertFromS c zftk v+      -- a rare case where we don't pull up but push down so that conversions+      -- don't end up interspersed with AstFromPrimal and similar+  (_, Ast.AstFromDual v) ->+    Ast.AstFromDual $ astConvertFromS c zftk v+  (FTKScalar, Ast.AstCond b v1 v2) ->+    astCond b (astConvertFromS c FTKScalar v1)+              (astConvertFromS c FTKScalar v2)+      -- for scalars, we don't pull up but push down+  (FTKScalar, Ast.AstLet var u v) ->+    astLet var u (astConvertFromS c FTKScalar v)+  (FTKScalar, AstPlusS u v) ->+    astConvertFromS c FTKScalar u + astConvertFromS c FTKScalar v+  (FTKScalar, AstTimesS u v) ->+    astConvertFromS c FTKScalar u * astConvertFromS c FTKScalar v+  (FTKScalar, Ast.AstN1S NegateOp u) ->+    negate (astConvertFromS c FTKScalar u)+  (FTKScalar, Ast.AstN1S AbsOp u) ->+    abs (astConvertFromS c FTKScalar u)+  (FTKScalar, Ast.AstN1S SignumOp u) ->+    signum (astConvertFromS c FTKScalar u)+  (FTKScalar @r1, Ast.AstI2S @r2 QuotOp u v)+    | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->+      astConvertFromS c FTKScalar u+      `quotH` astConvertFromS c FTKScalar v+  (FTKScalar @r1, Ast.AstI2S @r2 RemOp u v)+    | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->+      astConvertFromS c FTKScalar u+      `remH` astConvertFromS c FTKScalar v+  _ -> Ast.AstConvert c a  -- by default we pull up++-- We are pushing conversions to shaped tensors down, into concrete values+-- and towards variables, so that the conversions often cancel out.+astConvertSFromK :: forall r s. AstSpan s+                 => TKConversion (TKScalar r) (TKS '[] r)+                 -> FullShapeTK (TKS '[] r)+                 -> AstTensor AstMethodLet s (TKScalar r)+                 -> AstTensor AstMethodLet s (TKS '[] r)+astConvertSFromK c zftk@(FTKS ZSS FTKScalar) a0 = case a0 of+  Ast.AstConvert c2 a2 -> astConvert (ConvCmp c c2) a2+  {- TODO: this is the right thing to do, but it results in unreadable+     terms with big tconvert wrappers over product type variables+  Ast.AstProject1 t | FTKProduct _ ftk2 <- ftkAst t ->+    astProject1 $ astSFrom' (ftkToSTK $ FTKProduct zftk ftk2) t+  Ast.AstProject2 t | FTKProduct ftk1 _ <- ftkAst t ->+    astProject2 $ astSFrom' (ftkToSTK $ FTKProduct ftk1 zftk) t -}+  Ast.AstProject1{} -> Ast.AstConvert c a0+  Ast.AstProject2{} -> Ast.AstConvert c a0+  Ast.AstSum snat@SNat STKScalar a -> astSum snat (STKS ZSS STKScalar) a+  AstConcreteK k -> AstConcreteS $ Nested.sscalar k+  Ast.AstFloorK{} -> Ast.AstConvert c a0+  Ast.AstFromIntegralK{} -> Ast.AstConvert c a0+  Ast.AstCastK{} -> Ast.AstConvert c a0+  AstPlusK u v ->+    astConvertSFromK c zftk u + astConvertSFromK c zftk v+  AstTimesK u v ->+    astConvertSFromK c zftk u * astConvertSFromK c zftk v+  Ast.AstN1K NegateOp u ->+    negate (astConvertSFromK c zftk u)+  Ast.AstN1K AbsOp u ->+    abs (astConvertSFromK c zftk u)+  Ast.AstN1K SignumOp u ->+    signum (astConvertSFromK c zftk u)+  Ast.AstR1K opCode u -> Ast.AstR1S opCode (astConvertSFromK c zftk u)+  Ast.AstR2K opCode u v ->+    Ast.AstR2S opCode (astConvertSFromK c zftk u)+                      (astConvertSFromK c zftk v)+  Ast.AstI2K QuotOp u v ->+    astConvertSFromK c zftk u `quotH` astConvertSFromK c zftk v+  Ast.AstI2K RemOp u v ->+    astConvertSFromK c zftk u `remH` astConvertSFromK c zftk v+  Ast.AstApply{} -> Ast.AstConvert c a0+  Ast.AstVar{} -> Ast.AstConvert c a0+  Ast.AstCond b v w -> astCond b (astConvertSFromK c zftk v)+                                 (astConvertSFromK c zftk w)+  Ast.AstLet var u v -> astLet var u (astConvertSFromK c zftk v)+  Ast.AstPrimalPart a -> astPrimalPart $ astConvertSFromK c zftk a+  Ast.AstDualPart a -> astDualPart $ astConvertSFromK c zftk a+  Ast.AstFromPrimal a -> Ast.AstFromPrimal $ astConvertSFromK c zftk a+  Ast.AstFromDual a -> Ast.AstFromDual $ astConvertSFromK c zftk a++astConvertSFromR :: forall sh x s. AstSpan s+                 => TKConversion (TKR2 (Rank sh) x) (TKS2 sh x)+                 -> FullShapeTK (TKS2 sh x)+                 -> AstTensor AstMethodLet s (TKR2 (Rank sh) x)+                 -> AstTensor AstMethodLet s (TKS2 sh x)+astConvertSFromR c zftk@(FTKS sh x) a0 = case a0 of+  Ast.AstConvert c2 a2 -> astConvert (ConvCmp c c2) a2+  {- TODO: this is the right thing to do, but it results in unreadable+     terms with big tconvert wrappers over product type variables+  Ast.AstProject1 t | FTKProduct _ ftk2 <- ftkAst t ->+    astProject1 $ astSFrom' (ftkToSTK $ FTKProduct zftk ftk2) t+  Ast.AstProject2 t | FTKProduct ftk1 _ <- ftkAst t ->+    astProject2 $ astSFrom' (ftkToSTK $ FTKProduct ftk1 zftk) t -}+  Ast.AstProject1{} -> Ast.AstConvert c a0+  Ast.AstProject2{} -> Ast.AstConvert c a0+  -- TODO: here and elsewhere, make sure the generated c2 is unique/correct+  Ast.AstFromVector snat@SNat (STKR @n _ xstk) l -> case sh of+    snat2 :$$ rest | Just Refl <- sameNat snat snat2+                   , Refl <- lemRankReplicate (Proxy @n) ->+      let c2 = ConvCmp (ConvXS' (FTKS rest x)) ConvRX+      in astFromVector snat (STKS rest xstk)+                       (V.map (astConvert c2) l)+    _ -> error "astConvertSFromR: impossible shape"+  Ast.AstSum snat@SNat (STKR @n _ xstk) a+    | Refl <- lemRankReplicate (Proxy @(1 + n)) ->+      let c2 = ConvCmp (ConvXS' (FTKS (snat :$$ sh) x)) ConvRX+          !a2 = astConvert c2 a+      in astSum snat (STKS sh xstk) a2+  Ast.AstReplicate snat@SNat (STKR @n _ xstk) a -> case sh of+    snat2 :$$ rest | Just Refl <- sameNat snat snat2+                   , Refl <- lemRankReplicate (Proxy @n) ->+      let c2 = ConvCmp (ConvXS' (FTKS rest x)) ConvRX+          !a2 = astConvert c2 a+      in astReplicate snat (STKS rest xstk) a2+    _ -> error "astConvertSFromR: impossible shape"+  Ast.AstApply{} -> Ast.AstConvert c a0+  Ast.AstVar{} -> Ast.AstConvert c a0+  Ast.AstCond b v w -> astCond b (astConvertSFromR c zftk v)+                                 (astConvertSFromR c zftk w)+  Ast.AstBuild1 snat@SNat (STKR @n _ xstk) (var, a) -> case sh of+    snat2 :$$ rest | Just Refl <- sameNat snat snat2+                   , Refl <- lemRankReplicate (Proxy @n) ->+      let c2 = ConvCmp (ConvXS' (FTKS rest x)) ConvRX+          !a2 = astConvert c2 a+      in Ast.AstBuild1 snat (STKS rest xstk) (var, a2)+    _ -> error "astConvertSFromR: impossible shape"+  Ast.AstLet var u v -> astLet var u (astConvertSFromR c zftk v)+  Ast.AstPrimalPart a -> astPrimalPart $ astConvertSFromR c zftk a+  Ast.AstDualPart a -> astDualPart $ astConvertSFromR c zftk a+  Ast.AstFromPrimal a ->+    Ast.AstFromPrimal $ astConvertSFromR c zftk a+  Ast.AstFromDual a ->+    Ast.AstFromDual $ astConvertSFromR c zftk a++astConvertSFromX :: forall sh shx x s. (AstSpan s, Rank shx ~ Rank sh)+                 => TKConversion (TKX2 shx x) (TKS2 sh x)+                 -> FullShapeTK (TKS2 sh x)+                 -> AstTensor AstMethodLet s (TKX2 shx x)+                 -> AstTensor AstMethodLet s (TKS2 sh x)+astConvertSFromX c zftk@(FTKS sh x) a0 = case a0 of+  Ast.AstConvert c2 a2 -> astConvert (ConvCmp c c2) a2+  {- TODO: this is the right thing to do, but it results in unreadable+     terms with big tconvert wrappers over product type variables+  Ast.AstProject1 t | FTKProduct _ ftk2 <- ftkAst t ->+    astProject1 $ astSFrom' (ftkToSTK $ FTKProduct zftk ftk2) t+  Ast.AstProject2 t | FTKProduct ftk1 _ <- ftkAst t ->+    astProject2 $ astSFrom' (ftkToSTK $ FTKProduct ftk1 zftk) t -}+  Ast.AstProject1{} -> Ast.AstConvert c a0+  Ast.AstProject2{} -> Ast.AstConvert c a0+  Ast.AstFromVector snat@SNat (STKX @shx2 _ xstk) l -> case sh of+    (:$$) @_ @rest snat2 rest | Just Refl <- sameNat snat snat2 ->+      -- This is needed only for GHC 9.10.+      gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank shx2) $+      let c2 = ConvXS' (FTKS rest x)+      in astFromVector snat (STKS rest xstk)+                       (V.map (astConvert c2) l)+    _ -> error "astConvertSFromX: impossible shape"+  Ast.AstSum snat@SNat (STKX _ xstk) a ->+      let c2 = ConvXS' (FTKS (snat :$$ sh) x)+          !a2 = astConvert c2 a+      in astSum snat (STKS sh xstk) a2+  Ast.AstReplicate snat@SNat (STKX @shx2 _ xstk) a -> case sh of+    (:$$) @_ @rest snat2 rest | Just Refl <- sameNat snat snat2 ->+      -- This is needed only for GHC 9.10.+      gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank shx2) $+      let c2 = ConvXS' (FTKS rest x)+          !a2 = astConvert c2 a+      in astReplicate snat (STKS rest xstk) a2+    _ -> error "astConvertSFromX: impossible shape"+  Ast.AstApply{} -> Ast.AstConvert c a0+  Ast.AstVar{} -> Ast.AstConvert c a0+  Ast.AstCond b v w -> astCond b (astConvertSFromX c zftk v)+                                 (astConvertSFromX c zftk w)+  Ast.AstBuild1 snat@SNat (STKX @shx2 _ xstk) (var, a) -> case sh of+    (:$$) @_ @rest snat2 rest | Just Refl <- sameNat snat snat2 ->+      -- This is needed only for GHC 9.10.+      gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank shx2) $+      let c2 = ConvXS' (FTKS rest x)+          !a2 = astConvert c2 a+      in Ast.AstBuild1 snat (STKS rest xstk) (var, a2)+    _ -> error "astConvertSFromX: impossible shape"+  Ast.AstLet var u v -> astLet var u (astConvertSFromX c zftk v)+  Ast.AstPrimalPart a -> astPrimalPart $ astConvertSFromX c zftk a+  Ast.AstDualPart a -> astDualPart $ astConvertSFromX c zftk a+  Ast.AstFromPrimal a -> Ast.AstFromPrimal $ astConvertSFromX c zftk a+  Ast.AstFromDual a -> Ast.AstFromDual $ astConvertSFromX c zftk a++-- We are pushing conversions to shaped tensors down, into concrete values+-- and towards variables, so that the conversions often cancel out.+astConvertSFrom :: forall y z s. AstSpan s+                => TKConversion y z -> FullShapeTK z+                -> AstTensor AstMethodLet s y+                -> AstTensor AstMethodLet s z+astConvertSFrom c zftk t = case (zftk, ftkAst t) of+  (_, yftk) | Just Refl <- matchingFTK yftk zftk -> t+  _ | Ast.AstConvert c2 t2 <- t -> astConvert (ConvCmp c c2) t2+  (FTKS ZSS (FTKScalar @rz), FTKScalar @ry) ->+    case testEquality (typeRep @ry) (typeRep @rz) of+      Just Refl -> astConvertSFromK c zftk t+      Nothing -> error "astConvertSFrom: tensor kinds don't match"+  (FTKS shz zx, FTKR shr yx)->+    case (matchingFTK yx zx, testEquality (shsRank shz) (shrRank shr)) of+      (Just Refl, Just Refl) -> astConvertSFromR c zftk t+      _ -> error "astConvertSFrom: tensor kinds don't match"+  (FTKS shz zx, FTKX shy yx) ->+    case (matchingFTK yx zx, testEquality (shsRank shz) (shxRank shy)) of+      (Just Refl, Just Refl) -> astConvertSFromX c zftk t+      _ -> error "astConvertSFrom: tensor kinds don't match"+  (FTKProduct zftk1 zftk2, FTKProduct yftk1 yftk2) -> case t of+    Ast.AstPair a1 a2 ->+      -- Here we can't always use the c the user presumably wrote,+      -- so we always create a canonical one.+      astPair (astConvertSFrom (convSFrom yftk1 (ftkToSTK zftk1)) zftk1 a1)+              (astConvertSFrom (convSFrom yftk2 (ftkToSTK zftk2)) zftk2 a2)+    _ -> Ast.AstConvert c t  -- don't introduce let just to push a conversion+  (_, yftk) ->+    error $ "astConvertSFrom: wrong tensor kinds: " ++ show (yftk, zftk, t)++pattern AstSFromK' :: () => sh ~ '[]+                   => AstTensor AstMethodLet s (TKScalar r)+                   -> AstTensor AstMethodLet s (TKS sh r)+pattern AstSFromK' t <-+  Ast.AstConvert c (checkPatternAstSFromK' c -> Just (Refl, t))++checkPatternAstSFromK' :: TKConversion y (TKS2 sh (TKScalar r))+                       -> AstTensor AstMethodLet s y+                       -> Maybe ( sh :~: '[]+                                , AstTensor AstMethodLet s (TKScalar r) )+checkPatternAstSFromK' c t+  | FTKScalar @ry <- ftkAst t+  , FTKS ZSS (FTKScalar @r) <- convertFTK c (ftkAst t)+  , Just Refl <- testEquality (typeRep @ry) (typeRep @r) = Just (Refl, t)+checkPatternAstSFromK' _ _ = Nothing++-- The smart constructor with the prime suffix create canonical conversions+-- instead of taking a conversion as an argument.+astFromS' :: forall y z s. AstSpan s+          => FullShapeTK z -> AstTensor AstMethodLet s y+          -> AstTensor AstMethodLet s z+astFromS' zftk t = astConvertFromS (convFromS (ftkAst t) zftk) zftk t++astSFrom' :: forall y z s. AstSpan s+          => SingletonTK z -> AstTensor AstMethodLet s y+          -> AstTensor AstMethodLet s z+astSFrom' zstk t = astConvert (convSFrom (ftkAst t) zstk) t++astKFromS' :: forall r s. (AstSpan s, GoodScalar r)+           => AstTensor AstMethodLet s (TKS2 '[] (TKScalar r))+           -> AstTensor AstMethodLet s (TKScalar r)+astKFromS' t = astConvertFromS (ConvCmp ConvX0 ConvSX) FTKScalar t++-- Or should we take SNat (Rank sh) to help proving n ~ Rank sh?+astRFromS' :: forall sh x s. AstSpan s+           => AstTensor AstMethodLet s (TKS2 sh x)+           -> AstTensor AstMethodLet s (TKR2 (Rank sh) x)+astRFromS' t | FTKS sh x <- ftkAst t+             , Refl <- lemRankMapJust sh =+  let zftk = FTKR (shrFromShS sh) x+  in astConvertFromS (ConvCmp (ConvXR (ftkToSTK x)) ConvSX) zftk t++astXFromS' :: forall sh shx x s. (AstSpan s, Rank sh ~ Rank shx)+           => StaticShX shx -> AstTensor AstMethodLet s (TKS2 sh x)+           -> AstTensor AstMethodLet s (TKX2 shx x)+astXFromS' ssx t | FTKS sh x <- ftkAst t+                 , Refl <- lemRankMapJust sh =+  let zftk = FTKX (shCastSX ssx sh) x+  in astConvertFromS (ConvCmp (ConvXX' zftk) ConvSX) zftk t++astSFromK' :: forall r s. (GoodScalar r, AstSpan s)+           => AstTensor AstMethodLet s (TKScalar r)+           -> AstTensor AstMethodLet s (TKS '[] r)+astSFromK' a =+  let c2 = ConvCmp ConvXS (Conv0X STKScalar)+  in astConvertSFromK c2 (FTKS ZSS FTKScalar) a++astSFromR' :: forall sh s r. AstSpan s+           => ShS sh -> AstTensor AstMethodLet s (TKR2 (Rank sh) r)+           -> AstTensor AstMethodLet s (TKS2 sh r)+astSFromR' sh t = case ftkAst t of+  FTKR _ x | Refl <- lemRankReplicate (Proxy @(Rank sh)) ->+    let zftk = FTKS sh x+    in astConvertSFromR (ConvCmp (ConvXS' zftk) ConvRX) zftk t++astSFromX' :: forall sh sh' s x. (AstSpan s, Rank sh ~ Rank sh')+           => ShS sh -> AstTensor AstMethodLet s (TKX2 sh' x)+           -> AstTensor AstMethodLet s (TKS2 sh x)+astSFromX' sh t = case ftkAst t of+  FTKX _ x ->+    let zftk = FTKS sh x+    in astConvertSFromX (ConvXS' zftk) zftk t++astSum0S :: AstSpan s+         => AstTensor AstMethodLet s (TKS2 sh x)+         -> AstTensor AstMethodLet s (TKS2 '[] x)+astSum0S t = case t of+  Ast.AstSum SNat _ u -> astSum0S u+  Ast.AstReplicate snat (STKS _ STKScalar) u ->+    astSum0S u * (fromPrimal $ AstConcreteS+                  $ Nested.sscalar $ fromInteger $ fromSNat snat)+  Ast.AstReplicate snat STKScalar u ->+    astSFromK' u * (fromPrimal $ AstConcreteS+                    $ Nested.sscalar $ fromInteger $ fromSNat snat)+  Ast.AstLet var u v -> astLet var u (astSum0S v)+  AstTimesS t1 t2 -> astDot0S t1 t2+  AstConcreteS v ->+    withKnownShS (Nested.sshape v) $+    astConcreteS $ tssum0 (Concrete v)+  Ast.AstFromPrimal u -> Ast.AstFromPrimal $ astSum0S u+  Ast.AstFromDual u -> Ast.AstFromDual $ astSum0S u+  Ast.AstReverseS u -> astSum0S u+  Ast.AstTransposeS _ u -> astSum0S u+  Ast.AstReshapeS _ u -> astSum0S u+  Ast.AstN1S NegateOp u -> negate $ astSum0S u+  Ast.AstSum0S u -> astSum0S u+  Ast.AstDot0S{} -> t+  Ast.AstDot1InS _ _ t1 t2 -> astDot0S t1 t2+  Ast.AstMatmul2S m@SNat SNat p@SNat m1 m2 ->+    astDot0S (astTransposeS (Permutation.makePerm @'[1, 0])+                            (astReplicate p knownSTK m1))+             (astTransposeS (Permutation.makePerm @'[0, 2, 1])+                            (astReplicate m knownSTK m2))+  _ -> Ast.AstSum0S t++astDot0S :: (GoodScalar r, AstSpan s)+         => AstTensor AstMethodLet s (TKS sh r)+         -> AstTensor AstMethodLet s (TKS sh r)+         -> AstTensor AstMethodLet s (TKS '[] r)+astDot0S t1 t2 = case (t1, t2) of+  (AstConcreteS v1, AstConcreteS v2) ->+    withKnownShS (Nested.sshape v1) $+    astConcreteS $ tsdot0 (Concrete v1) (Concrete v2)+  _ | FTKS (snat :$$ _) _ <- ftkAst t1+    , Just u1 <- unRepl1 t1+    , Just u2 <- unRepl1 t2 ->+      astDot0S u1 u2 * (fromPrimal $ AstConcreteS+                        $ Nested.sscalar $ fromInteger $ fromSNat snat)+  (Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->+    Ast.AstFromPrimal $ astDot0S u1 u2+  (Ast.AstFromDual u1, Ast.AstFromDual u2) ->+    Ast.AstFromDual $ astDot0S u1 u2+  (Ast.AstTransposeS @_ @sh1 perm1 u1, Ast.AstTransposeS @_ @sh2 perm2 u2)+    | Just Refl <- testEquality perm1 perm2 ->+      gcastWith (unsafeCoerceRefl :: sh1 :~: sh2) $+      astDot0S u1 u2+  (Ast.AstReverseS u1, Ast.AstReverseS u2) -> astDot0S u1 u2+  (Ast.AstN1S NegateOp u1, Ast.AstN1S NegateOp u2) -> astDot0S u1 u2+  _ -> Ast.AstDot0S t1 t2++astDot1InS :: forall sh n r s. GoodScalar r+           => ShS sh -> SNat n+           -> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)+           -> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)+           -> AstTensor AstMethodLet s (TKS sh r)+astDot1InS sh n@SNat t1 t2 = case (t1, t2) of+  (AstConcreteS v1, AstConcreteS v2) ->+    withKnownShS sh $+    astConcreteS $ tsdot1In @_ @sh (SNat @n) (Concrete v1) (Concrete v2)+  (Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->+    Ast.AstFromPrimal $ astDot1InS sh n u1 u2+  (Ast.AstFromDual u1, Ast.AstFromDual u2) ->+    Ast.AstFromDual $ astDot1InS sh n u1 u2+  _ -> Ast.AstDot1InS sh n t1 t2++astMatmul2S :: GoodScalar r+            => SNat m -> SNat n -> SNat p+            -> AstTensor AstMethodLet s (TKS '[m, n] r)+            -> AstTensor AstMethodLet s (TKS '[n, p] r)+            -> AstTensor AstMethodLet s (TKS '[m, p] r)+astMatmul2S m@SNat n@SNat p@SNat t1 t2 = case (t1, t2) of+  (AstConcreteS v1, AstConcreteS v2) ->+    astConcreteS $ tsmatmul2 (Concrete v1) (Concrete v2)+  (Ast.AstFromPrimal u1, Ast.AstFromPrimal u2) ->+    Ast.AstFromPrimal $ astMatmul2S m n p u1 u2+  (Ast.AstFromDual u1, Ast.AstFromDual u2) ->+    Ast.AstFromDual $ astMatmul2S m n p u1 u2+  _ -> Ast.AstMatmul2S m n p t1 t2+++-- * ConvertTensor instances needed for unwinding in astConcrete++instance AstSpan s => ConvertTensor (AstTensor AstMethodLet s) where+  tconvert c _astk v = astConvert c v++  rfromX a = case ftkAst a of+    FTKX sh' _ ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        withKnownShS sh $+        rfromS $ sfromX @_ @sh a+  xfromR a = case ftkAst a of+    FTKR shr _ ->+      withShsFromShR shr $ \(sh :: ShS sh) ->+        withKnownShS sh $+        xfromS @_ @sh $ sfromR a++  sfromR = astSFromR' knownShS+  sfromX = astSFromX' knownShS+  xfromS = astXFromS' knownShX++  rzip @_ @_ @n a+   | Refl <- lemRankReplicate (Proxy @n) = case ftkAst a of+    FTKProduct (FTKR _sh y) (FTKR _ z) ->+      let c = ConvCmp+                (ConvXR (ftkToSTK (FTKProduct y z)))+                (ConvCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvRX ConvRX))+      in astConvert c a+  runzip @_ @_ @n a+   | Refl <- lemRankReplicate (Proxy @n) = case ftkAst a of+    FTKR _sh (FTKProduct y z) ->+      let c = ConvCmp+                (ConvT2 (ConvXR (ftkToSTK y)) (ConvXR (ftkToSTK z)))+                (ConvCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvRX)+      in astConvert c a+  szip a = case ftkAst a of+    FTKProduct (FTKS _sh y) (FTKS _ z) ->+      let c = ConvCmp+                ConvXS+                (ConvCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvSX ConvSX))+      in astConvert c a+  sunzip a = case ftkAst a of+    FTKS _sh (FTKProduct y z) ->+      let c = ConvCmp+                (ConvT2 ConvXS ConvXS)+                (ConvCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvSX)+      in astConvert c a+  xzip a = case ftkAst a of+    FTKProduct (FTKX _sh y) (FTKX _ z) ->+      let c = ConvZip (ftkToSTK y) (ftkToSTK z)+      in astConvert c a+  xunzip a = case ftkAst a of+    FTKX _sh (FTKProduct y z) ->+      let c = ConvUnzip (ftkToSTK y) (ftkToSTK z)+      in astConvert c a++  xnestR @sh1 @m @x sh1 a+    | Refl <- lemRankReplicate (Proxy @m) =+      let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)+                            (TKX2 sh1 (TKR2 m x))+          c = ConvCmp+                (ConvXX (ConvXR (knownSTK @x)))+                (ConvNest @_ @_ @(Replicate m Nothing)+                          (STKX sh1 (knownSTK @x)))+      in astConvert c a+  xnestS @_ @_ @x sh1 a =+    let c = ConvCmp (ConvXX ConvXS)+                    (ConvNest (STKX sh1 (knownSTK @x)))+    in astConvert c a+  xnest @_ @_ @x sh1 a =+    let c = ConvNest (STKX sh1 (knownSTK @x))+    in astConvert c a+  xunNestR a =+    let c = ConvCmp ConvUnnest+                    (ConvXX ConvRX)+    in astConvert c a+  xunNestS a =+    let c = ConvCmp ConvUnnest+                    (ConvXX ConvSX)+    in astConvert c a+  xunNest a =+    let c = ConvUnnest+    in astConvert c a++  tpairConv = astPair+  tunpairConv t = (astProject1 t, astProject2 t)+++-- * Helper combinators++-- All but the last case are shortcuts for common forms.+astConcrete :: FullShapeTK y -> Concrete y+            -> AstTensor AstMethodLet PrimalSpan y+astConcrete ftk v = case ftk of+  FTKScalar -> astConcreteK v+  FTKR sh' FTKScalar ->+    withShsFromShR sh' $ \(sh :: ShS sh) ->+      withKnownShS sh $+      astFromS' ftk $ astConcreteS (sfromR @_ @sh v)+  FTKS _ FTKScalar -> astConcreteS v+  FTKX sh' FTKScalar ->+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      withKnownShS sh $+      astFromS' ftk $ astConcreteS (sfromX @_ @sh v)+  FTKProduct ftk1 ftk2 ->+    astPair (astConcrete ftk1 (tproject1 v)) (astConcrete ftk2 (tproject2 v))+  _ -> concreteTarget astConcreteK astConcreteS astFromS' (ftkToSTK ftk) v++astLetFun :: forall y z s s2. (AstSpan s, AstSpan s2)+          => AstTensor AstMethodLet s y+          -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)+          -> AstTensor AstMethodLet s2 z+astLetFun = astLetFunBounds Nothing++astLetFunB :: forall z s s2. (AstSpan s, AstSpan s2)+           => AstTensor AstMethodLet s (TKScalar Int64)+           -> (AstTensor AstMethodLet s (TKScalar Int64)+               -> AstTensor AstMethodLet s2 z)+           -> AstTensor AstMethodLet s2 z+astLetFunB w = astLetFunBounds (Just $ bounds w) w++astLetFunBounds :: forall y z s s2. (AstSpan s, AstSpan s2)+                => Maybe (Int64, Int64) -> AstTensor AstMethodLet s y+                -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)+                -> AstTensor AstMethodLet s2 z+astLetFunBounds _ a f | astIsSmall True a = f a+astLetFunBounds mbs a f = case a of+  AstFromS' FTKScalar _ ->+    let (var, ast) = funToAst2 (ftkAst a) mbs f+    in astLet var a ast+  AstFromS' @y2 ftkz v ->+    let (var, ast) = funToAst2 (ftkAst v) mbs (f . astFromS' @y2 ftkz)+    in astLet var v ast+  _ -> case ftkAst a of+    ftk@(FTKR @_ @x sh' x) ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        let (var, ast) =+              funToAst2 (FTKS sh x) mbs+                        (f . astFromS' @(TKS2 sh x) ftk)+        in astLet var (astSFromR' sh a) ast+             -- safe, because subsitution ruled out above+    ftk@(FTKX @_ @x sh' x) ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        let (var, ast) =+              funToAst2 (FTKS sh x) mbs+                        (f . astFromS' @(TKS2 sh x) ftk)+        in astLet var (astSFromX' sh a) ast+    -- calling recursively for product may be not worth it+    ftk -> let (var, ast) = funToAst2 ftk mbs f+           in astLet var a ast++astReplicateNS :: forall shn shp s x. AstSpan s+               => ShS shn -> AstTensor AstMethodLet s (TKS2 shp x)+               -> AstTensor AstMethodLet s (TKS2 (shn ++ shp) x)+astReplicateNS shn v | STKS shp x <- ftkToSTK (ftkAst v) =+  let go :: ShS shn' -> AstTensor AstMethodLet s (TKS2 (shn' ++ shp) x)+      go ZSS = v+      go (snat :$$ shn2) =+        astReplicate snat (STKS (shn2 `shsAppend` shp) x) $ go shn2+  in go shn++unRepl :: AstSpan s+       => AstTensor AstMethodLet s (TKS2 sh x)+       -> Maybe (AstTensor AstMethodLet s (TKS2 '[] x))+unRepl (Ast.AstReplicate _ (STKS ZSS _) u) = Just u+unRepl (Ast.AstReplicate _ STKScalar u) = Just $ astSFromK' u+unRepl (Ast.AstReplicate _ STKS{} u) = unRepl u+unRepl (AstConcreteS a) = AstConcreteS . Nested.sscalar <$> sunReplicateScal a+unRepl (Ast.AstCond b v1 v2) = do+  u1 <- unRepl v1+  u2 <- unRepl v2+  return $! astCond b u1 u2+unRepl (Ast.AstLet var u t) = Ast.AstLet var u <$> unRepl t+unRepl (Ast.AstPrimalPart t) = astPrimalPart <$> unRepl t+unRepl (Ast.AstDualPart t) = astDualPart <$> unRepl t+unRepl _ = Nothing++unRepl1 :: AstSpan s+        => AstTensor AstMethodLet s (TKS2 (n ': sh) x)+        -> Maybe (AstTensor AstMethodLet s (TKS2 sh x))+unRepl1 (Ast.AstReplicate _ STKS{} u) = Just u+unRepl1 (Ast.AstReplicate _ STKScalar u) = Just $ astSFromK' u+unRepl1 (AstConcreteS a) = AstConcreteS <$> sunReplicate1 a+unRepl1 (Ast.AstCond b v1 v2) = do+  u1 <- unRepl1 v1+  u2 <- unRepl1 v2+  return $! astCond b u1 u2+unRepl1 (Ast.AstLet var u t) = Ast.AstLet var u <$> unRepl1 t+unRepl1 (Ast.AstPrimalPart t) = astPrimalPart <$> unRepl1 t+unRepl1 (Ast.AstDualPart t) = astDualPart <$> unRepl1 t+unRepl1 _ = Nothing++unReplN :: AstSpan s+        => ShS shm -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+        -> Maybe (AstTensor AstMethodLet s (TKS2 shn x))+unReplN ZSS a = Just a+unReplN (_ :$$ ZSS) (Ast.AstReplicate _ STKScalar u) = Just $ astSFromK' u+unReplN (_ :$$ sh) (Ast.AstReplicate _ STKS{} u) = unReplN sh u+unReplN shm (AstConcreteS a) = AstConcreteS <$> sunReplicateN shm a+unReplN shm (Ast.AstCond b v1 v2) = do+  u1 <- unReplN shm v1+  u2 <- unReplN shm v2+  return $! astCond b u1 u2+unReplN shm (Ast.AstLet var u t) = Ast.AstLet var u <$> unReplN shm t+unReplN shm (Ast.AstPrimalPart t) = astPrimalPart <$> unReplN shm t+unReplN shm (Ast.AstDualPart t) = astDualPart <$> unReplN shm t+unReplN _ _ = Nothing+++-- * Substitution wrappers++-- | We assume no variable is shared between a binding and its nested binding+-- and nobody substitutes into variables that are bound.+-- This keeps the substitution code simple, because we never need to compare+-- variables to any variable in the bindings.+substituteAst :: forall s s2 y z. (AstSpan s, AstSpan s2)+              => AstTensor AstMethodLet s2 z -> AstVarName s2 z+              -> AstTensor AstMethodLet s y+              -> AstTensor AstMethodLet s y+substituteAst i var v1 =+  fromMaybe v1 $ substitute1Ast i var v1++substituteAstIxS+  :: AstSpan s+  => AstTensor AstMethodLet s y -> AstVarName s y -> AstIxS AstMethodLet sh+  -> AstIxS AstMethodLet sh+substituteAstIxS i var ix =+  fromMaybe ix $ substitute1AstIxS i var ix++substituteAstBool+  :: AstSpan s+  => AstTensor AstMethodLet s y -> AstVarName s y -> AstBool AstMethodLet+  -> AstBool AstMethodLet+substituteAstBool i var v1 =+  fromMaybe v1 $ substitute1AstBool i var v1+++-- * Substitution workers++-- Invariant: if the variable has bounds, the expression can only have+-- values within the bounds (regardless of what the `bounds` call would say).+substitute1Ast :: forall s s2 y z. (AstSpan s, AstSpan s2)+               => AstTensor AstMethodLet s2 z -> AstVarName s2 z+               -> AstTensor AstMethodLet s y+               -> Maybe (AstTensor AstMethodLet s y)+substitute1Ast i var = subst where+ subst :: forall s3 y2. AstSpan s3+       => AstTensor AstMethodLet s3 y2+       -> Maybe (AstTensor AstMethodLet s3 y2)+ subst = \case+  Ast.AstPair u v ->+    case (subst u, subst v) of+      (Nothing, Nothing) -> Nothing+      (mu, mv) -> Just $ astPair (fromMaybe u mu) (fromMaybe v mv)+  Ast.AstProject1 a -> astProject1 <$> subst a+  Ast.AstProject2 a -> astProject2 <$> subst a+  Ast.AstFromVector snat stk args ->+    let margs = V.map subst args+    in if V.any isJust margs+       then Just $ astFromVector snat stk $ V.zipWith fromMaybe args margs+       else Nothing+  Ast.AstSum snat stk v -> astSum snat stk <$> subst v+  Ast.AstReplicate snat stk v ->+    astReplicate snat stk <$> subst v+  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->+      case ( substitute1AstHFun i var f, substitute1AstHFun i var df+           , substitute1AstHFun i var rf, subst acc0+           , subst es ) of+        (Nothing, Nothing, Nothing, Nothing, Nothing) -> Nothing+        (mf, mdf, mrf, macc0, mes) ->+          Just $ astMapAccumRDer k bftk eftk+                                 (fromMaybe f mf)+                                 (fromMaybe df mdf)+                                 (fromMaybe rf mrf)+                                 (fromMaybe acc0 macc0)+                                 (fromMaybe es mes)+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+      case ( substitute1AstHFun i var f, substitute1AstHFun i var df+           , substitute1AstHFun i var rf, subst acc0+           , subst es ) of+        (Nothing, Nothing, Nothing, Nothing, Nothing) -> Nothing+        (mf, mdf, mrf, macc0, mes) ->+          Just $ astMapAccumLDer k bftk eftk+                                 (fromMaybe f mf)+                                 (fromMaybe df mdf)+                                 (fromMaybe rf mrf)+                                 (fromMaybe acc0 macc0)+                                 (fromMaybe es mes)+  Ast.AstApply t ll ->+    case ( substitute1AstHFun i var t+         , subst ll ) of+      (Nothing, Nothing) -> Nothing+      (mt, mll) -> Just $ astApply (fromMaybe t mt) (fromMaybe ll mll)+  Ast.AstVar var2 ->+    -- We can't assert anything here, because only all runtime values need+    -- to be in bounds and bounds approximations don't have to agree.+    if varNameToAstVarId var == varNameToAstVarId var2+    then case sameAstSpan @s3 @s2 of+      Just Refl -> case testEquality var var2 of+        Just Refl -> case i of+          Ast.AstVar var3 | FTKScalar <- varNameToFTK var3 ->+            let (lb, ub) = fromMaybe (-1000000000, 1000000000)+                           $ varNameToBounds var+                (lb2, ub2) = fromMaybe (-1000000000, 1000000000)+                             $ varNameToBounds var2+                (lb3, ub3) = fromMaybe (-1000000000, 1000000000)+                             $ varNameToBounds var3+                bs = (max (max lb lb2) lb3, min (min ub ub2) ub3)+                  -- We know all bounds approximations have to be correct+                  -- so we can intersect them.+            in Just $ astVar $ mkAstVarName (varNameToFTK var3)+                                            (Just bs)+                                            (varNameToAstVarId var3)+          _ -> Just i+        _ -> error $ "substitute1Ast: kind of the variable "+                     ++ show var2 ++ ": " ++ show (varNameToFTK var)+                     ++ ", payload kind: " ++ show (varNameToFTK var2)+                     ++ ", payload: " ++ show i+      _ -> error "substitute1Ast: span"+    else Nothing+  Ast.AstCond b v w ->+    case ( substitute1AstBool i var b+         , subst v+         , subst w ) of+      (Nothing, Nothing, Nothing) -> Nothing+      (mb, mv, mw) ->+        Just $ astCond (fromMaybe b mb) (fromMaybe v mv) (fromMaybe w mw)+  Ast.AstBuild1 k stk (var2, v) ->+    assert (varNameToAstVarId var2 /= varNameToAstVarId var) $+    Ast.AstBuild1 k stk . (var2,) <$> subst v++  Ast.AstLet var2 u v ->+    case (subst u, subst v) of+      (Nothing, Nothing) -> Nothing+      (mu, mv) -> Just $ astLet var2 (fromMaybe u mu) (fromMaybe v mv)++  Ast.AstPrimalPart a -> astPrimalPart <$> subst a+  Ast.AstDualPart a -> astDualPart <$> subst a+  Ast.AstFromPrimal a -> Ast.AstFromPrimal <$> subst a+  Ast.AstFromDual a -> Ast.AstFromDual <$> subst a++  AstPlusK u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ fromMaybe u mu + fromMaybe v mv+       else Nothing+  AstTimesK u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ fromMaybe u mu * fromMaybe v mv+       else Nothing+  Ast.AstN1K NegateOp u -> negate <$> subst u+  Ast.AstN1K AbsOp u -> abs <$> subst u+  Ast.AstN1K SignumOp u -> signum <$> subst u+  Ast.AstR1K opCode u -> Ast.AstR1K opCode <$> subst u+  Ast.AstR2K opCode u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ Ast.AstR2K opCode (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstI2K QuotOp u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ quotH (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstI2K RemOp u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ remH (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstConcreteK{} -> Nothing+  Ast.AstFloorK a -> astFloorK <$> subst a+  Ast.AstFromIntegralK v -> astFromIntegralK <$> subst v+  Ast.AstCastK v -> astCastK <$> subst v++  AstPlusS u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ fromMaybe u mu + fromMaybe v mv+       else Nothing+  AstTimesS u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ fromMaybe u mu * fromMaybe v mv+       else Nothing+  Ast.AstN1S NegateOp u -> negate <$> subst u+  Ast.AstN1S AbsOp u -> abs <$> subst u+  Ast.AstN1S SignumOp u -> signum <$> subst u+  Ast.AstR1S opCode u -> Ast.AstR1S opCode <$> subst u+  Ast.AstR2S opCode u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ Ast.AstR2S opCode (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstI2S QuotOp u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ quotH (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstI2S RemOp u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ remH (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstConcreteS{} -> Nothing+  Ast.AstFloorS a -> astFloorS <$> subst a+  Ast.AstFromIntegralS a -> astFromIntegralS <$> subst a+  Ast.AstCastS v -> astCastS <$> subst v++  Ast.AstIndexS shn v ix ->+    case (subst v, substitute1AstIxS i var ix) of+      (Nothing, Nothing) -> Nothing+      (mv, mix) -> Just $ astIndexS shn (fromMaybe v mv) (fromMaybe ix mix)+  Ast.AstScatterS shn v (vars, ix) ->+    case (subst v, substitute1AstIxS i var ix) of+      (Nothing, Nothing) -> Nothing+      (mv, mix) -> Just $ astScatterS shn+                                      (fromMaybe v mv)+                                      (vars, fromMaybe ix mix)+  Ast.AstGatherS shn v (vars, ix) ->+    case (subst v, substitute1AstIxS i var ix) of+      (Nothing, Nothing) -> Nothing+      (mv, mix) -> Just $ astGatherS shn+                                     (fromMaybe v mv)+                                     (vars, fromMaybe ix mix)+  Ast.AstMinIndexS a -> Ast.AstMinIndexS <$> subst a+  Ast.AstMaxIndexS a -> Ast.AstMaxIndexS <$> subst a+  Ast.AstIotaS{} -> Nothing+  Ast.AstAppendS x y ->+    case (subst x, subst y) of+      (Nothing, Nothing) -> Nothing+      (mx, my) -> Just $ astAppendS (fromMaybe x mx) (fromMaybe y my)+  Ast.AstSliceS i2 n k v -> astSliceS i2 n k <$> subst v+  Ast.AstReverseS v -> astReverseS <$> subst v+  Ast.AstTransposeS perm v -> astTransposeS perm <$> subst v+  Ast.AstReshapeS sh v -> astReshapeS sh <$> subst v++  Ast.AstConvert c v -> astConvert c <$> subst v++  Ast.AstSum0S v -> astSum0S <$> subst v+  Ast.AstDot0S u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astDot0S (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstDot1InS sh n u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astDot1InS sh n (fromMaybe u mu) (fromMaybe v mv)+       else Nothing+  Ast.AstMatmul2S m n p u v ->+    let mu = subst u+        mv = subst v+    in if isJust mu || isJust mv+       then Just $ astMatmul2S m n p (fromMaybe u mu) (fromMaybe v mv)+       else Nothing++substitute1AstIxS+  :: AstSpan s2+  => AstTensor AstMethodLet s2 y -> AstVarName s2 y -> AstIxS AstMethodLet sh+  -> Maybe (AstIxS AstMethodLet sh)+substitute1AstIxS i var ix =+  let mix = fmap (substitute1Ast i var) ix+  in if any isJust mix+     then Just $ ixsZipWith fromMaybe ix mix+     else Nothing++substitute1AstHFun+  :: forall s s2 s3 x y z.+     AstTensor AstMethodLet s3 z -> AstVarName s3 z -> AstHFun s s2 x y+  -> Maybe (AstHFun s s2 x y)+substitute1AstHFun _i _var AstLambda{} =+  Nothing  -- no outside free variables++substitute1AstBool :: AstSpan s2+                   => AstTensor AstMethodLet s2 y -> AstVarName s2 y+                   -> AstBool AstMethodLet+                   -> Maybe (AstBool AstMethodLet)+substitute1AstBool i var = subst where+ subst :: AstBool AstMethodLet+       -> Maybe (AstBool AstMethodLet)+ subst = \case+  Ast.AstBoolConst{} -> Nothing+  Ast.AstBoolNot arg -> notB <$> subst arg+  Ast.AstBoolAnd arg1 arg2 ->+    let mb1 = subst arg1+        mb2 = subst arg2+    in if isJust mb1 || isJust mb2+       then Just $ fromMaybe arg1 mb1 &&* fromMaybe arg2 mb2+       else Nothing+  Ast.AstLeqK arg1 arg2 ->+    let mr1 = substitute1Ast i var arg1+        mr2 = substitute1Ast i var arg2+    in if isJust mr1 || isJust mr2+       then Just $ fromMaybe arg1 mr1 <=. fromMaybe arg2 mr2+       else Nothing+  Ast.AstLeqS arg1 arg2 ->+    let mr1 = substitute1Ast i var arg1+        mr2 = substitute1Ast i var arg2+    in if isJust mr1 || isJust mr2+       then Just $ fromMaybe arg1 mr1 <=. fromMaybe arg2 mr2+       else Nothing
+ src/HordeAd/Core/AstTools.hs view
@@ -0,0 +1,523 @@+{-# LANGUAGE ViewPatterns #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+-- | An assortment of operations working on AST of the code to be differentiated+-- or the code resulting from differentiation.+module HordeAd.Core.AstTools+  ( -- * Full tensor kind derivation+    ftkAst, isTensorInt+    -- * Variable occurrence detection+  , varInAst, varInAstBool, varInIxS, varNameInAst, varNameInIxS+    -- * Determining if a term is too small to require sharing+  , astIsSmall, ixIsSmall+    -- * Odds and ends+  , bounds+  , liftRFromS1, liftRFromS2, liftXFromS1, liftXFromS2+  , cAstConvert, cAstSFromR, cAstSFromX, cAstXFromS+  , pattern AstFromS', checkAstFromS, checkFtkAstFromS, checkFtkAstSFrom+  , cAstFromS, cAstSFrom, convFromS, convSFrom+  , setTotalSharing+  ) where++import Prelude hiding (foldl')++import Control.Exception.Assert.Sugar+import Data.Int (Int64)+import Data.IORef+import Data.Maybe (isJust)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import System.IO.Unsafe (unsafePerformIO)+import Type.Reflection (typeRep)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (snatPlus)++import HordeAd.Core.Ast+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * Full tensor kind derivation++-- | This is cheap and dirty. We don't shape-check the terms and we don't+-- unify or produce (partial) results with unknowns. Instead, we investigate+-- only one path and fail if it doesn't contain enough information+-- to determine shape (which rarely happens in our AST that is based+-- on shaped tensors and always indicates an ill-typed term).+ftkAst :: forall s y ms. AstTensor ms s y -> FullShapeTK y+ftkAst t = case t of+  AstPair t1 t2 -> FTKProduct (ftkAst t1) (ftkAst t2)+  AstProject1 v -> case ftkAst v of+    FTKProduct ftk1 _ -> ftk1+  AstProject2 v -> case ftkAst v of+    FTKProduct _ ftk2 -> ftk2+  AstFromVector snat _ l -> case V.uncons l of+    Nothing -> error "ftkAst: empty vector"+    Just (v, _) -> buildFTK snat (ftkAst v)+  AstSum snat stk v -> razeFTK snat stk (ftkAst v)+  AstReplicate snat _ v -> buildFTK snat (ftkAst v)+  AstMapAccumRDer k bftk _eftk _f _df _rf acc0 _es ->+    FTKProduct (ftkAst acc0) (buildFTK k bftk)+  AstMapAccumLDer k bftk _eftk _f _df _rf acc0 _es ->+    FTKProduct (ftkAst acc0) (buildFTK k bftk)+  AstApply (AstLambda !_ !l) _ -> ftkAst l+  AstVar var -> varNameToFTK var+  AstCond _b v _w -> ftkAst v+  AstBuild1 snat _ (_var, v) -> buildFTK snat (ftkAst v)++  AstLet _ _ v -> ftkAst v+  AstShare var _ -> varNameToFTK var+  AstToShare v -> ftkAst v++  AstPrimalPart a -> ftkAst a+  AstDualPart a -> ftkAst a+  AstFromPrimal a -> ftkAst a+  AstFromDual a -> ftkAst a++  AstPlusK{} -> FTKScalar+  AstTimesK{} -> FTKScalar+  AstN1K{} -> FTKScalar+  AstR1K{} -> FTKScalar+  AstR2K{} -> FTKScalar+  AstI2K{} -> FTKScalar+  AstConcreteK _ -> FTKScalar+  AstFloorK{} -> FTKScalar+  AstFromIntegralK{} -> FTKScalar+  AstCastK{} -> FTKScalar++  AstPlusS v _ -> ftkAst v+  AstTimesS v _ -> ftkAst v+  AstN1S _ v -> ftkAst v+  AstR1S _ v -> ftkAst v+  AstR2S _ v _ -> ftkAst v+  AstI2S _ v _ -> ftkAst v+  AstConcreteS a -> FTKS (Nested.sshape a) FTKScalar+  AstFloorS v -> case ftkAst v of+    FTKS sh FTKScalar -> FTKS sh FTKScalar+  AstFromIntegralS v -> case ftkAst v of+    FTKS sh FTKScalar -> FTKS sh FTKScalar+  AstCastS v -> case ftkAst v of+    FTKS sh FTKScalar -> FTKS sh FTKScalar++  AstIndexS shn v _ix -> case ftkAst v of+    FTKS _ x -> FTKS shn x+  AstScatterS shn v (_ , ix) -> case ftkAst v of+    FTKS _ x -> FTKS (shsFromIxS ix `shsAppend` shn) x+  AstGatherS shn v (vars, _) -> case ftkAst v of+    FTKS _ x -> FTKS (shsFromListS vars `shsAppend` shn) x+  AstMinIndexS v -> case ftkAst v of+    FTKS sh FTKScalar -> FTKS (shsInit sh) FTKScalar+  AstMaxIndexS v -> case ftkAst v of+    FTKS sh FTKScalar -> FTKS (shsInit sh) FTKScalar+  AstIotaS n@SNat -> FTKS (n :$$ ZSS) FTKScalar+  AstAppendS a b -> case (ftkAst a, ftkAst b) of+    (FTKS (m :$$ sh) x, FTKS (n :$$ _) _) -> FTKS (snatPlus m n :$$ sh) x+  AstSliceS _ n@SNat _ a -> case ftkAst a of+    FTKS (_ :$$ sh) x -> FTKS (n :$$ sh) x+  AstReverseS v -> ftkAst v+  AstTransposeS perm v -> case ftkAst v of+    FTKS sh x -> FTKS (shsPermutePrefix perm sh) x+  AstReshapeS sh2 v -> case ftkAst v of+    FTKS _ x -> FTKS sh2 x++  AstConvert c u -> convertFTK c $ ftkAst u++  AstSum0S v ->  case ftkAst v of+    FTKS _ x -> FTKS ZSS x+  AstDot0S _u _v -> FTKS ZSS FTKScalar+  AstDot1InS sh _ _u _v -> FTKS sh FTKScalar+  AstMatmul2S m@SNat _ p@SNat _u _v -> FTKS (m :$$ p :$$ ZSS) FTKScalar++isTensorInt :: forall s y ms. AstSpan s+            => Proxy s -> FullShapeTK y+            -> Maybe (AstTensor ms s y :~: AstInt ms)+isTensorInt _ ftk = case ftk of+  FTKScalar @r -> case ( testEquality (typeRep @r) (typeRep @Int64)+                       , sameAstSpan @s @PrimalSpan ) of+                    (Just Refl, Just Refl) -> Just Refl+                    _ -> Nothing+  _ -> Nothing+++-- * Variable occurrence detection++-- | We assume no variable is shared between a binding and its nested binding+-- and nobody asks about occurrences of variables that are bound.+-- This keeps the occurrence checking code simple, because we never need+-- to compare variables to any variable in the bindings.+varInAst :: AstVarId -> AstTensor ms s y -> Bool+varInAst var = \case+  AstPair t1 t2 -> varInAst var t1 || varInAst var t2+  AstProject1 t -> varInAst var t+  AstProject2 t -> varInAst var t+  AstFromVector _ _ vl -> any (varInAst var) vl+  AstSum _ _ v -> varInAst var v+  AstReplicate _ _ v -> varInAst var v+  AstMapAccumRDer _k _bftk _eftk _f _df _rf acc0 es ->+    varInAst var acc0 || varInAst var es+  AstMapAccumLDer _k _bftk _eftk _f _df _rf acc0 es ->+    varInAst var acc0 || varInAst var es+  AstApply t ll -> varInAstHFun var t || varInAst var ll+  AstVar var2 -> var == varNameToAstVarId var2+  AstCond b v w -> varInAstBool var b || varInAst var v || varInAst var w+  AstBuild1 _ _ (var2, v) ->+    assert (varNameToAstVarId var2 /= var) $+    varInAst var v++  AstLet _ u v -> varInAst var u || varInAst var v+  AstShare _ v -> varInAst var v+  AstToShare v -> varInAst var v++  AstPrimalPart a -> varInAst var a+  AstDualPart a -> varInAst var a+  AstFromPrimal v -> varInAst var v+  AstFromDual v -> varInAst var v++  AstPlusK t u -> varInAst var t || varInAst var u+  AstTimesK t u -> varInAst var t || varInAst var u+  AstN1K _ t -> varInAst var t+  AstR1K _ t -> varInAst var t+  AstR2K _ t u -> varInAst var t || varInAst var u+  AstI2K _ t u -> varInAst var t || varInAst var u+  AstConcreteK{} -> False+  AstFloorK a -> varInAst var a+  AstFromIntegralK t -> varInAst var t+  AstCastK t -> varInAst var t++  AstPlusS t u -> varInAst var t || varInAst var u+  AstTimesS t u -> varInAst var t || varInAst var u+  AstN1S _ t -> varInAst var t+  AstR1S _ t -> varInAst var t+  AstR2S _ t u -> varInAst var t || varInAst var u+  AstI2S _ t u -> varInAst var t || varInAst var u+  AstConcreteS{} -> False+  AstFloorS a -> varInAst var a+  AstFromIntegralS a -> varInAst var a+  AstCastS t -> varInAst var t++  AstIndexS _ v ix -> varInAst var v || varInIxS var ix+  AstScatterS _ v (_vars, ix) -> varInIxS var ix || varInAst var v+  AstGatherS _ v (_vars, ix) -> varInIxS var ix || varInAst var v+  AstMinIndexS a -> varInAst var a+  AstMaxIndexS a -> varInAst var a+  AstIotaS{} -> False+  AstAppendS v u -> varInAst var v || varInAst var u+  AstSliceS _ _ _ v -> varInAst var v+  AstReverseS v -> varInAst var v+  AstTransposeS _perm v -> varInAst var v+  AstReshapeS _ v -> varInAst var v++  AstConvert _ v -> varInAst var v++  AstSum0S v -> varInAst var v+  AstDot0S u v -> varInAst var u || varInAst var v+  AstDot1InS _ _ u v -> varInAst var u || varInAst var v+  AstMatmul2S _ _ _ u v -> varInAst var u || varInAst var v++varInIxS :: AstVarId -> AstIxS ms sh -> Bool+varInIxS var = any (varInAst var)++varInAstHFun :: AstVarId -> AstHFun s s2 x y -> Bool+varInAstHFun _var AstLambda{} =+  False  -- we take advantage of the term being closed++varInAstBool :: AstVarId -> AstBool ms -> Bool+varInAstBool var = \case+  AstBoolConst{} -> False+  AstBoolNot b -> varInAstBool var b+  AstBoolAnd arg1 arg2 -> varInAstBool var arg1 || varInAstBool var arg2+  AstLeqK arg1 arg2 -> varInAst var arg1 || varInAst var arg2+  AstLeqS arg1 arg2 -> varInAst var arg1 || varInAst var arg2++varNameInAst :: AstVarName f y -> AstTensor ms s2 y2 -> Bool+varNameInAst var = varInAst (varNameToAstVarId var)++varNameInIxS :: AstVarName f y -> AstIxS ms sh -> Bool+varNameInIxS var = varInIxS (varNameToAstVarId var)+++-- * Determining if a term requires sharing++-- Turns off all but the most trivial cases of astIsSmall.+-- For tests only. Affects all simplification and inlining taking place+-- in parallel in the program at the time it's changed.+unsafeTotalSharingRef :: IORef Bool+{-# NOINLINE unsafeTotalSharingRef #-}+unsafeTotalSharingRef = unsafePerformIO $ newIORef False++setTotalSharing :: Bool -> IO ()+setTotalSharing b = atomicWriteIORef unsafeTotalSharingRef b++-- | A term requires sharing if it's too large as a term and so duplicating+-- it could affect the performance of simplification+-- or if it's too expensive when interpreted and so duplicating it+-- would increase the work done at runtime.+astIsSmall :: Bool -> AstTensor ms s y -> Bool+astIsSmall _ AstVar{} = True+astIsSmall _ AstShare{} = True+astIsSmall _ AstConcreteK{} = True+astIsSmall _ (AstConcreteS a) | sNatValue (Nested.srank a) == 0 = True+astIsSmall lax t = unsafePerformIO $ do+  unsafeTotalSharing <- readIORef unsafeTotalSharingRef+  return $! if | unsafeTotalSharing -> False+               | lax -> astIsSmallN 50 t > 0+               | otherwise -> astIsSmallN 20 t > 0++-- The cases with n <= 20 are usually good redex candidates,+-- so we expose them, but only if they are not burried too deeply.+-- Some of these constructors change tensor metadata into+-- a non-canonical form, which sometimes incurs the cost of converting+-- the vector to canonical form. The cost can be shared only+-- when the constructor is not the root of the shared term,+-- so when inlining (the False argument) we share them+-- unless they are at the root of the term tree.+astIsSmallN :: Int -> AstTensor ms s y -> Int+astIsSmallN n _ | n <= 0 = 0+astIsSmallN n t0 = case t0 of+  AstPair t1 t2 -> astIsSmallN (astIsSmallN (n - 1) t1) t2+  AstProject1 t -> astIsSmallN (n - 1) t+  AstProject2 t -> astIsSmallN (n - 1) t+  AstFromVector (SNat' @1) _ v -> astIsSmallN (n - 1) $ v V.! 0+  AstSum (SNat' @1) _ v -> astIsSmallN (n - 1) v+  AstReplicate _ _ v ->+    astIsSmallN (n - 1) v  -- a really good redex and often in series+      -- executed as a metadata change, which is however not free+  AstVar{} -> n+  AstCond b u v -> astIsSmallN (astIsSmallN (astBoolIsSmallN (n - 1) b) u) v+  AstConcreteK _ -> n+  AstConcreteS _ -> n  -- small term with zero interpretation cost;+                       -- the physical arrays is shared on GHC heap+  AstShare{} -> n+  AstPrimalPart v -> astIsSmallN (n - 1) v+  AstDualPart v -> astIsSmallN (n - 1) v+  AstFromPrimal v -> astIsSmallN (n - 1) v+  AstFromDual v -> astIsSmallN (n - 1) v++  AstIotaS{} -> n+  AstSliceS _ _ _ v ->+    if n <= 20 then 0 else astIsSmallN (n - 1) v  -- executed as metadata change+  AstReverseS v ->+    astIsSmallN (n - 1) v  -- executed as a cheap metadata change+  AstTransposeS _perm v ->+    if n <= 20 then 0 else astIsSmallN (n - 1) v  -- executed as metadata change++  AstConvert _ v -> astIsSmallN (n - 1) v++  _ -> 0++astBoolIsSmallN :: Int -> AstBool ms -> Int+astBoolIsSmallN n _ | n <= 0 = 0+astBoolIsSmallN n t0 = case t0 of+  AstBoolConst{} -> n+  AstBoolNot v -> astBoolIsSmallN (n - 1) v+  AstBoolAnd u v -> astBoolIsSmallN (astBoolIsSmallN (n - 1) u) v+  AstLeqK u v -> astIsSmallN (astIsSmallN (n - 1) u) v+  AstLeqS u v -> astIsSmallN (astIsSmallN (n - 1) u) v++ixIsSmall :: AstIxS ms sh -> Bool+ixIsSmall = all (astIsSmall True)++-- * Odds and ends++-- An approximation: lower and upper bound.+-- TODO: extend, e.g., to general quot and rem.+bounds :: GoodScalar r => AstTensor ms s (TKScalar r) -> (r, r)+bounds (AstConcreteK u) = (u, u)+bounds (AstVar var) = case varNameToBounds var of+  Nothing -> (-1000000000, 1000000000)+  Just (u1, u2) -> (fromIntegral u1, fromIntegral u2)+bounds (AstFromPrimal u) = bounds u+bounds (AstPrimalPart u) = bounds u+bounds (AstCond _b u v) = let (u1, u2) = bounds u+                              (v1, v2) = bounds v+                          in (min u1 v1, max u2 v2)+bounds (AstLet _ _ u) = bounds u  -- TODO: substitute?+bounds (AstPlusK u v) = let (u1, u2) = bounds u+                            (v1, v2) = bounds v+                        in (u1 + v1, u2 + v2)+bounds (AstN1K NegateOp u) = let (u1, u2) = bounds u in (- u2, - u1)+bounds (AstTimesK u v) =+  let (u1, u2) = bounds u+      (v1, v2) = bounds v+      l = [u1 * v1, u1 * v2, u2 * v1, u2 * v2]+  in (minimum l, maximum l)+bounds (AstI2K QuotOp u (AstConcreteK v)) | v > 0 =  -- a common case+  let (u1, u2) = bounds u+  in (u1 `quotH` v, u2 `quotH` v)+bounds (AstI2K RemOp u (AstConcreteK v)) | v > 0 =+  let (u1, u2) = bounds u+  in if | u1 >= 0 -> (0, min u2 (v - 1))  -- very crude+        | u2 <= 0 -> (max u1 (- v + 1), 0)+        | otherwise -> (- v + 1, v - 1)+bounds _ = (-1000000000, 1000000000)++liftRFromS1 :: forall n x ms s.+               (forall sh.+                   AstTensor ms s (TKS2 sh x)+                -> AstTensor ms s (TKS2 sh x))+            -> AstTensor ms s (TKR2 n x)+            -> AstTensor ms s (TKR2 n x)+liftRFromS1 f a = case ftkAst a of+  ftk@(FTKR sh' _) ->+    withShsFromShR sh' $ \(sh :: ShS sh) ->+      cAstFromS @(TKS2 sh x) ftk+      $ f (cAstSFromR @sh sh a)++liftRFromS2 :: forall n x ms s.+               (forall sh.+                   AstTensor ms s (TKS2 sh x) -> AstTensor ms s (TKS2 sh x)+                -> AstTensor ms s (TKS2 sh x))+            -> AstTensor ms s (TKR2 n x) -> AstTensor ms s (TKR2 n x)+            -> AstTensor ms s (TKR2 n x)+liftRFromS2 f a b  = case ftkAst a of+  ftk@(FTKR sh' _) ->+    withShsFromShR sh' $ \(sh :: ShS sh) ->+      cAstFromS @(TKS2 sh x) ftk+      $ f (cAstSFromR @sh sh a) (cAstSFromR @sh sh b)++liftXFromS1 :: forall sh' x ms s.+               (forall sh.+                   AstTensor ms s (TKS2 sh x)+                -> AstTensor ms s (TKS2 sh x))+            -> AstTensor ms s (TKX2 sh' x)+            -> AstTensor ms s (TKX2 sh' x)+liftXFromS1 f a = case ftkAst a of+  ftk@(FTKX sh' _) ->+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      cAstFromS @(TKS2 sh x) ftk+      $ f (cAstSFromX @sh @sh' sh a)++liftXFromS2 :: forall sh' x ms s.+               (forall sh.+                   AstTensor ms s (TKS2 sh x) -> AstTensor ms s (TKS2 sh x)+                -> AstTensor ms s (TKS2 sh x))+            -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKX2 sh' x)+            -> AstTensor ms s (TKX2 sh' x)+liftXFromS2 f a b = case ftkAst a of+  ftk@(FTKX sh' _) ->+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      cAstFromS @(TKS2 sh x) ftk+      $ f (cAstSFromX @sh @sh' sh a) (cAstSFromX @sh @sh' sh b)++cAstConvert :: TKConversion x z -> AstTensor ms s x -> AstTensor ms s z+cAstConvert c t+  | Just Refl <- matchingFTK (ftkAst t) (convertFTK c (ftkAst t)) = t+cAstConvert c1 (AstConvert c2 t2) = cAstConvert (ConvCmp c1 c2) t2+cAstConvert c1 (AstFromPrimal (AstConvert c2 t2)) =+  AstFromPrimal (cAstConvert (ConvCmp c1 c2) t2)+cAstConvert c t = AstConvert c t++cAstSFromR :: forall sh x ms s.+              ShS sh -> AstTensor ms s (TKR2 (Rank sh) x)+           -> AstTensor ms s (TKS2 sh x)+cAstSFromR sh v = case ftkAst v of+  FTKR _ x | Refl <- lemRankReplicate (Proxy @(Rank sh)) ->+    let c2 = ConvCmp (ConvXS' (FTKS sh x)) ConvRX+    in cAstConvert c2 v++cAstSFromX :: forall sh sh' x ms s. Rank sh ~ Rank sh'+           => ShS sh -> AstTensor ms s (TKX2 sh' x)+           -> AstTensor ms s (TKS2 sh x)+cAstSFromX sh v = case ftkAst v of+  FTKX _ x -> let c2 = ConvXS' (FTKS sh x)+              in cAstConvert c2 v++cAstXFromS :: forall sh sh' x ms s. Rank sh ~ Rank sh'+           => StaticShX sh' -> AstTensor ms s (TKS2 sh x)+           -> AstTensor ms s (TKX2 sh' x)+cAstXFromS ssx v+  | FTKS sh x <- ftkAst v+  , let shx = shCastSX ssx sh+  , Refl <- lemRankMapJust sh =+    let c2 = ConvCmp (ConvXX' (FTKX shx x)) ConvSX+    in cAstConvert c2 v++pattern AstFromS' :: forall {z1} {ms1} {s1}.+                     forall y z ms s. (z ~ z1, ms ~ ms1, s ~ s1)+                  => FullShapeTK z -> AstTensor ms s y+                  -> AstTensor ms1 s1 z1+pattern AstFromS' zftk a <-+  AstConvert c (checkPatternAstFromS c -> Just (zftk, a))++checkPatternAstFromS :: TKConversion y z -> AstTensor ms s y+                     -> Maybe (FullShapeTK z, AstTensor ms s y)+checkPatternAstFromS c t =+  let zftk = convertFTK c (ftkAst t)+  in if checkFtkAstFromS (ftkAst t) zftk then Just (zftk, t) else Nothing++checkAstFromS :: TKConversion a b -> AstTensor ms s a -> Bool+checkAstFromS c t = isJust $ checkPatternAstFromS c t++-- TODO: this is too lax, since it accepts nests/unnests. Add tests and fix.+checkFtkAstFromS :: FullShapeTK y -> FullShapeTK z -> Bool+checkFtkAstFromS yftk zftk | Just Refl <- matchingFTK yftk zftk = True+checkFtkAstFromS FTKS{} FTKS{} = False+checkFtkAstFromS FTKS{} _ = True+checkFtkAstFromS (FTKProduct yftk1 yftk2) (FTKProduct zftk1 zftk2) =+  checkFtkAstFromS yftk1 zftk1 && checkFtkAstFromS yftk2 zftk2+checkFtkAstFromS _ _ = False++checkFtkAstSFrom :: FullShapeTK y -> FullShapeTK z -> Bool+checkFtkAstSFrom yftk zftk | Just Refl <- matchingFTK yftk zftk = True+checkFtkAstSFrom FTKS{} FTKS{} = False+checkFtkAstSFrom _ FTKS{} = True+checkFtkAstSFrom (FTKProduct yftk1 yftk2) (FTKProduct zftk1 zftk2) =+  checkFtkAstSFrom yftk1 zftk1 && checkFtkAstSFrom yftk2 zftk2+checkFtkAstSFrom _ _ = False++cAstFromS :: forall y z ms s.+             FullShapeTK z -> AstTensor ms s y+          -> AstTensor ms s z+cAstFromS zftk t = cAstConvert (convFromS (ftkAst t) zftk) t++cAstSFrom :: forall y z ms s.+             FullShapeTK z -> AstTensor ms s y+          -> AstTensor ms s z+cAstSFrom zftk t = cAstConvert (convSFrom (ftkAst t) (ftkToSTK zftk)) t++convFromS :: FullShapeTK y0 -> FullShapeTK z0 -> TKConversion y0 z0+convFromS yftk0 zftk0 = case (yftk0, zftk0) of+  _ | Just Refl <- matchingFTK yftk0 zftk0 -> ConvId+  (FTKS ZSS (FTKScalar @ry), FTKScalar @rz)+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      ConvCmp ConvX0 ConvSX+  (FTKS sh x, FTKR rsh rx)+    | Just Refl <- matchingFTK x rx+    , Just Refl <- testEquality (shsRank sh) (shrRank rsh)+    , Refl <- lemRankMapJust sh ->+      ConvCmp (ConvXR (ftkToSTK x)) ConvSX+  (FTKS sh x, FTKX xsh xx)+    | Just Refl <- matchingFTK x xx+    , Just Refl <- testEquality (shsRank sh) (shxRank xsh)+    , Refl <- lemRankMapJust sh ->+      ConvCmp (ConvXX' zftk0) ConvSX+  (FTKProduct yftk1 yftk2, FTKProduct zftk1 zftk2) ->+    ConvT2 (convFromS yftk1 zftk1) (convFromS yftk2 zftk2)+  _ -> error $ "convFromS': unexpected types "  -- TODO: try nevertheless+               ++ "(" ++ show yftk0 ++ ", " ++ show zftk0 ++ ")"++convSFrom :: FullShapeTK y0 -> SingletonTK z0 -> TKConversion y0 z0+convSFrom yftk0 zstk0 = case (zstk0, yftk0) of+  _ | Just Refl <- sameSTK (ftkToSTK yftk0) zstk0 -> ConvId+  (STKS ZSS (STKScalar @ry), FTKScalar @rz)+    | Just Refl <- testEquality (typeRep @ry) (typeRep @rz) ->+      ConvCmp ConvXS (Conv0X STKScalar)+  (STKS @sh sh x, FTKR rsh rx)+    | Just Refl <- sameSTK x (ftkToSTK rx)+    , Just Refl <- testEquality (shsRank sh) (shrRank rsh)+    , Refl <- lemRankReplicate (Proxy @(Rank sh)) ->+      ConvCmp (ConvXS' (FTKS sh rx)) ConvRX+  (STKS sh x, FTKX xsh xx)+    | Just Refl <- sameSTK x (ftkToSTK xx)+    , Just Refl <- testEquality (shsRank sh) (shxRank xsh)+    , Refl <- lemRankMapJust sh ->+      ConvXS' (FTKS sh xx)+  (STKProduct zstk1 zstk2, FTKProduct yftk1 yftk2) ->+    ConvT2 (convSFrom yftk1 zstk1) (convSFrom yftk2 zstk2)+  _ -> error $ "convSFrom': unexpected types "  -- TODO: try nevertheless+               ++ "(" ++ show yftk0 ++ ", " ++ show zstk0 ++ ")"
+ src/HordeAd/Core/AstTraverse.hs view
@@ -0,0 +1,646 @@+-- | This module implements a few complete bottom-up simplifying passes+-- over any AST expression.+module HordeAd.Core.AstTraverse+  ( -- * The expansion (e.g., into gather expressions) bottom-up pass+    expandAst+    -- * The simplifying bottom-up pass+  , simplifyAst+    -- * The contraction (e.g., from gather expressions) bottom-up pass+  , contractAst+  ) where++import Prelude++import Data.Int (Int64)+import Data.Maybe (fromMaybe)+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import Foreign.C (CInt)+import GHC.TypeLits (KnownNat)+import Type.Reflection (typeRep)++import Data.Array.Nested (type (++))+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation (Perm (..))+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.Ast+  ( AstBool (AstBoolConst)+  , AstTensor (AstConcreteK, AstConcreteS, AstPlusK, AstPlusS, AstTimesK, AstTimesS)+  )+import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))+import HordeAd.Core.Ast qualified as Ast+import HordeAd.Core.AstSimplify+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * The expansion (e.g., into gather expressions) bottom-up pass++expandAstIxS :: AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh+expandAstIxS = fmap expandAst++-- | This pass expands terms, e.g., into @AstGather@ terms, in order+-- to expose redexes and enable fusion. It assumes that a contraction+-- pass follows that undoes some of the remaining expansion and applies+-- fusion rules that would be immediately counteracted by expansion rules+-- if applied earlier.+expandAst+  :: forall s y. AstSpan s+  => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y+expandAst t = case t of+  Ast.AstPair t1 t2 -> astPair (expandAst t1) (expandAst t2)+  Ast.AstProject1 v -> astProject1 (expandAst v)+  Ast.AstProject2 v -> astProject2 (expandAst v)+  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map expandAst l)+  Ast.AstSum snat stk v -> astSum snat stk (expandAst v)+  Ast.AstReplicate snat stk v -> astReplicate snat stk (expandAst v)+  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->+    astMapAccumRDer k bftk eftk+                    (expandAstHFun f)+                    (expandAstHFun df)+                    (expandAstHFun rf)+                    (expandAst acc0)+                    (expandAst es)+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    astMapAccumLDer k bftk eftk+                    (expandAstHFun f)+                    (expandAstHFun df)+                    (expandAstHFun rf)+                    (expandAst acc0)+                    (expandAst es)+  Ast.AstApply v ll -> astApply (expandAstHFun v) (expandAst ll)+  Ast.AstVar var -> astVar var+  Ast.AstCond b a2 a3 ->+    astCond (expandAstBool b) (expandAst a2) (expandAst a3)+  Ast.AstBuild1 k stk (var, v) ->+    let !v2 = expandAst v+    in Ast.AstBuild1 k stk (var, v2)++  Ast.AstLet var u v -> astLet var (expandAst u) (expandAst v)++  Ast.AstPrimalPart v -> astPrimalPart (expandAst v)+  Ast.AstDualPart v -> astDualPart (expandAst v)+  Ast.AstFromPrimal v -> Ast.AstFromPrimal (expandAst v)+  Ast.AstFromDual v -> Ast.AstFromDual (expandAst v)++  AstPlusK u v -> expandAst u + expandAst v+  AstTimesK u v -> expandAst u * expandAst v+  Ast.AstN1K NegateOp u -> negate (expandAst u)+  Ast.AstN1K AbsOp u -> abs (expandAst u)+  Ast.AstN1K SignumOp u -> signum (expandAst u)+  Ast.AstR1K opCode u -> Ast.AstR1K opCode (expandAst u)+  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (expandAst u) (expandAst v)+  Ast.AstI2K QuotOp u v -> quotH (expandAst u) (expandAst v)+  Ast.AstI2K RemOp u v -> remH (expandAst u) (expandAst v)+  AstConcreteK k -> AstConcreteK k+  Ast.AstFloorK a -> astFloorK (expandAst a)+  Ast.AstFromIntegralK v -> astFromIntegralK $ expandAst v+  Ast.AstCastK v -> astCastK $ expandAst v++  AstPlusS u v -> expandAst u + expandAst v+  AstTimesS u v -> expandAst u * expandAst v+  Ast.AstN1S NegateOp u -> negate (expandAst u)+  Ast.AstN1S AbsOp u -> abs (expandAst u)+  Ast.AstN1S SignumOp u -> signum (expandAst u)+  Ast.AstR1S opCode u -> Ast.AstR1S opCode (expandAst u)+  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (expandAst u) (expandAst v)+  Ast.AstI2S QuotOp u v -> quotH (expandAst u) (expandAst v)+  Ast.AstI2S RemOp u v -> remH (expandAst u) (expandAst v)+  AstConcreteS a -> AstConcreteS a+  Ast.AstFloorS a -> astFloorS (expandAst a)+  Ast.AstFromIntegralS v -> astFromIntegralS $ expandAst v+  Ast.AstCastS v -> astCastS $ expandAst v++  Ast.AstIndexS shn v ix ->+    astIndexKnobsS (defaultKnobs {knobPhase = PhaseExpansion})+                   shn (expandAst v) (expandAstIxS ix)+  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->+    astScatterS @shm @shn @shp shn (expandAst v) (vars, expandAstIxS ix)+  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->+    astGatherKnobsS @shm @shn @shp+                    (defaultKnobs {knobPhase = PhaseExpansion})+                    shn (expandAst v) (vars, expandAstIxS ix)+  Ast.AstMinIndexS a -> Ast.AstMinIndexS (expandAst a)+  Ast.AstMaxIndexS a -> Ast.AstMaxIndexS (expandAst a)+  Ast.AstIotaS{} -> t+  Ast.AstAppendS x y -> astAppendS (expandAst x) (expandAst y)+  Ast.AstSliceS i n k v -> astSliceS i n k (expandAst v)+  Ast.AstReverseS v -> astReverseS (expandAst v)+  Ast.AstTransposeS perm v -> {-+   -- disabled until we can reliably fuse back to transpose+   case expandAst v of+    Ast.AstVar{} -> t  -- normal form+    Ast.AstPrimalPart Ast.AstVar{} -> t  -- normal form+    Ast.AstDualPart Ast.AstVar{} -> t  -- normal form+    Ast.AstFromPrimal Ast.AstVar{} -> t  -- normal form+    Ast.AstFromDual Ast.AstVar{} -> t  -- normal form+    Ast.AstProject1{} -> t  -- normal form+    Ast.AstProject2{} -> t  -- normal form+    Ast.AstFromIntegralS{} -> t  -- normal form+    Ast.AstCastS{} -> t  -- normal form+    Ast.AstReplicate{} -> t  -- normal form+      -- TODO: this nf is silly, but right now transposes of replicates+      -- are small arrays and equivalent gathers are large terms and arrays,+      -- so this has to stay. Maybe we should contract gathers back+      -- to transposes of replicates (not only to replicates). Or maybe+      -- we should extend orthotope to any gather schemes, not only+      -- the simple ones.+      -- TODO: review also other nfs here and for AstReshapeS below+    Ast.AstScatterS _ _ (_, ix)+     | gcompare (Permutation.permRank perm) (ixsRank ix) == GGT -> t  -- nf+    v2 ->  -- not nf, let's express all as a gather+      astTransposeAsGatherS (defaultKnobs {knobExpand = True})+                            perm v2  -- TODO: (normalizePermutation perm)+        -- this is expensive but the only way to guarantee full simplification+    -} astTransposeS perm (expandAst v)+  Ast.AstReshapeS sh v -> {-  -- too hard to fuse back to reshape+   case expandAst v of+    Ast.AstVar{} -> t  -- normal form+    Ast.AstPrimalPart Ast.AstVar{} -> t  -- normal form+    Ast.AstDualPart Ast.AstVar{} -> t  -- normal form+    Ast.AstFromPrimal Ast.AstVar{} -> t  -- normal form+    Ast.AstFromDual Ast.AstVar{} -> t  -- normal form+    Ast.AstProject1{} -> t  -- normal form+    Ast.AstProject2{} -> t  -- normal form+    Ast.AstFromIntegralS{} -> t  -- normal form+    Ast.AstCastS{} -> t  -- normal form+    AstPlusS{} -> t  -- normal form+    AstTimesS{} -> t  -- normal form+    Ast.AstR2S{} -> t  -- normal form+    Ast.AstScatterS{} -> t  -- normal form+    v2 ->  -- not nf, let's express all as a gather+      astReshapeAsGatherS (defaultKnobs {knobExpand = True})+                          sh v2+        -- this is expensive but the only way to guarantee full simplification+    -} astReshapeS sh (expandAst v)++  Ast.AstConvert c v -> astConvert c $ expandAst v++  -- These should not appear in this context unless via wacky tests.+  Ast.AstSum0S{} -> t+  Ast.AstDot0S{} -> t+  Ast.AstDot1InS{} -> t+  Ast.AstMatmul2S{} -> t++expandAstHFun :: AstSpan s2+              => AstHFun s s2 x y -> AstHFun s s2 x y+expandAstHFun (AstLambda var l) = AstLambda var (expandAst l)++expandAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet+expandAstBool t = case t of+  AstBoolConst{} -> t+  Ast.AstBoolNot arg -> notB $ expandAstBool arg+  Ast.AstBoolAnd arg1 arg2 -> expandAstBool arg1 &&* expandAstBool arg2+  Ast.AstLeqK arg1 arg2 -> expandAst arg1 <=. expandAst arg2+  Ast.AstLeqS arg1 arg2 -> expandAst arg1 <=. expandAst arg2+++-- * The simplifying bottom-up pass++simplifyAstIxS :: AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh+simplifyAstIxS = fmap simplifyAst++-- | This function guarantees full simplification (unless redexes are obscured,+-- for which the expansion pass is sometimes a remedy): every redex+-- is visited and each combinator applied. The most exhaustive and costly+-- variants of each combinator are used, e.g., astIndexR.+simplifyAst+  :: forall s y. AstSpan s+  => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y+simplifyAst t = case t of+  Ast.AstPair t1 t2 -> astPair (simplifyAst t1) (simplifyAst t2)+  Ast.AstProject1 v -> astProject1 (simplifyAst v)+  Ast.AstProject2 v -> astProject2 (simplifyAst v)+  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map simplifyAst l)+  Ast.AstSum snat stk v -> astSum snat stk (simplifyAst v)+  Ast.AstReplicate snat stk v ->  astReplicate snat stk (simplifyAst v)+  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->+    astMapAccumRDer k bftk eftk+                    (simplifyAstHFun f)+                    (simplifyAstHFun df)+                    (simplifyAstHFun rf)+                    (simplifyAst acc0)+                    (simplifyAst es)+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    astMapAccumLDer k bftk eftk+                    (simplifyAstHFun f)+                    (simplifyAstHFun df)+                    (simplifyAstHFun rf)+                    (simplifyAst acc0)+                    (simplifyAst es)+  Ast.AstApply v ll -> astApply (simplifyAstHFun v) (simplifyAst ll)+  Ast.AstVar var -> astVar var+  Ast.AstCond b a2 a3 ->+    astCond (simplifyAstBool b) (simplifyAst a2) (simplifyAst a3)+  Ast.AstBuild1 k stk (var, v) ->+    let !v2 = simplifyAst v+    in Ast.AstBuild1 k stk (var, v2)++  Ast.AstLet var u v -> astLet var (simplifyAst u) (simplifyAst v)++  Ast.AstPrimalPart v -> astPrimalPart (simplifyAst v)+  Ast.AstDualPart v -> astDualPart (simplifyAst v)+  Ast.AstFromPrimal v -> Ast.AstFromPrimal (simplifyAst v)+  Ast.AstFromDual v -> Ast.AstFromDual (simplifyAst v)++  AstPlusK u v -> simplifyAst u + simplifyAst v+  AstTimesK u v -> simplifyAst u * simplifyAst v+  Ast.AstN1K NegateOp u -> negate (simplifyAst u)+  Ast.AstN1K AbsOp u -> abs (simplifyAst u)+  Ast.AstN1K SignumOp u -> signum (simplifyAst u)+  Ast.AstR1K opCode u -> Ast.AstR1K opCode (simplifyAst u)+  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (simplifyAst u) (simplifyAst v)+  Ast.AstI2K QuotOp u v -> quotH (simplifyAst u) (simplifyAst v)+  Ast.AstI2K RemOp u v -> remH (simplifyAst u) (simplifyAst v)+  AstConcreteK k -> AstConcreteK k+  Ast.AstFloorK a -> astFloorK (simplifyAst a)+  Ast.AstFromIntegralK v -> astFromIntegralK $ simplifyAst v+  Ast.AstCastK v -> astCastK $ simplifyAst v++  AstPlusS u v -> simplifyAst u + simplifyAst v+  AstTimesS u v -> simplifyAst u * simplifyAst v+  Ast.AstN1S NegateOp u -> negate (simplifyAst u)+  Ast.AstN1S AbsOp u -> abs (simplifyAst u)+  Ast.AstN1S SignumOp u -> signum (simplifyAst u)+  Ast.AstR1S opCode u -> Ast.AstR1S opCode (simplifyAst u)+  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (simplifyAst u) (simplifyAst v)+  Ast.AstI2S QuotOp u v -> quotH (simplifyAst u) (simplifyAst v)+  Ast.AstI2S RemOp u v -> remH (simplifyAst u) (simplifyAst v)+  AstConcreteS a -> AstConcreteS a+  Ast.AstFloorS a -> astFloorS (simplifyAst a)+  Ast.AstFromIntegralS v -> astFromIntegralS $ simplifyAst v+  Ast.AstCastS v -> astCastS $ simplifyAst v++  Ast.AstIndexS shn v ix ->+    astIndexKnobsS (defaultKnobs {knobPhase = PhaseSimplification})+                   shn (simplifyAst v) (simplifyAstIxS ix)+  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->+    astScatterS @shm @shn @shp shn (simplifyAst v) (vars, simplifyAstIxS ix)+  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->+    astGatherKnobsS @shm @shn @shp+                    (defaultKnobs {knobPhase = PhaseSimplification})+                    shn (simplifyAst v) (vars, simplifyAstIxS ix)+  Ast.AstMinIndexS a -> Ast.AstMinIndexS (simplifyAst a)+  Ast.AstMaxIndexS a -> Ast.AstMaxIndexS (simplifyAst a)+  Ast.AstIotaS{} -> t+  Ast.AstAppendS x y -> astAppendS (simplifyAst x) (simplifyAst y)+  Ast.AstSliceS i n k v -> astSliceS i n k (simplifyAst v)+  Ast.AstReverseS v -> astReverseS (simplifyAst v)+  Ast.AstTransposeS perm v -> astTransposeS perm $ simplifyAst v  -- TODO:(normalizePermutation perm)+  Ast.AstReshapeS sh v -> astReshapeS sh $ simplifyAst v++  Ast.AstConvert c v -> astConvert c $ simplifyAst v++  -- These should not appear in this context unless via wacky tests.+  Ast.AstSum0S{} -> t+  Ast.AstDot0S{} -> t+  Ast.AstDot1InS{} -> t+  Ast.AstMatmul2S{} -> t++simplifyAstHFun :: AstSpan s2+                => AstHFun s s2 x y -> AstHFun s s2 x y+simplifyAstHFun (AstLambda var l) = AstLambda var (simplifyAst l)++simplifyAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet+simplifyAstBool t = case t of+  AstBoolConst{} -> t+  Ast.AstBoolNot arg -> notB $ simplifyAstBool arg+  Ast.AstBoolAnd arg1 arg2 -> simplifyAstBool arg1 &&* simplifyAstBool arg2+  Ast.AstLeqK arg1 arg2 -> simplifyAst arg1 <=. simplifyAst arg2+  Ast.AstLeqS arg1 arg2 -> simplifyAst arg1 <=. simplifyAst arg2+++-- * The contraction (e.g., from gather expressions) bottom-up pass++contractAstIxS :: AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh+contractAstIxS = fmap contractAst++-- | When we have multiple backends, there should be one such pass+-- per backend that chooses a representation that is best for the backend.+-- The interpreter would interpret all of the backend-specific term+-- constructors, but the simplifier would ignore all and the user API+-- would not make them available.+--+-- Note that unlike all the other code in this module, this function+-- is not written in a compositional style nor close to it,+-- but it's instead defined in an ad-hoc way based on benchmarks.+contractAst+  :: forall s y. AstSpan s+  => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y+contractAst t0 = case t0 of+  Ast.AstPair t1 t2 -> astPair (contractAst t1) (contractAst t2)+  Ast.AstProject1 v -> astProject1 (contractAst v)+  Ast.AstProject2 v -> astProject2 (contractAst v)+  Ast.AstFromVector snat stk l -> astFromVector snat stk (V.map contractAst l)+  Ast.AstSum _ (STKS ZSS _) t2 -> astSum0S (contractAst t2)+  Ast.AstSum _ STKScalar t2 -> astKFromS' $ astSum0S (contractAst t2)+  Ast.AstSum+    snat@(SNat @m2)+    stk@(STKS (SNat @n2 :$$ SNat @p2 :$$ ZSS) STKScalar)+    v@(AstTimesS (Ast.AstTransposeS @permt permt+                    (Ast.AstReplicate (SNat @kt) (STKS @sht _ _) t2))+                 (Ast.AstTransposeS @permu permu+                    (Ast.AstReplicate (SNat @ku) (STKS @shu _ _) u2))) ->+    let perm10 = Permutation.makePerm @'[1, 0]+    in fromMaybe (astSum snat stk (contractAst v))+       $ case (permt, permu) of+      ( SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix permt (kt ': sht)+                      :~: [m2, n2, p2]) $+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix permu (ku ': shu)+                      :~: [m2, n2, p2]) $+        -- Sadly, the casts below, though implied by the permutations+        -- (as redundantly spelled out by the casts above) are required+        -- to make it type-check and they easily mask bugs, too.+        -- In the result, this is as type-unsafe as ranked code would be.+        gcastWith (unsafeCoerceRefl :: sht :~: [n2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, p2]) $+        attemptMatmul2 t2 u2+      ( SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, p2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [n2, m2]) $+        attemptMatmul2 u2 t2+      ( SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [n2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [p2, m2]) $+        attemptMatmul2 t2 (astTransposeS perm10 u2)+      ( SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [p2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [n2, m2]) $+        attemptMatmul2 u2 (astTransposeS perm10 t2)+      ( SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, n2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, p2]) $+        attemptMatmul2 (astTransposeS perm10 t2) u2+      ( SNat' @1 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, p2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, n2]) $+        attemptMatmul2 (astTransposeS perm10 u2) t2+      ( SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil+       ,SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [m2, n2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [p2, m2]) $+        attemptMatmul2 (astTransposeS perm10 t2)+                       (astTransposeS perm10 u2)+      ( SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` PNil+       ,SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` PNil ) ->+        gcastWith (unsafeCoerceRefl :: sht :~: [p2, m2]) $+        gcastWith (unsafeCoerceRefl :: shu :~: [m2, n2]) $+        attemptMatmul2 (astTransposeS perm10 u2)+                       (astTransposeS perm10 t2)+      _ -> Nothing+  Ast.AstSum n@(SNat @n) (STKS @sh sh _) (AstTimesS t2 u) ->+    let cpermR = backpermCycle $ 1 + shsLength sh+    in Permutation.permFromList cpermR $ \(cperm :: Permutation.Perm cperm) ->+         gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (n : sh)) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix cperm (n : sh)+                       :~: sh ++ '[n]) $+         fromMaybe (error "contractAst: impossible non-permutation")+         $ Permutation.permCheckPermutation cperm+         $ astDot1InS sh n (contractAst $ Ast.AstTransposeS cperm t2)+                           (contractAst $ Ast.AstTransposeS cperm u)+  Ast.AstSum snat stk (AstTimesS (Ast.AstLet vart vt t2) u) ->+    {- TODO: do we really need this check?+    | not (varNameInAst vart u) ->+        -- The varNameInAst check is needed, because although variable+        -- capture is impossible, because we don't create nested lets+        -- with the same variable, we could create such nested lets+        -- if we omitted this check.+    -}+    astLet vart+           (contractAst vt)+           (contractAst $ Ast.AstSum snat stk  -- the crucial exposed redex+                                     (AstTimesS t2 u))+  Ast.AstSum snat stk (AstTimesS t2 (Ast.AstLet varu vu u)) ->+    astLet varu+           (contractAst vu)+           (contractAst $ Ast.AstSum snat stk (AstTimesS t2 u))+  Ast.AstSum snat stk (Ast.AstLet var v t2) ->+    astLet var (contractAst v) (contractAst (Ast.AstSum snat stk t2))+  Ast.AstSum snat stk (Ast.AstSum snat2 stk2 (Ast.AstLet var v t2)) ->+    astLet var (contractAst v)+               (contractAst (Ast.AstSum snat stk (Ast.AstSum snat2 stk2 t2)))+  Ast.AstSum snat stk v -> astSum snat stk (contractAst v)+  Ast.AstReplicate snat stk v -> astReplicate snat stk (contractAst v)+  Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->+    astMapAccumRDer k bftk eftk+                    (contractAstHFun f)+                    (contractAstHFun df)+                    (contractAstHFun rf)+                    (contractAst acc0)+                    (contractAst es)+  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->+    astMapAccumLDer k bftk eftk+                    (contractAstHFun f)+                    (contractAstHFun df)+                    (contractAstHFun rf)+                    (contractAst acc0)+                    (contractAst es)+  Ast.AstApply v ll -> astApply (contractAstHFun v) (contractAst ll)+  Ast.AstVar var -> astVar var+  Ast.AstCond b a2 a3 ->+    astCond (contractAstBool b) (contractAst a2) (contractAst a3)+  -- These are only needed for tests that don't vectorize Ast.+  Ast.AstBuild1 snat (STKS ZSS _)  -- generalize+                (var, Ast.AstSum+                        n _+                        (AstTimesS+                           t2+                           (Ast.AstIndexS _shn+                              u (((:.$) @m (AstIntVar var2) ZIS)))))+    | Just Refl <- testEquality snat (SNat @m)+    , var == var2+    , not (varNameInAst var t2), not (varNameInAst var u) ->+        astDot1InS (snat :$$ ZSS) n+                   (contractAst u)+                   (contractAst+                    $ Ast.AstReplicate snat (ftkToSTK (ftkAst t2)) t2)+  Ast.AstBuild1 snat (STKS ZSS _)+                (var, Ast.AstSum _ _+                        (Ast.AstReshapeS+                           _sh (AstTimesS+                                  t2+                                  (Ast.AstIndexS _shn+                                     u (((:.$) @m (AstIntVar var2) ZIS))))))+    | ftk2@(FTKS (n :$$ ZSS) _) <- ftkAst t2+    , Just Refl <- testEquality snat (SNat @m)+    , var == var2+    , not (varNameInAst var t2), not (varNameInAst var u) ->+        astDot1InS (snat :$$ ZSS) n+                   (contractAst u)+                   (contractAst $ Ast.AstReplicate snat (ftkToSTK ftk2) t2)+  Ast.AstBuild1 k stk (var, v) ->+    let !v2 = contractAst v+    in Ast.AstBuild1 k stk (var, v2)++  Ast.AstLet var u v -> astLet var (contractAst u) (contractAst v)++  Ast.AstPrimalPart v -> astPrimalPart (contractAst v)+  Ast.AstDualPart v -> astDualPart (contractAst v)+  Ast.AstFromPrimal v -> Ast.AstFromPrimal (contractAst v)+  Ast.AstFromDual v -> Ast.AstFromDual (contractAst v)++  AstPlusK u v -> contractAst u + contractAst v+  AstTimesK u v -> contractAst u * contractAst v+  Ast.AstN1K NegateOp u -> negate (contractAst u)+  Ast.AstN1K AbsOp u -> abs (contractAst u)+  Ast.AstN1K SignumOp u -> signum (contractAst u)+  Ast.AstR1K RecipOp u -> recip (contractAst u)+  Ast.AstR1K ExpOp u -> exp (contractAst u)+  Ast.AstR1K LogOp u -> log (contractAst u)+  Ast.AstR1K SqrtOp u -> sqrt (contractAst u)+  Ast.AstR1K SinOp u -> sin (contractAst u)+  Ast.AstR1K CosOp u -> cos (contractAst u)+  Ast.AstR1K TanOp u -> tan (contractAst u)+  Ast.AstR1K AsinOp u -> asin (contractAst u)+  Ast.AstR1K AcosOp u -> acos (contractAst u)+  Ast.AstR1K AtanOp u -> atan (contractAst u)+  Ast.AstR1K SinhOp u -> sinh (contractAst u)+  Ast.AstR1K CoshOp u -> cosh (contractAst u)+  Ast.AstR1K TanhOp u -> tanh (contractAst u)+  Ast.AstR1K AsinhOp u -> asinh (contractAst u)+  Ast.AstR1K AcoshOp u -> acosh (contractAst u)+  Ast.AstR1K AtanhOp u -> atanh (contractAst u)+  Ast.AstR2K DivideOp u v -> contractAst u / contractAst v+  Ast.AstR2K PowerOp u v -> contractAst u ** contractAst v+  Ast.AstR2K LogBaseOp u v -> logBase (contractAst u) (contractAst v)+  Ast.AstR2K Atan2Op u v -> atan2H (contractAst u) (contractAst v)+  Ast.AstI2K QuotOp u v -> quotH (contractAst u) (contractAst v)+  Ast.AstI2K RemOp u v -> remH (contractAst u) (contractAst v)+  AstConcreteK k -> AstConcreteK k+  Ast.AstFloorK a -> astFloorK (contractAst a)+  Ast.AstFromIntegralK v -> astFromIntegralK $ contractAst v+  Ast.AstCastK v -> astCastK $ contractAst v++  AstPlusS u v -> contractAst u + contractAst v+  AstTimesS u v -> contractAst u * contractAst v+  Ast.AstN1S NegateOp u -> negate (contractAst u)+  Ast.AstN1S AbsOp u -> abs (contractAst u)+  Ast.AstN1S SignumOp u -> signum (contractAst u)+  Ast.AstR1S RecipOp u -> recip (contractAst u)+  Ast.AstR1S ExpOp u -> exp (contractAst u)+  Ast.AstR1S LogOp u -> log (contractAst u)+  Ast.AstR1S SqrtOp u -> sqrt (contractAst u)+  Ast.AstR1S SinOp u -> sin (contractAst u)+  Ast.AstR1S CosOp u -> cos (contractAst u)+  Ast.AstR1S TanOp u -> tan (contractAst u)+  Ast.AstR1S AsinOp u -> asin (contractAst u)+  Ast.AstR1S AcosOp u -> acos (contractAst u)+  Ast.AstR1S AtanOp u -> atan (contractAst u)+  Ast.AstR1S SinhOp u -> sinh (contractAst u)+  Ast.AstR1S CoshOp u -> cosh (contractAst u)+  Ast.AstR1S TanhOp u -> tanh (contractAst u)+  Ast.AstR1S AsinhOp u -> asinh (contractAst u)+  Ast.AstR1S AcoshOp u -> acosh (contractAst u)+  Ast.AstR1S AtanhOp u -> atanh (contractAst u)+  Ast.AstR2S DivideOp u v -> contractAst u / contractAst v+  Ast.AstR2S PowerOp u v -> contractAst u ** contractAst v+  Ast.AstR2S LogBaseOp u v -> logBase (contractAst u) (contractAst v)+  Ast.AstR2S Atan2Op u v -> atan2H (contractAst u) (contractAst v)+  Ast.AstI2S QuotOp u v -> quotH (contractAst u) (contractAst v)+  Ast.AstI2S RemOp u v -> remH (contractAst u) (contractAst v)+  AstConcreteS a -> AstConcreteS a+  Ast.AstFloorS t -> case contractAst t of+    AstConcreteS a -> astConcreteS (tsfloor $ Concrete a)+    t2 -> astFloorS t2+  Ast.AstFromIntegralS t -> case contractAst t of+    AstConcreteS a -> astConcreteS (tsfromIntegral $ Concrete a)+    t2 -> astFromIntegralS t2+  Ast.AstCastS t -> case contractAst t of+    AstConcreteS a -> astConcreteS (tscast $ Concrete a)+    t2 -> astCastS t2++  Ast.AstIndexS shn v ix ->+    astIndexKnobsS (defaultKnobs {knobPhase = PhaseContraction})+                   shn (contractAst v) (contractAstIxS ix)+  Ast.AstScatterS @shm @shn @shp shn v (vars, ix) ->+    astScatterS @shm @shn @shp shn (contractAst v) (vars, contractAstIxS ix)+  -- This rule is reverted in vectorization, so contraction phase may be fine.+  Ast.AstGatherS shn v (vars, Ast.AstCond b i1 i2 :.$ prest)+    | not $ any ((`varInAstBool` b) . varNameToAstVarId) (listsToList vars) ->+      contractAst+      $ Ast.AstCond b+                    (Ast.AstGatherS shn v (vars, i1 :.$ prest))+                    (Ast.AstGatherS shn v (vars, i2 :.$ prest))+  Ast.AstGatherS @shm @shn @shp shn v (vars, ix) ->+    astGatherKnobsS @shm @shn @shp+                    (defaultKnobs {knobPhase = PhaseContraction})+                    shn (contractAst v) (vars, contractAstIxS ix)+{- TODO, but sbuild is tricky, so only if benchmarks show it's worth it:+  AstGatherS @shm AstIotaS (vars, i :.$ ZIS) | Refl <- lemAppNil @shm ->+    gcastWith (unsafeCoerceRefl :: Drop (Rank shm) shm :~: '[]) $+    gcastWith (unsafeCoerceRefl :: Take (Rank shm) shm :~: shm) $+    sbuild @_ @_ @(Rank shm)+           (interpretLambdaIndexS+              interpretAst env+              (vars, fromPrimal @s $ AstFromIntegralS $ AstSFromK i)) -}+  Ast.AstMinIndexS a -> Ast.AstMinIndexS (contractAst a)+  Ast.AstMaxIndexS a -> Ast.AstMaxIndexS (contractAst a)+  Ast.AstIotaS (SNat @n) -> astConcreteS $ tsiota @_ @n+  Ast.AstAppendS x y -> astAppendS (contractAst x) (contractAst y)+  Ast.AstSliceS i n k t -> astSliceS i n k (contractAst t)+  Ast.AstReverseS t -> astReverseS (contractAst t)+  Ast.AstTransposeS perm v -> astTransposeS perm $ contractAst v  -- TODO:(normalizePermutation perm)+  Ast.AstReshapeS sh2 t -> case contractAst t of+    AstConcreteS v -> astConcreteS (tsreshape sh2 $ Concrete v)+    t2 -> astReshapeS sh2 t2++  Ast.AstConvert c v -> astConvert c $ contractAst v++  -- These should not appear in this context unless via wacky tests.+  Ast.AstSum0S{} -> t0+  Ast.AstDot0S{} -> t0+  Ast.AstDot1InS{} -> t0+  Ast.AstMatmul2S{} -> t0++attemptMatmul2+  :: forall m n p r s.+     (KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s)+  => AstTensor AstMethodLet s (TKS '[m, n] r)+  -> AstTensor AstMethodLet s (TKS '[n, p] r)+  -> Maybe (AstTensor AstMethodLet s (TKS '[m, p] r))+attemptMatmul2 t3 u3 = Just $+  let t4 = contractAst t3+      u4 = contractAst u3+  in case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl ->+      astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl ->+        astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+      _ -> case testEquality (typeRep @r) (typeRep @Int64) of+        Just Refl ->+          astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+        _ -> case testEquality (typeRep @r) (typeRep @CInt) of+          Just Refl ->+            astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+          _ -> case testEquality (typeRep @r) (typeRep @Z1) of+            Just Refl ->+              astMatmul2S (SNat @m) (SNat @n) (SNat @p) t4 u4+            _ -> error "attemptMatmul2: unexpected scalar"++contractAstHFun :: AstSpan s2+                => AstHFun s s2 x y -> AstHFun s s2 x y+contractAstHFun (AstLambda var l) = AstLambda var (contractAst l)++contractAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet+contractAstBool t = case t of+  AstBoolConst{} -> t+  Ast.AstBoolNot arg -> notB $ contractAstBool arg+  Ast.AstBoolAnd arg1 arg2 -> contractAstBool arg1 &&* contractAstBool arg2+  Ast.AstLeqK arg1 arg2 -> contractAst arg1 <=. contractAst arg2+  Ast.AstLeqS arg1 arg2 -> contractAst arg1 <=. contractAst arg2
+ src/HordeAd/Core/AstVectorize.hs view
@@ -0,0 +1,571 @@+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | BOT (bulk-operation transformation) of the AST, which is a kind+-- of vectorization. It eliminates the build operation and, consequently,+-- any occurrence of indexing under build, which would cause delta expression+-- explosion and afterwards one-hot explosion when evaluating deltas.+module HordeAd.Core.AstVectorize+  ( build1Vectorize, traceRuleEnabledRef+  ) where++import Prelude++import Control.Exception.Assert.Sugar+import Control.Monad (when)+import Data.Functor.Const+import Data.IORef+import Data.Maybe (fromMaybe)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import GHC.TypeLits (type (+), type (<=?))+import System.IO (Handle, hFlush, hPutStrLn, stderr, stdout)+import System.IO.Unsafe (unsafePerformIO)++import Data.Array.Nested (type (++))+import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (Tail, unsafeCoerceRefl)++import HordeAd.AstEngine+import HordeAd.Core.Ast (AstTensor)+import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))+import HordeAd.Core.Ast qualified as Ast+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstSimplify+import HordeAd.Core.AstTools+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * The top-level wrapper++-- | The application @build1Vectorize k (var, v)@ vectorizes+-- the term @AstBuild1 k (var, v)@, that is, rewrites it to a term+-- with the same value, but not containing the outermost @AstBuild1@+-- constructor and not increasing (and potentially decreasing)+-- the total number of @AstBuild1@ occuring in the term.+-- If no @AstBuild1@ terms occur in @v@, the resulting term won't+-- have any, either.+build1Vectorize+  :: forall y k s. AstSpan s+  => SNat k -> SingletonTK y -> (IntVarName, AstTensor AstMethodLet s y)+  -> AstTensor AstMethodLet s (BuildTensorKind k y)+{-# NOINLINE build1Vectorize #-}+build1Vectorize snat@SNat stk (!var, !v0) = unsafePerformIO $ do+  enabled <- readIORef traceRuleEnabledRef+  let width = 1000 * traceWidth+      startTerm = Ast.AstBuild1 snat stk (var, v0)+  when enabled $ do+    writeIORef traceNestingLevel 0+    hPutStrLnFlush stderr $+      "\n"+      ++ "START of vectorization for term "+      ++ ellipsisString width (printAstSimple startTerm)+      ++ "\n"+  let !endTerm = build1VOccurrenceUnknown snat (var, v0)+  when enabled $ do+    hPutStrLnFlush stderr $+      "\n"+      ++ "END of vectorization yields "+      ++ ellipsisString width (printAstSimple endTerm)+      ++ "\n"+  let !_A = assert (ftkAst startTerm == ftkAst endTerm+                   `blame` "build1Vectorize: term shape changed"+                   `swith`( ftkAst startTerm+                          , ftkAst endTerm) ) ()+  return endTerm+++-- * Vectorization++-- | The application @build1VOccurrenceUnknown k (var, v)@ vectorizes+-- term @AstBuild1 k (var, v)@, where it's unknown whether+-- @var@ occurs in @v@.+build1VOccurrenceUnknown+  :: forall y k s. AstSpan s+  => SNat k -> (IntVarName, AstTensor AstMethodLet s y)+  -> AstTensor AstMethodLet s (BuildTensorKind k y)+build1VOccurrenceUnknown snat@SNat (!var, !v0)+  | stk0 <- ftkToSTK (ftkAst v0) =+    let traceRule =+          mkTraceRule "build1VOcc" (Ast.AstBuild1 snat stk0 (var, v0)) v0 1+    in if varNameInAst var v0+       then build1V snat (var, v0)+       else traceRule $+         astReplicate snat stk0 v0++-- This refreshes the indexing variable in a build body.+-- This is used to avoid biding the same variable twice in the code,+-- (unless in very safe situations, e.g., different branches+-- of an arithmetic expression) which may end up as nested bindings eventually+-- and break our invariants that we need for simplified handling of bindings+-- when rewriting terms.+build1VOccurrenceUnknownRefresh+  :: forall y k s. AstSpan s+  => SNat k -> (IntVarName, AstTensor AstMethodLet s y)+  -> AstTensor AstMethodLet s (BuildTensorKind k y)+{-# NOINLINE build1VOccurrenceUnknownRefresh #-}+build1VOccurrenceUnknownRefresh snat@SNat (!var, !v0) =+  funToAstIntVar (varNameToBounds var) $ \ (!varFresh, !astVarFresh) ->+    let !v2 = substituteAst astVarFresh var v0+                -- cheap subst, because only a renaming+    in build1VOccurrenceUnknown snat (varFresh, v2)++-- | The application @build1V k (var, v)@ vectorizes term+-- @AstBuild1 k (var, v)@, where it's known that- @var@ occurs in @v@,+-- see above for what it means precisely.+build1V+  :: forall y k s. AstSpan s+  => SNat k -> (IntVarName, AstTensor AstMethodLet s y)+  -> AstTensor AstMethodLet s (BuildTensorKind k y)+build1V snat@SNat (!var, !v0) | ftk0 <- ftkAst v0 =+  let bv = Ast.AstBuild1 snat (ftkToSTK ftk0) (var, v0)+      traceRule = mkTraceRule "build1V" bv v0 1+  in case v0 of+    Ast.AstPair t1 t2 -> traceRule $+      astPair (build1VOccurrenceUnknown snat (var, t1))+              (build1VOccurrenceUnknown snat (var, t2))+    Ast.AstProject1 t -> traceRule $+      astProject1 (build1V snat (var, t))+    Ast.AstProject2 t -> traceRule $+      astProject2 (build1V snat (var, t))+    Ast.AstFromVector snat1@(SNat @k1) stk l -> traceRule $+      astTrBuild (SNat @k1) (SNat @k) stk+      $ astFromVector snat1 (buildSTK snat stk) (V.map (\v ->+          build1VOccurrenceUnknown snat (var, v)) l)+    Ast.AstSum (SNat @k1) stk v -> traceRule $+      astSum (SNat @k1) (buildSTK snat stk)+      $ astTrBuild (SNat @k) (SNat @k1) stk $ build1V snat (var, v)+    Ast.AstReplicate snat2@(SNat @k2) stk v -> traceRule $+      astTrBuild (SNat @k2) SNat stk+      $ astReplicate snat2 (buildSTK snat stk) $ build1V snat (var, v)+    Ast.AstMapAccumRDer k5@(SNat @k5) bftk eftk f df rf acc0 es+     | Refl <- lemBuildOfAD snat (ftkToSTK (ftkAst acc0))+     , Refl <- lemBuildOfAD snat (ftkToSTK bftk)+     , Refl <- lemBuildOfAD snat (ftkToSTK eftk) -> traceRule $+      astLetFun+        (astMapAccumRDer+           k5+           (buildFTK snat bftk)+           (buildFTK snat eftk)+           (build1VHFun snat (var, f))+           (build1VHFun snat (var, df))+           (build1VHFun snat (var, rf))+           (build1VOccurrenceUnknown snat (var, acc0))+           (astTrBuild (SNat @k) (SNat @k5) (ftkToSTK eftk)+            $ build1VOccurrenceUnknown snat (var, es)))+        (\x1bs1 -> astPair (astProject1 x1bs1)+                           (astTrBuild (SNat @k5) (SNat @k)+                                       (ftkToSTK bftk) (astProject2 x1bs1)))+    Ast.AstMapAccumLDer k5@(SNat @k5) bftk eftk f df rf acc0 es+     | Refl <- lemBuildOfAD snat (ftkToSTK (ftkAst acc0))+     , Refl <- lemBuildOfAD snat (ftkToSTK bftk)+     , Refl <- lemBuildOfAD snat (ftkToSTK eftk) -> traceRule $+      astLetFun+        (astMapAccumLDer+           k5+           (buildFTK snat bftk)+           (buildFTK snat eftk)+           (build1VHFun snat (var, f))+           (build1VHFun snat (var, df))+           (build1VHFun snat (var, rf))+           (build1VOccurrenceUnknown snat (var, acc0))+           (astTrBuild (SNat @k) (SNat @k5) (ftkToSTK eftk)+            $ build1VOccurrenceUnknown snat (var, es)))+        (\x1bs1 -> astPair (astProject1 x1bs1)+                           (astTrBuild (SNat @k5) (SNat @k)+                                       (ftkToSTK bftk) (astProject2 x1bs1)))+    Ast.AstApply t ll -> traceRule $+      astApply (build1VHFun snat (var, t))+               (build1VOccurrenceUnknown snat (var, ll))+    Ast.AstVar var2 -> traceRule $+      if varNameToAstVarId var2 == varNameToAstVarId var+      then case isTensorInt (Proxy @s) (varNameToFTK var2) of+        Just Refl -> fromPrimal @s $ Ast.AstIotaS (SNat @k)+        _ -> error "build1V: build variable is not an index variable"+      else error "build1V: AstVar can't contain other free variables"+    Ast.AstCond b u v -> traceRule $+      let uv = astFromVector (SNat @2) (ftkToSTK (ftkAst u)) (V.fromList [u, v])+          t = astIndexBuild (SNat @2) ftk0 uv (astCond b 0 1)+      in build1VOccurrenceUnknown snat (var, t)+    Ast.AstBuild1 snat2 _ (var2, v2) -> traceRule $+      assert (var2 /= var) $+      build1VOccurrenceUnknown+        snat (var, build1VOccurrenceUnknown snat2 (var2, v2))+          -- happens only when testing and mixing different pipelines++    Ast.AstLet var1 u v -> traceRule $+      let (var3, v2) = substProjRep snat var var1 v+      in astLet var3 (build1VOccurrenceUnknown snat (var, u))+                     (build1VOccurrenceUnknownRefresh snat (var, v2))+           -- ensures no duplicated bindings, see below++    Ast.AstPrimalPart v -> traceRule $+      astPrimalPart $ build1V snat (var, v)+    Ast.AstDualPart v -> traceRule $+      astDualPart $ build1V snat (var, v)+    Ast.AstFromPrimal v -> traceRule $+      Ast.AstFromPrimal $ build1V snat (var, v)+    Ast.AstFromDual v -> traceRule $+      Ast.AstFromDual $ build1V snat (var, v)++    Ast.AstPlusK u v -> traceRule $+      build1VOccurrenceUnknown snat (var, u)+      + build1VOccurrenceUnknown snat (var, v)+    Ast.AstTimesK u v -> traceRule $+      build1VOccurrenceUnknown snat (var, u)+      * build1VOccurrenceUnknown snat (var, v)+        -- we permit duplicated bindings, because they can't easily+        -- be substituted into one another unlike. e.g., inside a let,+        -- which may get inlined+    Ast.AstN1K opCode u -> traceRule $+      Ast.AstN1S opCode (build1V snat (var, u))+    Ast.AstR1K opCode u -> traceRule $+      Ast.AstR1S opCode (build1V snat (var, u))+    Ast.AstR2K opCode u v -> traceRule $+      Ast.AstR2S opCode (build1VOccurrenceUnknown snat (var, u))+                        (build1VOccurrenceUnknown snat (var, v))+    Ast.AstI2K opCode u v -> traceRule $+      Ast.AstI2S opCode (build1VOccurrenceUnknown snat (var, u))+                        (build1VOccurrenceUnknown snat (var, v))+    Ast.AstConcreteK{} ->+      error "build1V: AstConcreteK can't have free index variables"+    Ast.AstFloorK v -> traceRule $+      astFloorS $ build1V snat (var, v)+    Ast.AstFromIntegralK v -> traceRule $+      astFromIntegralS $ build1V snat (var, v)+    Ast.AstCastK v -> traceRule $+      astCastS $ build1V snat (var, v)++    Ast.AstPlusS u v -> traceRule $+      build1VOccurrenceUnknown snat (var, u)+      + build1VOccurrenceUnknown snat (var, v)+    Ast.AstTimesS u v -> traceRule $+      build1VOccurrenceUnknown snat (var, u)+      * build1VOccurrenceUnknown snat (var, v)+        -- we permit duplicated bindings, because they can't easily+        -- be substituted into one another unlike. e.g., inside a let,+        -- which may get inlined+    Ast.AstN1S opCode u -> traceRule $+      Ast.AstN1S opCode (build1V snat (var, u))+    Ast.AstR1S opCode u -> traceRule $+      Ast.AstR1S opCode (build1V snat (var, u))+    Ast.AstR2S opCode u v -> traceRule $+      Ast.AstR2S opCode (build1VOccurrenceUnknown snat (var, u))+                        (build1VOccurrenceUnknown snat (var, v))+    Ast.AstI2S opCode u v -> traceRule $+      Ast.AstI2S opCode (build1VOccurrenceUnknown snat (var, u))+                        (build1VOccurrenceUnknown snat (var, v))+    Ast.AstConcreteS{} ->+      error "build1V: AstConcreteS can't have free index variables"+    Ast.AstFloorS v -> traceRule $+      astFloorS $ build1V snat (var, v)+    Ast.AstFromIntegralS v -> traceRule $+      astFromIntegralS $ build1V snat (var, v)+    Ast.AstCastS v -> traceRule $+      astCastS $ build1V snat (var, v)++    Ast.AstIndexS shn v ix -> traceRule $+      build1VIndexS snat shn (var, v, ix)  -- @var@ is in @v@ or @ix@+    Ast.AstScatterS @shm @shn @shp shn v (vars, ix) -> traceRule $+      let (varFresh, astVarFresh, ix2) = intBindingRefreshS (var, ix)+      in astScatterS @(k ': shm) @shn @(k ': shp) shn+                     (build1VOccurrenceUnknown snat (var, v))+                     (Const varFresh ::$ vars, astVarFresh :.$ ix2)+    Ast.AstGatherS @shm @shn @shp shn v (vars, ix) -> traceRule $+      let (varFresh, astVarFresh, ix2) = intBindingRefreshS (var, ix)+      in astGatherS @(k ': shm) @shn @(k ': shp) shn+                    (build1VOccurrenceUnknown snat (var, v))+                    (Const varFresh ::$ vars, astVarFresh :.$ ix2)+    Ast.AstMinIndexS v -> traceRule $+      Ast.AstMinIndexS $ build1V snat (var, v)+    Ast.AstMaxIndexS v -> traceRule $+      Ast.AstMaxIndexS $ build1V snat (var, v)+    Ast.AstIotaS{} ->+      error "build1V: AstIotaS can't have free variables"+    Ast.AstAppendS v w -> traceRule $+      astTrS $ astAppendS (astTrS $ build1VOccurrenceUnknown snat (var, v))+                          (astTrS $ build1VOccurrenceUnknown snat (var, w))+    Ast.AstSliceS i n k v -> traceRule $+      astTrS $ astSliceS i n k $ astTrS $ build1V snat (var, v)+    Ast.AstReverseS v -> traceRule $+      astTrS $ astReverseS $ astTrS $ build1V snat (var, v)+    Ast.AstTransposeS @perm @sh1 perm v -> traceRule $ case ftk0 of+      FTKS @sh _ _ ->+        let zsuccPerm :: Permutation.Perm (0 : Permutation.MapSucc perm)+            zsuccPerm = Permutation.permShift1 perm+        in gcastWith (unsafeCoerceRefl+                      :: Permutation.PermutePrefix+                           (0 : Permutation.MapSucc perm) (k : sh1)+                         :~: k : sh) $+           gcastWith (unsafeCoerceRefl+                      :: Rank (0 : Permutation.MapSucc perm)+                         :~: 1 + Rank perm) $+           fromMaybe (error "build1V: impossible non-permutation")+           $ Permutation.permCheckPermutation zsuccPerm+           $ astTransposeS zsuccPerm $ build1V snat (var, v)+    Ast.AstReshapeS sh v -> traceRule $+      astReshapeS (snat :$$ sh) $ build1V snat (var, v)++    Ast.AstConvert c v -> traceRule $+      astConvert (buildTKConversion snat (ftkAst v) c)+      $ build1V snat (var, v)++    Ast.AstSum0S{} -> error "build1V: term not accessible from user API"+    Ast.AstDot0S{} -> error "build1V: term not accessible from user API"+    Ast.AstDot1InS{} -> error "build1V: term not accessible from user API"+    Ast.AstMatmul2S{} -> error "build1V: term not accessible from user API"++-- This refreshes an index variable in a list of index expressions.+intBindingRefreshS+  :: (IntVarName, AstIxS AstMethodLet sh)+  -> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh)+{-# NOINLINE intBindingRefreshS #-}+intBindingRefreshS (!var, !ix) =+  funToAstIntVar (varNameToBounds var) $ \ (!varFresh, !astVarFresh) ->+    let !ix2 = substituteAstIxS astVarFresh var ix+                 -- cheap subst, because only a renaming+    in (varFresh, astVarFresh, ix2)++-- | The application @build1VIndex snat (var, v, ix)@ vectorizes+-- the term @AstBuild1 snat (var, AstIndex v ix)@, where it's unknown whether+-- @var@ occurs in any of @v@, @ix@.+--+-- We try to push indexing down as far as needed to eliminate any occurrences+-- of @var@ from @v@ (but not necessarily from @ix@), which is enough+-- to replace @AstBuild1@ with @AstGather@ and so complete+-- the vectorization.+--+-- This pushing down is performed by alternating steps of simplification,+-- in @astIndex@, that eliminates indexing from the top of a term+-- position (except for two permissible normal forms) and vectorization,+-- @build1VOccurrenceUnknown@, that recursively goes down under constructors+-- until it encounter indexing again. We have to do this in lockstep+-- so that we simplify terms only as much as needed to vectorize.+--+-- If simplification can't proceed, which means that we reached one of the few+-- normal forms wrt simplification, we invoke the pure desperation rule (D)+-- which produces large tensors, which are hard to simplify even when+-- eventually proven unnecessary. The rule changes the index to a gather+-- and pushes the build down the gather, getting the vectorization unstuck.+build1VIndexS+  :: forall k shm shn x s. AstSpan s+  => SNat k -> ShS shn+  -> ( IntVarName  -- bounds (0, k - 1)+     , AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)+     , AstIxS AstMethodLet shm )+  -> AstTensor AstMethodLet s (TKS2 (k ': shn) x)+build1VIndexS k _ (!var, !v0, ZIS) =+  build1VOccurrenceUnknown k (var, v0)+build1VIndexS k@SNat shn (var, v0, ix) | STKS _ x <- ftkToSTK (ftkAst v0) =+  let vTrace = Ast.AstBuild1 k (STKS shn x) (var, Ast.AstIndexS shn v0 ix)+      traceRule = mkTraceRule "build1VIndexS" vTrace v0 1+  in if varNameInAst var v0+     then case astIndexKnobsS (defaultKnobs {knobPhase = PhaseVectorization})+                              shn v0 ix of  -- push deeper+       Ast.AstIndexS _ v1 ZIS -> traceRule $+         build1VOccurrenceUnknown k (var, v1)+       v@(Ast.AstIndexS shn1 v1 ix1) -> traceRule $+         let (varFresh, astVarFresh, ix2) = intBindingRefreshS (var, ix1)+             ruleD :: AstTensor AstMethodLet s (TKS2 (k ': shn) x)+             ruleD = astGatherS+                       shn1 (build1VOccurrenceUnknown k (var, v1))+                       (Const varFresh ::$ ZS, astVarFresh :.$ ix2)+             len = ixsLength ix1+         in if varNameInAst var v1+            then case v1 of  -- try to avoid ruleD if not a normal form+              Ast.AstFromVector{} | len == 1 -> ruleD+              Ast.AstScatterS{} -> ruleD+              Ast.AstGatherS{} -> ruleD+              -- These can only be simplified to the AstFromVector NF above.+              Ast.AstReplicate{} -> ruleD+              Ast.AstAppendS{} -> ruleD+              -- These, in general, simplify to gathers, so as bad as ruleD.+              Ast.AstTransposeS{} -> ruleD+              Ast.AstReshapeS{} -> ruleD+              -- Rarely these don't simplify enough; left as an escape hatch:+              -- (TODO: simplify better)+              Ast.AstConvert{} -> ruleD+              _ -> build1VOccurrenceUnknown k (var, v)  -- not a normal form+            else build1VOccurrenceUnknown k (var, v)  -- shortcut+       v -> traceRule $+         build1VOccurrenceUnknown k (var, v)+           -- peel off yet another constructor+     else traceRule $+            astGatherS shn v0 (Const var ::$ ZS, ix)++build1VHFun+  :: forall k x z s s2. (AstSpan s, AstSpan s2)+  => SNat k -> (IntVarName, AstHFun s s2 x z)+  -> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)+build1VHFun snat@SNat (!var, !v0) = case v0 of+  Ast.AstLambda var1 t ->+    -- This handles the case of l having free variables beyond var1,+    -- which is not possible for lambdas used in folds, etc.+    -- But note that, due to substProjVars, l2 has var occurrences,+    -- so build1VOccurrenceUnknownRefresh is neccessary to handle+    -- them and to eliminate them so that the function is closed again.+    let (var2, t2) = substProjRep snat var var1 t+    in Ast.AstLambda var2 (build1VOccurrenceUnknownRefresh snat (var, t2))+++-- * Auxiliary operations++astTr :: forall n s r. AstSpan s+      => AstTensor AstMethodLet s (TKR2 (2 + n) r)+      -> AstTensor AstMethodLet s (TKR2 (2 + n) r)+astTr a = case Permutation.makePerm @'[1, 0] of+  (perm :: Permutation.Perm perm) -> case ftkAst a of+    FTKR sh'@(k :$: m :$: shr) x | SNat <- shrRank sh' ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        gcastWith (unsafeCoerceRefl :: (Rank perm <=? Rank sh) :~: True) $+        astFromS' (FTKR (m :$: k :$: shr) x)+        . astTransposeS perm . astSFromR' sh $ a+    _ -> error "astTr: impossible shape"++astTrS :: forall n m sh s r. AstSpan s+       => AstTensor AstMethodLet s (TKS2 (n ': m ': sh) r)+       -> AstTensor AstMethodLet s (TKS2 (m ': n ': sh) r)+astTrS a | FTKS (_ :$$ _ :$$ sh) _ <- ftkAst a+         , SNat <- shsRank sh =  -- why on Earth is this needed?+  astTransposeS (Permutation.makePerm @'[1, 0]) a++astTrX :: forall n m shx s r. AstSpan s+       => AstTensor AstMethodLet s (TKX2 (Just n ': Just m ': shx) r)+       -> AstTensor AstMethodLet s (TKX2 (Just m ': Just n ': shx) r)+astTrX a = case Permutation.makePerm @'[1, 0] of+  (perm :: Permutation.Perm perm) -> case ftkAst a of+    FTKX sh'@(mn :$% mm :$% shx) x ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        gcastWith (unsafeCoerceRefl :: (Rank perm <=? Rank sh) :~: True) $+        astFromS' (FTKX (mm :$% mn :$% shx) x)+        . astTransposeS perm . astSFromX' sh $ a++astTrBuild+  :: forall k1 k2 s y. AstSpan s+  => SNat k1 -> SNat k2 -> SingletonTK y+  -> AstTensor AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))+  -> AstTensor AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))+astTrBuild SNat SNat stk t = case stk of+  STKScalar -> astTrS t+  STKR{} -> astTr t+  STKS{} -> astTrS t+  STKX{} -> astTrX t+  STKProduct stk1 stk2 ->+    astLetFun t $ \ !tShared ->+      let (u1, u2) = (astProject1 tShared, astProject2 tShared)+      in astPair (astTrBuild (SNat @k1) (SNat @k2) stk1 u1)+                 (astTrBuild (SNat @k1) (SNat @k2) stk2 u2)++astIndexBuild :: forall y k s. AstSpan s+              => SNat k -> FullShapeTK y+              -> AstTensor AstMethodLet s (BuildTensorKind k y)+              -> AstInt AstMethodLet+              -> AstTensor AstMethodLet s y+astIndexBuild snat@SNat ftk u i = case ftk of+  FTKScalar -> astFromS' ftk $ astIndexS ZSS u (i :.$ ZIS)+  FTKR{} -> case ftkAst u of+    FTKR shmshn _ ->+      withShsFromShR shmshn $ \(sh :: ShS sh) ->+        gcastWith (unsafeCoerceRefl :: k ': Tail sh :~: sh) $+        astFromS' ftk $ astIndexS (shsTail sh) (astSFromR' sh u) (i :.$ ZIS)+  FTKS sh _ -> astIndexS sh u (i :.$ ZIS)+  FTKX{} -> case ftkAst u of+   FTKX shBuild' _->+    withShsFromShX shBuild' $ \shBuild -> case shBuild of+      _ :$$ rest ->+        astFromS' ftk $ astIndexS rest (astSFromX' shBuild u) (i :.$ ZIS)+  FTKProduct ftk1 ftk2 ->+    astLetFun u $ \ !u3 ->+      astPair (astIndexBuild snat ftk1 (astProject1 u3) i)+              (astIndexBuild snat ftk2 (astProject2 u3) i)++substProjRep+  :: forall k s s2 y2 y. (AstSpan s, AstSpan s2)+  => SNat k -> IntVarName+  -> AstVarName s2 y2 -> AstTensor AstMethodLet s y+  -> (AstVarName s2 (BuildTensorKind k y2), AstTensor AstMethodLet s y)+substProjRep snat@SNat var var1 v =+  let ftk3 = buildFTK snat $ varNameToFTK var1+      var3 :: AstVarName s2 (BuildTensorKind k y2)+      var3 = mkAstVarName ftk3 (varNameToBounds var1) (varNameToAstVarId var1)+      astVar3 = astVar var3+      v2 = substituteAst+             (astIndexBuild snat (varNameToFTK var1)+                            astVar3 (astVar var))+             var1 v+        -- The subsitutions of projections don't break sharing,+        -- because they don't duplicate variables and the added var+        -- is eventually being eliminated instead of substituted for.+  in (var3, v2)+++-- * Rule tracing machinery++traceRuleEnabledRef :: IORef Bool+{-# NOINLINE traceRuleEnabledRef #-}+traceRuleEnabledRef = unsafePerformIO $ newIORef False++traceNestingLevel :: IORef Int+{-# NOINLINE traceNestingLevel #-}+traceNestingLevel = unsafePerformIO $ newIORef 0++traceWidth :: Int+traceWidth = 90++padString :: Int -> String -> String+padString width full = let cropped = take width full+                       in if length full <= width+                          then take width $ cropped ++ repeat ' '+                          else take (width - 3) cropped ++ "..."++ellipsisString :: Int -> String -> String+ellipsisString width full = let cropped = take width full+                            in if length full <= width+                               then cropped+                               else take (width - 3) cropped ++ "..."++mkTraceRule :: forall y z s. AstSpan s+            => String+            -> AstTensor AstMethodLet s y+            -> AstTensor AstMethodLet s z+            -> Int+            -> AstTensor AstMethodLet s y+            -> AstTensor AstMethodLet s y+{-# NOINLINE mkTraceRule #-}+mkTraceRule !prefix !from !caseAnalysed !nwords ~to = unsafePerformIO $ do+  enabled <- readIORef traceRuleEnabledRef+  let width = traceWidth+      constructorName =+        unwords $ take nwords $ words $ take 21+        $ case caseAnalysed of+          Ast.AstVar{} -> "variable"+          Ast.AstIndexS{} -> "sindex"+          _ -> printAstSimple caseAnalysed+      ruleName = prefix ++ "." ++ constructorName+      ruleNamePadded = take 21 $ ruleName ++ repeat ' '+  when enabled $ do+    nestingLevel <- readIORef traceNestingLevel+    modifyIORef' traceNestingLevel succ+    -- Force in the correct order:+    let !paddedNesting = take 3 $ show nestingLevel ++ repeat ' '+    let !stringFrom = printAstSimple from+    let !stringTo = printAstSimple to+    hPutStrLnFlush stderr $ paddedNesting ++ "rule " ++ ruleNamePadded+                            ++ " sends " ++ padString width stringFrom+                            ++ " to " ++ padString width stringTo+    modifyIORef' traceNestingLevel pred+  let !_A = assert (ftkAst from == ftkAst to+                    `blame` "mkTraceRule: term shape changed"+                    `swith`( ftkAst from, ftkAst to+                           , from, to )) ()+  return $! to++hPutStrLnFlush :: Handle -> String -> IO ()+hPutStrLnFlush target s = do+  hFlush stdout >> hFlush stderr+  hPutStrLn target s+  hFlush stdout >> hFlush stderr
+ src/HordeAd/Core/CarriersADVal.hs view
@@ -0,0 +1,430 @@+{-# LANGUAGE UndecidableInstances #-}+-- | Dual numbers type and operations on it.+-- These definitions, mostly class instances, are needed to make dual numbers+-- a valid carrier for a tensor class algebra (instance) defined+-- in "HordeAd.Core.OpsADVal" so that user programs can be instantiated to+-- or interpreted into it (which corresponds to the forward pass)+-- and subsequently differentiated by evaluating the resulting+-- delta expressions (which corresponds to the backpropagation phase+-- in case of reverse derivatives).+module HordeAd.Core.CarriersADVal+  ( -- * The dual number type+    ADVal, pattern D, dD, dDnotShared+    -- * Auxiliary definitions+  , unDeltaPair, unDeltaPairUnshared, dScale, dAdd+  , dSFromR, dSFromX, dXFromS+  , ensureToplevelSharing, scaleNotShared, addNotShared, multNotShared+  , generateDeltaInputs+  ) where++import Prelude++import Data.Int (Int64)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality ((:~:) (Refl))++import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Shaped.Shape++import Data.Array.Nested.Lemmas+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Delta+import HordeAd.Core.DeltaFreshId+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * The dual number type++-- | The type of dual numbers that are the values of objective functions+-- when they are being differentiated. In general, the primal parts+-- of the dual numbers represent tensors or tuples of tensors.+-- The dual parts are fixed to be delta expressions.+-- The first type argument is the functor determining the nature+-- of the tensors (concrete, symbolic, etc.).+-- The second argument is the tensor kind, which constraints the shapes+-- of the particular tensor (or tensor tuple).+type role ADVal nominal nominal+data ADVal (f :: Target) y = ADVal (f y) (Delta f y)++pattern D :: f z -> Delta f z -> ADVal f z+pattern D t u <- ADVal t u  -- enforces only pattern matching+{-# COMPLETE D #-}++deriving instance (Show (f z), Show (Delta f z))+                  => Show (ADVal f z)++-- | Smart constructor for t'ADVal' that additionally records delta+-- expression sharing information (regardless if the primal part+-- of the dual number is an AST term or not).+-- The bare constructor should not (and cannot) be used for constructing+-- values, but only for deconstructing dual numbers via pattern-matching.+dD :: forall f z.+      f z -> Delta f z -> ADVal f z+dD !a !dual = dDnotShared a (shareDelta dual)++-- | This is a not so smart constructor for t'ADVal' that does not record+-- sharing information. If used in contexts where duplication occurs,+-- it may cause exponential blowup when a delta expression is evaluated+-- in backpropagation phase. In contexts without duplication, it saves+-- some evaluation time and memory (in delta term structure, but even more+-- in the per-node data stored while evaluating).+dDnotShared :: f z -> Delta f z -> ADVal f z+dDnotShared = ADVal+++-- * Auxiliary definitions++-- TODO: maybe create a separate module of Delta smart constructors+-- and then use the following haddocks for the module:++-- The instances are impure, because 'shareDelta'+-- adorns terms with an @Int@ identifier from a counter that is afterwards+-- incremented (and never changed in any other way).+--+-- The identifiers are not part of any non-internal module API+-- and the impure counter that gets incremented is not exposed+-- (except for low level tests). The identifiers are assigned here once+-- and ever accessed read-only.+-- Their uniqueness ensures that subterms that are shared in memory+-- are evaluated only once. If pointer equality worked efficiently+-- (e.g., if compact regions with sharing were cheaper), we wouldn't need+-- the impurity.+--+-- Given that we have to use impurity anyway, we make the implementation+-- faster by ensuring the order of identifiers reflects data dependency,+-- that is, parent nodes always have higher identifier than child nodes.+-- The @StrictData@ extension ensures that the delta constructors+-- are call by value, which is needed for that identifier ordering.+--+-- As long as "HordeAd.Core.Delta" is used exclusively through+-- smart constructors from this API, the impurity is completely safe+-- (currently, it's enough that @DeltaShare@ is used exclusively+-- via @shareDelta@). Even compiler optimizations, e.g., cse and full-laziness,+-- can't break the required invariants. On the contrary,+-- they increase sharing and make evaluation yet cheaper.++unDeltaPair :: Delta target (TKProduct x y) -> (Delta target x, Delta target y)+unDeltaPair (DeltaPair a b) = (a, b)+unDeltaPair (DeltaZero (FTKProduct ftk1 ftk2)) =+  (DeltaZero ftk1, DeltaZero ftk2)+unDeltaPair d = let dShared = shareDelta d+                in (DeltaProject1 dShared, DeltaProject2 dShared)++unDeltaPairUnshared :: Delta target (TKProduct x y)+                    -> (Delta target x, Delta target y)+unDeltaPairUnshared (DeltaPair a b) = (a, b)+unDeltaPairUnshared (DeltaZero (FTKProduct ftk1 ftk2)) =+  (DeltaZero ftk1, DeltaZero ftk2)+unDeltaPairUnshared d = (DeltaProject1 d, DeltaProject2 d)  -- duplicated++dScale :: Num (f z) => f z -> Delta f z -> Delta f z+dScale _ (DeltaZero ftk) = DeltaZero ftk+dScale v u' = DeltaScale (NestedTarget v) u'++dAdd :: Num (f z) => Delta f z -> Delta f z -> Delta f z+dAdd DeltaZero{} w = w+dAdd v DeltaZero{} = v+dAdd v w = DeltaAdd v w++-- Prevents building huge Delta terms, not only evaluating them.+dSFromR :: forall sh x target.+           ShS sh -> Delta target (TKR2 (Rank sh) x)+        -> Delta target (TKS2 sh x)+dSFromR sh w@(DeltaConvert _c d)+  | FTKR _ x <- ftkDelta w+  , Just Refl <- matchingFTK (ftkDelta d) (FTKS sh x) = d+dSFromR sh d | FTKR _ x <- ftkDelta d+             , Refl <- lemRankReplicate (Proxy @(Rank sh)) =+  let c2 = ConvCmp (ConvXS' (FTKS sh x)) ConvRX+  in DeltaConvert c2 d++dSFromX :: forall sh sh' x target. Rank sh ~ Rank sh'+        => ShS sh -> Delta target (TKX2 sh' x)+        -> Delta target (TKS2 sh x)+dSFromX sh w@(DeltaConvert _c d)+  | FTKX _ x <- ftkDelta w+  , Just Refl <- matchingFTK (ftkDelta d) (FTKS sh x) = d+dSFromX sh d | FTKX _ x <- ftkDelta d =+  let c2 = ConvXS' (FTKS sh x)+  in DeltaConvert c2 d++dXFromS :: forall sh sh' x target. Rank sh ~ Rank sh'+        => StaticShX sh' -> Delta target (TKS2 sh x)+        -> Delta target (TKX2 sh' x)+dXFromS ssx w@(DeltaConvert _c d)+  | FTKS sh x <- ftkDelta w+  , let shx = shCastSX ssx sh+  , Just Refl <- matchingFTK (ftkDelta d) (FTKX shx x) = d+dXFromS ssx d+  | FTKS sh x <- ftkDelta d+  , let shx = shCastSX ssx sh+  , Refl <- lemRankMapJust sh =+    let c2 = ConvCmp (ConvXX' (FTKX shx x)) ConvSX+    in DeltaConvert c2 d++-- This hack is needed to recover shape from tensors,+-- in particular in case of numeric literals and also for forward derivative.+intOfShape :: forall z f. ADReadyNoLet f+           => Delta f z -> Int -> f z+intOfShape tsh c = treplTarget (fromIntegral c) (ftkDelta tsh)++-- | Add sharing information to the top level of a term, presumably+-- constructed using multiple applications of the `dDnotShared` operation.+-- The resulting term may not have sharing information inside,+-- but is ready to be shared as a whole.+ensureToplevelSharing :: ADVal f z -> ADVal f z+ensureToplevelSharing (D u u') = dD u u'++scaleNotShared :: Num (f z)+               => f z -> ADVal f z -> ADVal f z+scaleNotShared !a (D u u') = dDnotShared (a * u) (dScale a u')++addNotShared :: forall f z. Num (f z)+             => ADVal f z -> ADVal f z -> ADVal f z+addNotShared (D u u') (D v v') = dDnotShared (u + v) (dAdd u' v')++multNotShared :: forall f z. Num (f z)+              => ADVal f z -> ADVal f z -> ADVal f z+multNotShared (D u u') (D v v') =+  dDnotShared (u * v) (dAdd (dScale v u') (dScale u v'))++generateDeltaInputs :: forall x target.+                       FullShapeTK x -> Delta target x+generateDeltaInputs =+  let gen :: Int -> FullShapeTK y -> (Delta target y, Int)+      gen j ftk = case ftk of+        FTKProduct ftk1 ftk2 ->+          let (d1, j1) = gen j ftk1+              (d2, j2) = gen j1 ftk2+          in (DeltaPair d1 d2, j2)+        _ | differentiableFTK ftk -> (DeltaInput (mkInputId ftk j), j + 1)+        _ -> (DeltaZero ftk, j)+  in fst . gen 0+++-- * Assorted instances++type instance BoolOf (ADVal f) = BoolOf f++instance EqH f (TKScalar r) => EqH (ADVal f) (TKScalar r) where+  D u _ ==. D v _ = u ==. v++instance OrdH f (TKScalar r) => OrdH (ADVal f) (TKScalar r) where+  D u _ <=. D v _ = u <=. v++instance EqH f (TKR n r) => EqH (ADVal f) (TKR n r) where+  D u _ ==. D v _ = u ==. v++instance OrdH f (TKR n r) => OrdH (ADVal f) (TKR n r) where+  D u _ <=. D v _ = u <=. v++instance EqH f (TKS sh r) => EqH (ADVal f) (TKS sh r) where+  D u _ ==. D v _ = u ==. v++instance OrdH f (TKS sh r) => OrdH (ADVal f) (TKS sh r) where+  D u _ <=. D v _ = u <=. v++instance EqH f (TKX sh r) => EqH (ADVal f) (TKX sh r) where+  D u _ ==. D v _ = u ==. v++instance OrdH f (TKX sh r) => OrdH (ADVal f) (TKX sh r) where+  D u _ <=. D v _ = u <=. v++type instance HFunOf (ADVal f) x y = HFun x y++type instance PrimalOf (ADVal f) = f++type instance DualOf (ADVal f) = Delta f++type instance ShareOf (ADVal f) = ADVal f+  -- Maybe this should be ADVal (ShareOf f), but we'd need tests+  -- that use this, probably tests with ADVal (AST) nested in ADVal+++-- * Numeric instances++-- These two instances are required for the numeric tensor instances.+-- They can't be made valid for AST, because they require interpretation before+-- they can be compared with an instant Bool result, so let's fail early+-- also here.+instance Eq (ADVal f z) where+  (==) = error "Eq is not defined for ADVal; please use EqH instead"+  (/=) = error "Eq is not defined for ADVal; please use EqH instead"++instance Ord (ADVal f z) where+  (<=) = error "Ord is not defined for ADVal; please use OrdH instead"++-- This is copied from below to permit fromInteger for TKScalar.+-- This OVERLAPPABLE seems to work 100% reliably for indexes+-- and not at all for a variant of rfromListLinear that takes scalars.+instance (GoodScalar r, ShareTensor f, ADReadyNoLet f)+         => Num (ADVal f (TKScalar r)) where+  D u u' + D v v' = dD (u + v) (dAdd u' v')+  D u u' - D v v' =+    dD (u - v) (dAdd u' (dScale (intOfShape v' (-1)) v'))+  D ue u' * D ve v' =+    -- The bangs are neccessary for GHC 9.2.7 test results to match 9.4.+    let !u = tshare ue in+    let !v = tshare ve+    in dD (u * v) (dAdd (dScale v u') (dScale u v'))+  negate (D v v') = dD (negate v) (dScale (intOfShape v' (-1)) v')+  abs (D ve v') = let !v = tshare ve+                  in dD (abs v) (dScale (signum v) v')+  signum (D v v') = dDnotShared (signum v) (DeltaZero $ ftkDelta v')+  fromInteger i = dDnotShared (fromInteger i) (DeltaZero FTKScalar)+  -- The constraints in the pragmas below are needed only to avoid+  -- module import cycles.+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKScalar Double)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKScalar Float)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKScalar Int64)) #-}++instance {-# OVERLAPPABLE #-}+         (Num (f z), ShareTensor f, ADReadyNoLet f)+         => Num (ADVal f z) where+  D u u' + D v v' = dD (u + v) (dAdd u' v')+  D u u' - D v v' =+    dD (u - v) (dAdd u' (dScale (intOfShape v' (-1)) v'))+  D ue u' * D ve v' =+    -- The bangs are neccessary for GHC 9.2.7 test results to match 9.4.+    let !u = tshare ue in+    let !v = tshare ve+    in dD (u * v) (dAdd (dScale v u') (dScale u v'))+  negate (D v v') = dD (negate v) (dScale (intOfShape v' (-1)) v')+  abs (D ve v') = let !v = tshare ve+                  in dD (abs v) (dScale (signum v) v')+  signum (D v v') = dDnotShared (signum v) (DeltaZero $ ftkDelta v')+  fromInteger = error "fromInteger is not defined for tensors in general"+  -- The constraints in the pragmas below are needed only to avoid+  -- module import cycles.+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKR n Double)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKR n Float)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKR n Int64)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKS sh Double)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKS sh Float)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKS sh Int64)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKX sh Double)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKX sh Float)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => Num (ADVal Concrete (TKX sh Int64)) #-}++instance (Real (f z), ShareTensor f, ADReadyNoLet f)+         => Real (ADVal f z) where+  toRational (D v _) = toRational v+    -- this is most probably not what the user expects, but the type+    -- of the result (Rational) doesn't permit any better solution++instance (IntegralH (f z), ShareTensor f, ADReadyNoLet f)+         => IntegralH (ADVal f z) where+  quotH (D u _) (D v v') = dDnotShared (quotH u v) (DeltaZero $ ftkDelta v')+  remH (D u _) (D v v') = dDnotShared (remH u v) (DeltaZero $ ftkDelta v')+  -- The constraints in the pragmas below are needed only to avoid+  -- module import cycles.+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => IntegralH (ADVal Concrete (TKR n Int64)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => IntegralH (ADVal Concrete (TKS sh Int64)) #-}+  {-# SPECIALIZE instance (ShareTensor Concrete, ADReadyNoLet Concrete) => IntegralH (ADVal Concrete (TKX sh Int64)) #-}++-- This is copied from below to permit fromRational for TKScalar.+instance ( GoodScalar r, Fractional (f (TKScalar r)), ShareTensor f+         , ADReadyNoLet f )+         => Fractional (ADVal f (TKScalar r)) where+  D ue u' / D ve v' =+    let !u = tshare ue in+    let !v = tshare ve+    in dD (u / v) (dAdd (dScale (recip v) u') (dScale ((- u) / (v * v)) v'))+  recip (D ve v') =+    let !v = tshare ve+        minusRecipSq = - recip (v * v)+    in dD (recip v) (dScale minusRecipSq v')+  fromRational r = dDnotShared (fromRational r) (DeltaZero FTKScalar)++instance {-# OVERLAPPABLE #-}+         (Fractional (f z), ShareTensor f, ADReadyNoLet f)+         => Fractional (ADVal f z) where+  D ue u' / D ve v' =+    let !u = tshare ue in+    let !v = tshare ve+    in dD (u / v) (dAdd (dScale (recip v) u') (dScale ((- u) / (v * v)) v'))+  recip (D ve v') =+    let !v = tshare ve+        minusRecipSq = - recip (v * v)+    in dD (recip v) (dScale minusRecipSq v')+  fromRational = error "fromRational is not defined for tensors in general"++instance (Floating (f z), ShareTensor f, ADReadyNoLet f)+         => Floating (ADVal f z) where+  pi = error "pi is not defined for tensors"+  exp (D ue u') = let !expU = tshare (exp ue)+                  in dD expU (dScale expU u')+  log (D ue u') = let !u = tshare ue+                  in dD (log u) (dScale (recip u) u')+  sqrt (D ue u') = let !sqrtU = tshare (sqrt ue)+                   in dD sqrtU (dScale (recip (sqrtU + sqrtU)) u')+  D ue u' ** D ve v' =+    let !u = tshare ue in+    let !v = tshare ve+    in dD (u ** v) (dAdd (dScale (v * (u ** (v - intOfShape v' 1))) u')+                         (dScale ((u ** v) * log u) v'))+  -- logBase x y = log y / log x+  sin (D ue u') = let !u = tshare ue+                  in dD (sin u) (dScale (cos u) u')+  cos (D ue u') = let !u = tshare ue+                  in dD (cos u) (dScale (- (sin u)) u')+  tan (D ue u') = let !u = tshare ue in+                  let !cosU = tshare (cos u)+                  in dD (tan u) (dScale (recip (cosU * cosU)) u')+  asin (D ue u') = let !u = tshare ue+                   in dD (asin u)+                         (dScale (recip (sqrt (intOfShape u' 1 - u * u))) u')+  acos (D ue u') = let !u = tshare ue+                   in dD (acos u)+                         (dScale (- recip (sqrt (intOfShape u' 1 - u * u))) u')+  atan (D ue u') = let !u = tshare ue+                   in dD (atan u)+                         (dScale (recip (intOfShape u' 1 + u * u)) u')+  sinh (D ue u') = let !u = tshare ue+                   in dD (sinh u) (dScale (cosh u) u')+  cosh (D ue u') = let !u = tshare ue+                   in dD (cosh u) (dScale (sinh u) u')+  tanh (D ue u') = let !y = tshare (tanh ue)+                   in dD y (dScale (intOfShape u' 1 - y * y) u')+  asinh (D ue u') = let !u = tshare ue+                    in dD (asinh u)+                          (dScale (recip (sqrt (intOfShape u' 1 + u * u))) u')+  acosh (D ue u') = let !u = tshare ue+                    in dD (acosh u)+                          (dScale (recip (sqrt (u * u - intOfShape u' 1))) u')+  atanh (D ue u') = let !u = tshare ue+                    in dD (atanh u)+                          (dScale (recip (intOfShape u' 1 - u * u)) u')++instance (RealFrac (f z), ShareTensor f, ADReadyNoLet f)+         => RealFrac (ADVal f z) where+  properFraction = error "properFraction is not defined for tensors"+    -- The integral type doesn't have a Storable constraint,+    -- so we can't implement this (nor RealFracB from Boolean package).++instance (Fractional (f z), RealFloatH (f z), ShareTensor f, ADReadyNoLet f)+         => RealFloatH (ADVal f z) where+  atan2H (D ue u') (D ve v') =+    let !u = tshare ue in+    let !v = tshare ve in+    let !t = tshare (recip (u * u + v * v))+    in dD (atan2H u v) (dAdd (dScale ((- u) * t) v') (dScale (v * t) u'))++instance (RealFloat (f z), ShareTensor f, ADReadyNoLet f)+         => RealFloat (ADVal f z) where+  atan2 (D ue u') (D ve v') =+    let !u = tshare ue in+    let !v = tshare ve in+    let !t = tshare (recip (u * u + v * v))+    in dD (atan2 u v) (dAdd (dScale ((- u) * t) v') (dScale (v * t) u'))+  floatRadix (D u _) = floatRadix u+  floatDigits (D u _) = floatDigits u+  floatRange (D u _) = floatRange u+  decodeFloat (D u _) = decodeFloat u+  encodeFloat _i _j = error "encodeFloat is not defined for tensors"+  isNaN (D u _) = isNaN u+  isInfinite (D u _) = isInfinite u+  isDenormalized (D u _) = isDenormalized u+  isNegativeZero (D u _) = isNegativeZero u+  isIEEE (D u _) = isIEEE u
+ src/HordeAd/Core/CarriersAst.hs view
@@ -0,0 +1,1149 @@+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Definitions, mostly class instances, needed to make AST a valid+-- carrier for a tensor class algebra (instance) defined+-- in "HordeAd.Core.OpsAst".+-- This algebra permits user programs to be instantiated as AST terms,+-- as well as to be interpreted into AST terms and it also permits derivatives+-- to be expressed as AST terms.+module HordeAd.Core.CarriersAst+  ( AstRaw(..), AstNoVectorize(..), AstNoSimplify(..)+  , astLetFunNoSimplify, sunReplicateScal, sunReplicate1, sunReplicateN+  ) where++import Prelude hiding (foldl')++import Data.Int (Int64)+import Data.Type.Equality (testEquality, (:~:) (Refl))+import Foreign.C (CInt)+import Type.Reflection (typeRep)++import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)+import Data.Array.Nested.Mixed qualified as Mixed+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Shaped.Shape++import HordeAd.Core.Ast+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.OpsConcrete ()+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * Type family instances for AstTensor++type instance PrimalOf (AstTensor ms s) = AstTensor ms PrimalSpan+type instance DualOf (AstTensor ms s) = AstTensor ms DualSpan+type instance ShareOf (AstTensor ms s) = AstRaw s++-- This can't be just HFun, because they need to be vectorized+-- and vectorization applies such functions to the variable from build1+-- and the variable has to be eliminated via vectorization to preserve+-- the closed form of the function. Just applying a Haskell closure+-- to the build1 variable and then duplicating the result of the function+-- would not eliminate the variable and also would likely results+-- in more costly computations. Also, that would prevent simplification+-- of the instances, especially after applied to arguments that are terms.+type instance HFunOf (AstTensor AstMethodLet s) x z = AstHFun s s x z++type instance BoolOf (AstTensor ms s) = AstBool ms+++-- * Unlawful numeric instances for AST scalars; they are lawful modulo evaluation++-- These are, unfortunately, required by some numeric instances.+instance Eq (AstTensor ms s y) where+  (==) = error "Eq is not defined for AST; please use EqH instead"+  (/=) = error "Eq is not defined for AST; please use EqH instead"++instance Ord (AstTensor ms s y) where+  (<=) = error "Ord is not defined for AST; please use OrdH instead"++-- TODO: perhaps aim for a polynomial normal form? but that requires global+-- inspection of the whole expression+-- TODO: let's aim at SOP (Sum-of-Products) form, just as+-- ghc-typelits-natnormalise does. Also, let's associate to the right+-- and let's push negation down.+--+-- Not considered are rules that would require comparing non-constant terms+-- or that would duplicate a non-constant term, as well as most rules+-- informed by inequalities, expressed via max or min, such as+-- max n (signum (abs x)) | n <= 0 --> signum (abs x).+-- We could use sharing via @tlet@ if terms are duplicated, but it's+-- unclear if the term bloat is worth it and also we'd need to restrict+-- this extended simplification to AstMethodLet.+--+-- | Integer terms need to be simplified, because large ones are sometimes+-- created due to vectorization, e.g., via astTransposeAsGather+-- or astReshapeAsGather and can be a deciding factor in whether+-- the other tensor terms can be simplified in turn.+--+-- The normal form has AstConcreteK, if any, as the first argument+-- of the constructor. No flattening is performed beyond that.+instance (GoodScalar r, AstSpan s)+         => Num (AstTensor ms s (TKScalar r)) where+  AstFromPrimal u + AstFromPrimal v = AstFromPrimal $ u + v+  AstFromDual u + AstFromDual v = AstFromDual $ u + v+  -- TODO: define a pattern synonym that captures the below. Also elsewhere.+  AstConvert c u + AstConvert _ v+    | FTKS ZSS x <- ftkAst u+    , FTKS ZSS y <- ftkAst v+    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u))+    , Just Refl <- matchingFTK x y =+      AstConvert c $ u + v+  AstConcreteK 0 + u = u+  u + AstConcreteK 0 = u+  AstConcreteK n + AstConcreteK k = AstConcreteK (n + k)+  AstConcreteK n + AstPlusK (AstConcreteK k) u = AstConcreteK (n + k) + u+  AstPlusK (AstConcreteK n) u + AstConcreteK k = AstConcreteK (n + k) + u+  AstPlusK (AstConcreteK n) u + AstPlusK (AstConcreteK k) v =+    AstConcreteK (n + k) + AstPlusK u v  -- u and v can cancel, but unlikely++  -- Unfortunately, these only fire if the required subterms are at the top+  -- of the reduced term, which happens rarely except in small terms.+  -- We could keep variables at the top, but they'd compete with AstConcreteK.+  AstN1K NegateOp (AstVar var) + AstVar var'+    | var == var' = 0+  AstN1K NegateOp (AstConvert _ (AstVar var))+    + AstConvert _ (AstVar var')+    | varNameToAstVarId var == varNameToAstVarId var' = 0+  AstN1K NegateOp (AstVar var) + AstPlusK (AstVar var') u+    | var == var' = u+  AstN1K NegateOp (AstConvert _ (AstVar var))+    + AstPlusK (AstConvert _ (AstVar var')) u+    | varNameToAstVarId var == varNameToAstVarId var' = u+  AstVar var' + AstN1K NegateOp (AstVar var)+    | var == var' = 0+  AstConvert _ (AstVar var')+    + AstN1K NegateOp (AstConvert _ (AstVar var))+    | varNameToAstVarId var == varNameToAstVarId var' = 0+  AstVar var' + AstPlusK (AstN1K NegateOp (AstVar var)) u+    | var == var' = u+  AstConvert _ (AstVar var')+    + AstPlusK (AstN1K NegateOp (AstConvert _ (AstVar var))) u+    | varNameToAstVarId var == varNameToAstVarId var' = u++  AstI2K RemOp (AstN1K NegateOp (AstVar var)) (AstConcreteK n)+   + AstI2K RemOp (AstVar var') (AstConcreteK n')+     | var == var' && n == n' = 0+  AstI2K RemOp (AstN1K NegateOp (AstConvert _ (AstVar var)))+               (AstConcreteK n)+   + AstI2K RemOp (AstConvert _ (AstVar var')) (AstConcreteK n')+     | varNameToAstVarId var == varNameToAstVarId var' && n == n' = 0+  AstI2K RemOp (AstN1K NegateOp (AstVar var)) (AstConcreteK n)+   + AstPlusK (AstI2K RemOp (AstVar var') (AstConcreteK n')) u+     | var == var' && n == n' = u+  AstI2K RemOp (AstN1K NegateOp (AstConvert _ (AstVar var)))+               (AstConcreteK n)+   + AstPlusK (AstI2K RemOp (AstConvert _ (AstVar var'))+                            (AstConcreteK n')) u+     | varNameToAstVarId var == varNameToAstVarId var' && n == n' = u+  AstI2K RemOp (AstVar var') (AstConcreteK n')+   + AstI2K RemOp (AstN1K NegateOp (AstVar var)) (AstConcreteK n)+     | var == var' && n == n' = 0+  AstI2K RemOp (AstConvert _ (AstVar var')) (AstConcreteK n')+   + AstI2K RemOp (AstN1K NegateOp (AstConvert _ (AstVar var)))+                                   (AstConcreteK n)+     | varNameToAstVarId var == varNameToAstVarId var' && n == n' = 0+  AstI2K RemOp (AstVar var') (AstConcreteK n')+   + AstPlusK (AstI2K RemOp (AstN1K NegateOp (AstVar var)) (AstConcreteK n)) u+     | var == var' && n == n' = u+  AstI2K RemOp (AstConvert _ (AstVar var')) (AstConcreteK n')+   + AstPlusK (AstI2K RemOp (AstN1K NegateOp (AstConvert _ (AstVar var)))+                            (AstConcreteK n)) u+     | varNameToAstVarId var == varNameToAstVarId var' && n == n' = u++  AstPlusK u@AstConcreteK{} v + w = AstPlusK u (AstPlusK v w)  -- as above+  u + v@AstConcreteK{} = AstPlusK v u+  u + AstPlusK v@AstConcreteK{} w = AstPlusK v (AstPlusK u w)  -- as above+  t1 + t2 | eqK t1 t2 = 2 * t1+  t1 + AstTimesK (AstConcreteK n) t2 | eqK t1 t2 = AstConcreteK (n + 1) * t1+  AstTimesK (AstConcreteK n) t2 + t1 | eqK t1 t2 = AstConcreteK (n + 1) * t1+  AstTimesK (AstConcreteK n1) t1 + AstTimesK (AstConcreteK n2) t2+    | eqK t1 t2 = AstConcreteK (n1 + n2) * t1+  u + v = AstPlusK u v++  AstFromPrimal u * AstFromPrimal v = AstFromPrimal $ u * v+  AstFromDual{} * AstFromDual{} = 0+  AstConcreteK 0 * _ = 0+  _ * AstConcreteK 0 = 0+  AstConcreteK 1 * u = u+  u * AstConcreteK 1 = u+  AstConvert c u * AstConvert _ v+    | FTKS ZSS x <- ftkAst u+    , FTKS ZSS y <- ftkAst v+    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u))+    , Just Refl <- matchingFTK x y =+      AstConvert c $ u * v+  AstConcreteK n * AstConcreteK k = AstConcreteK (n * k)+  AstConcreteK n * AstTimesK (AstConcreteK k) u = AstConcreteK (n * k) * u+  AstTimesK (AstConcreteK n) u * AstConcreteK k = AstConcreteK (n * k) * u+  AstTimesK (AstConcreteK n) u * AstTimesK (AstConcreteK k) v =+    AstConcreteK (n * k) * AstTimesK u v  -- u and v can cancel, but unlikely++  u@AstConcreteK{} * AstPlusK v w = AstPlusK (u * v) (u * w)+  AstTimesK u@AstConcreteK{} x * AstPlusK v w =+    AstTimesK x (AstPlusK (u * v) (u * w))+  AstPlusK v w * u@AstConcreteK{} = AstPlusK (v * u) (w * u)+  AstPlusK v w * AstTimesK u@AstConcreteK{} x =+    AstTimesK (AstPlusK (v * u) (w * u)) x++  AstN1K NegateOp u * AstN1K NegateOp v = AstTimesK u v++  {- TODO: these rules increase the number of occurrences of a variable+     and trade multiplication and quotient for an equally problematic remnant,+     so they are disabled until we find a way to profit from them.+  -- With static shapes, the second argument to QuotOp and RemOp+  -- is often a constant, which makes such rules worth including,+  -- since they are likely to fire. To help them fire, we avoid changing+  -- that constant, if possible, e.g., in rules for NegateOp.+  AstConcreteK n * AstI2K QuotOp (AstVar var) (AstConcreteK n')+    | n == n' =+      AstPlusK+        (AstVar var)+        (negate (AstI2K RemOp (AstVar var) (AstConcreteK n)))+  AstTimesK (AstConcreteK n) x * AstI2K QuotOp (AstVar var)+                                               (AstConcreteK n')+    | n == n' =+      AstTimesK+        x+        (AstPlusK+           (AstVar var)+           (negate (AstI2K RemOp (AstVar var) (AstConcreteK n))))+  AstI2K QuotOp (AstVar var) (AstConcreteK n') * AstConcreteK n+    | n == n' =+      AstPlusK+        (AstVar var)+        (negate (AstI2K RemOp (AstVar var) (AstConcreteK n)))+  AstI2K QuotOp (AstVar var)+                (AstConcreteK n') * AstTimesK (AstConcreteK n) x+    | n == n' =+      AstTimesK+        (AstPlusK+           (AstVar var)+           (negate (AstI2K RemOp (AstVar var) (AstConcreteK n))))+        x+  -}++  AstTimesK u@AstConcreteK{} v * w = AstTimesK u (AstTimesK v w)  -- as above+  u * v@AstConcreteK{} = AstTimesK v u+  u * AstTimesK v@AstConcreteK{} w = AstTimesK v (AstTimesK u w)  -- as above+  u * v = AstTimesK u v++  negate (AstCond b n k) = AstCond b (negate n) (negate k)+  negate (AstLet var n k) = AstLet var n (negate k)+  negate (AstFromPrimal n) = AstFromPrimal (negate n)+  negate (AstFromDual n) = AstFromDual (negate n)+  negate (AstPlusK u v) = AstPlusK (negate u) (negate v)+  negate (AstTimesK u v) = negate u * v+  negate (AstN1K NegateOp u) = u+  negate (AstN1K SignumOp u) = AstN1K SignumOp (negate u)+  negate (AstI2K QuotOp u v) = AstI2K QuotOp (negate u) v+    -- v is likely positive and let's keep it so+  negate (AstI2K RemOp u v) = AstI2K RemOp (negate u) v+    -- v is likely positive and let's keep it so+  negate (AstConcreteK n) = AstConcreteK (negate n)+  negate (AstConvert c u)+    | FTKS ZSS x <- ftkAst u+    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u)) =+      AstConvert c (negate u)+  negate u = AstN1K NegateOp u+  abs (AstFromPrimal n) = AstFromPrimal (abs n)+  abs (AstFromDual n) = AstFromDual (abs n)+  abs (AstConcreteK n) = AstConcreteK (abs n)+  abs (AstN1K AbsOp u) = AstN1K AbsOp u+  abs (AstN1K NegateOp u) = abs u+  abs (AstConvert c u)+    | FTKS ZSS x <- ftkAst u+    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u)) =+      AstConvert c (abs u)+  abs u = AstN1K AbsOp u+  signum (AstFromPrimal n) = AstFromPrimal (signum n)+  signum (AstFromDual n) = AstFromDual (signum n)+  signum (AstConcreteK n) = AstConcreteK (signum n)+  signum (AstN1K SignumOp u) = AstN1K SignumOp u+  signum (AstConvert c u)+    | FTKS ZSS x <- ftkAst u+    , Just Refl <- matchingFTK x (convertFTK c (ftkAst u)) =+      AstConvert c (signum u)+  signum u = AstN1K SignumOp u+  fromInteger i = fromPrimal $ AstConcreteK (fromInteger i)+  {-# SPECIALIZE instance Num (AstTensor ms FullSpan (TKScalar Int64)) #-}+  {-# SPECIALIZE instance Num (AstTensor ms PrimalSpan (TKScalar Int64)) #-}+  {-# SPECIALIZE instance Num (AstTensor ms FullSpan (TKScalar CInt)) #-}+  {-# SPECIALIZE instance Num (AstTensor ms PrimalSpan (TKScalar CInt)) #-}+  {-# SPECIALIZE instance Num (AstTensor ms FullSpan (TKScalar Double)) #-}+  {-# SPECIALIZE instance Num (AstTensor ms PrimalSpan (TKScalar Double)) #-}+  {-# SPECIALIZE instance Num (AstTensor ms FullSpan (TKScalar Float)) #-}+  {-# SPECIALIZE instance Num (AstTensor ms PrimalSpan (TKScalar Float)) #-}++-- An approximation. False doesn't imply terms have different semantics,+-- but True implies they have equal semantics.+eqK :: AstTensor ms s (TKScalar r) -> AstTensor ms s (TKScalar r) -> Bool+eqK (AstVar var1) (AstVar var2) = var1 == var2+eqK (AstLet @_ @_ @s1 var1 u1 v1) (AstLet @_ @_ @s2 var2 u2 v2)+  | FTKScalar @r1 <- ftkAst u1, FTKScalar @r2 <- ftkAst u2+  , Just Refl <- testEquality (typeRep @r1) (typeRep @r2)+  , Just Refl <- sameAstSpan @s1 @s2 =+    var1 == var2 && eqK u1 u2 && eqK v1 v2+eqK (AstPrimalPart u1) (AstPrimalPart u2) = eqK u1 u2+eqK (AstDualPart u1) (AstDualPart u2) = eqK u1 u2+eqK (AstFromPrimal u1) (AstFromPrimal u2) = eqK u1 u2+eqK (AstFromDual u1) (AstFromDual u2) = eqK u1 u2+eqK (AstPlusK u1 v1) (AstPlusK u2 v2) =+  eqK u1 u2 && eqK v1 v2 || eqK u1 v2 && eqK v1 u2+eqK (AstTimesK u1 v1) (AstTimesK u2 v2) =+  eqK u1 u2 && eqK v1 v2 || eqK u1 v2 && eqK v1 u2+eqK (AstN1K opCode1 u1) (AstN1K opCode2 u2) = opCode1 == opCode2 && eqK u1 u2+eqK (AstR1K opCode1 u1) (AstR1K opCode2 u2) = opCode1 == opCode2 && eqK u1 u2+eqK (AstR2K opCode1 u1 v1) (AstR2K opCode2 u2 v2) =+  opCode1 == opCode2 && eqK u1 u2 && eqK v1 v2+eqK (AstI2K opCode1 u1 v1) (AstI2K opCode2 u2 v2) =+  opCode1 == opCode2 && eqK u1 u2 && eqK v1 v2+eqK (AstConcreteK u1) (AstConcreteK u2) = u1 == u2+eqK (AstFloorK @r1 u1) (AstFloorK @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) = eqK u1 u2+eqK (AstFromIntegralK @r1 u1) (AstFromIntegralK @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) = eqK u1 u2+eqK (AstCastK @r1 u1) (AstCastK @r2 u2)+  | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) = eqK u1 u2+eqK _ _ = False++-- Div and mod operations are very costly (simplifying them requires+-- constructing conditionals, etc), so they are not included in IntegralH.+instance (GoodScalar r, IntegralH r, Nested.IntElt r, AstSpan s)+         => IntegralH (AstTensor ms s (TKScalar r)) where+  quotH (AstFromPrimal n) (AstFromPrimal k) = AstFromPrimal (quotH n k)+  quotH (AstConvert c n) (AstConvert _ k)+    | FTKS ZSS x <- ftkAst n+    , FTKS ZSS y <- ftkAst k+    , Just Refl <- matchingFTK x (convertFTK c (ftkAst n))+    , Just Refl <- matchingFTK x y =+      AstConvert c (quotH n k)+  quotH (AstConcreteK n) (AstConcreteK k) = AstConcreteK (quotH n k)+  quotH (AstConcreteK 0) _ = 0+  quotH u (AstConcreteK 1) = u+  quotH (AstI2K RemOp _ (AstConcreteK k)) (AstConcreteK k')+    | k' >= k && k >= 0 = 0+  quotH (AstI2K QuotOp u v) w = quotH u (v * w)+  quotH (AstTimesK (AstConcreteK n) v) (AstConcreteK n')+    | n == n' = v+  quotH u v =+    let t = AstI2K QuotOp u v+        (u1, u2) = bounds t+    in if u1 == u2 then fromPrimal $ AstConcreteK u1 else t++  remH (AstFromPrimal n) (AstFromPrimal k) = AstFromPrimal (remH n k)+  remH (AstConvert c n) (AstConvert _ k)+    | FTKS ZSS x <- ftkAst n+    , FTKS ZSS y <- ftkAst k+    , Just Refl <- matchingFTK x (convertFTK c (ftkAst n))+    , Just Refl <- matchingFTK x y =+      AstConvert c (remH n k)+  remH (AstConcreteK n) (AstConcreteK k) = AstConcreteK (remH n k)+  remH (AstConcreteK 0) _ = 0+  remH _ (AstConcreteK 1) = 0+  remH (AstI2K RemOp u (AstConcreteK k)) (AstConcreteK k')+    | k' >= k && k >= 0 = AstI2K RemOp u (AstConcreteK k)+  remH (AstI2K RemOp u (AstConcreteK k)) (AstConcreteK k')+    | remH k k' == 0 && k > 0 = remH u (AstConcreteK k')+  remH (AstTimesK (AstConcreteK n) _) (AstConcreteK n')+    | remH n n' == 0 = 0+  remH u v =+    let t = AstI2K RemOp u v+        (u1, u2) = bounds t+    in if u1 == u2 then fromPrimal $ AstConcreteK u1 else t+  {-# SPECIALIZE instance IntegralH (AstTensor ms FullSpan (TKScalar Int64)) #-}+  {-# SPECIALIZE instance IntegralH (AstTensor ms PrimalSpan (TKScalar Int64)) #-}+  {-# SPECIALIZE instance IntegralH (AstTensor ms FullSpan (TKScalar CInt)) #-}+  {-# SPECIALIZE instance IntegralH (AstTensor ms PrimalSpan (TKScalar CInt)) #-}++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => Fractional (AstTensor ms s (TKScalar r)) where+  AstFromPrimal u / AstFromPrimal v = AstFromPrimal $ u / v+  AstConcreteK u / AstConcreteK v = AstConcreteK $ u / v+  u / v = AstR2K DivideOp u v+  recip (AstFromPrimal u) = AstFromPrimal (recip u)+  recip (AstConcreteK u) = AstConcreteK (recip u)+  recip u = AstR1K RecipOp u+  fromRational r = fromPrimal $ AstConcreteK (fromRational r)++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => Floating (AstTensor ms s (TKScalar r)) where+  pi = error "pi is not defined for tensors"+  exp (AstFromPrimal u) = AstFromPrimal $ exp u+  exp (AstConcreteK u) = AstConcreteK $ exp u+  exp u = AstR1K ExpOp u+  log (AstFromPrimal u) = AstFromPrimal $ log u+  log (AstConcreteK u) = AstConcreteK $ log u+  log u = AstR1K LogOp u+  sqrt (AstFromPrimal u) = AstFromPrimal $ sqrt u+  sqrt (AstConcreteK u) = AstConcreteK $ sqrt u+  sqrt u = AstR1K SqrtOp u+  (AstFromPrimal u) ** (AstFromPrimal v) = AstFromPrimal $ u ** v+  (AstConcreteK u) ** (AstConcreteK v) = AstConcreteK $ u ** v+  u ** v = AstR2K PowerOp u v+  logBase (AstFromPrimal u) (AstFromPrimal v) = AstFromPrimal $ logBase u v+  logBase (AstConcreteK u) (AstConcreteK v) = AstConcreteK $ logBase u v+  logBase u v = AstR2K LogBaseOp u v+  sin (AstFromPrimal u) = AstFromPrimal $ sin u+  sin (AstConcreteK u) = AstConcreteK $ sin u+  sin u = AstR1K SinOp u+  cos (AstFromPrimal u) = AstFromPrimal $ cos u+  cos (AstConcreteK u) = AstConcreteK $ cos u+  cos u = AstR1K CosOp u+  tan (AstFromPrimal u) = AstFromPrimal $ tan u+  tan (AstConcreteK u) = AstConcreteK $ tan u+  tan u = AstR1K TanOp u+  asin (AstFromPrimal u) = AstFromPrimal $ asin u+  asin (AstConcreteK u) = AstConcreteK $ asin u+  asin u = AstR1K AsinOp u+  acos (AstFromPrimal u) = AstFromPrimal $ acos u+  acos (AstConcreteK u) = AstConcreteK $ acos u+  acos u = AstR1K AcosOp u+  atan (AstFromPrimal u) = AstFromPrimal $ atan u+  atan (AstConcreteK u) = AstConcreteK $ atan u+  atan u = AstR1K AtanOp u+  sinh (AstFromPrimal u) = AstFromPrimal $ sinh u+  sinh (AstConcreteK u) = AstConcreteK $ sinh u+  sinh u = AstR1K SinhOp u+  cosh (AstFromPrimal u) = AstFromPrimal $ cosh u+  cosh (AstConcreteK u) = AstConcreteK $ cosh u+  cosh u = AstR1K CoshOp u+  tanh (AstFromPrimal u) = AstFromPrimal $ tanh u+  tanh (AstConcreteK u) = AstConcreteK $ tanh u+  tanh u = AstR1K TanhOp u+  asinh (AstFromPrimal u) = AstFromPrimal $ asinh u+  asinh (AstConcreteK u) = AstConcreteK $ asinh u+  asinh u = AstR1K AsinhOp u+  acosh (AstFromPrimal u) = AstFromPrimal $ acosh u+  acosh (AstConcreteK u) = AstConcreteK $ acosh u+  acosh u = AstR1K AcoshOp u+  atanh (AstFromPrimal u) = AstFromPrimal $ atanh u+  atanh (AstConcreteK u) = AstConcreteK $ atanh u+  atanh u = AstR1K AtanhOp u++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => RealFloatH (AstTensor ms s (TKScalar r)) where+  atan2H (AstFromPrimal u) (AstFromPrimal v) = AstFromPrimal $ atan2H u v+  atan2H (AstConcreteK u) (AstConcreteK v) = AstConcreteK $ atan2H u v+  atan2H u v = AstR2K Atan2Op u v+++-- * Unlawful numeric instances for ranked AST; lawful modulo evaluation++instance (GoodScalar r, AstSpan s)+         => Num (AstTensor ms s (TKR n r)) where+  (+) = liftRFromS2 (+)+  (-) = liftRFromS2 (-)+  (*) = liftRFromS2 (*)+  negate = liftRFromS1 negate+  abs = liftRFromS1 abs+  signum = liftRFromS1 signum+  fromInteger i = error $ "fromInteger is not defined for ranked tensors: "+                          ++ show i++instance (GoodScalar r, IntegralH r, Nested.IntElt r, AstSpan s)+         => IntegralH (AstTensor ms s (TKR n r)) where+  quotH = liftRFromS2 quotH+  remH = liftRFromS2 remH++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => Fractional (AstTensor ms s (TKR n r)) where+  (/) = liftRFromS2 (/)+  recip = liftRFromS1 recip+  fromRational r = error $ "fromRational is not defined for ranked tensors: "+                           ++ show r++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => Floating (AstTensor ms s (TKR n r)) where+  pi = error "pi is not defined for tensors"+  exp = liftRFromS1 exp+  log = liftRFromS1 log+  sqrt = liftRFromS1 sqrt+  (**) = liftRFromS2 (**)+  logBase = liftRFromS2 logBase+  sin = liftRFromS1 sin+  cos = liftRFromS1 cos+  tan = liftRFromS1 tan+  asin = liftRFromS1 asin+  acos = liftRFromS1 acos+  atan = liftRFromS1 atan+  sinh = liftRFromS1 sinh+  cosh = liftRFromS1 cosh+  tanh = liftRFromS1 tanh+  asinh = liftRFromS1 asinh+  acosh = liftRFromS1 acosh+  atanh = liftRFromS1 atanh++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => RealFloatH (AstTensor ms s (TKR n r)) where+  atan2H = liftRFromS2 atan2H+++-- * Unlawful numeric instances for shaped AST; lawful modulo evaluation++instance (GoodScalar r, AstSpan s)+         => Num (AstTensor ms s (TKS sh r)) where+  AstReplicate snat stk@STKS{} u + AstReplicate _ STKS{} v =+    AstReplicate snat stk $ u + v+  AstFromPrimal u + AstFromPrimal v = AstFromPrimal $ u + v+  AstFromDual u + AstFromDual v = AstFromDual $ u + v+  AstConvert c u + AstConvert _ v+    | FTKS ZSS x <- convertFTK c (ftkAst u)+    , Just Refl <- matchingFTK x (ftkAst u)+    , Just Refl <- matchingFTK x (ftkAst v) =+      AstConvert c $ u + v+  AstConcreteS z + u | Just 0 <- sunReplicateScal z = u+  u + AstConcreteS z | Just 0 <- sunReplicateScal z = u+  AstConcreteS n + AstConcreteS k = AstConcreteS (n + k)+  AstConcreteS n + AstPlusS (AstConcreteS k) u =+    AstPlusS (AstConcreteS (n + k)) u+  AstPlusS (AstConcreteS n) u + AstConcreteS k =+    AstPlusS (AstConcreteS (n + k)) u+  AstPlusS (AstConcreteS n) u + AstPlusS (AstConcreteS k) v =+    AstPlusS (AstConcreteS (n + k)) (AstPlusS u v)++--  AstN1S NegateOp (AstVar var) + AstVar var'+--    | var == var' = 0+  AstN1S NegateOp (AstVar var) + AstPlusS (AstVar var') u+    | var == var' = u+--  AstVar var' + AstN1S NegateOp (AstVar var)+--    | var == var' = 0+  AstVar var' + AstPlusS (AstN1S NegateOp (AstVar var)) u+    | var == var' = u++  AstPlusS u@AstConcreteS{} v + w = AstPlusS u (AstPlusS v w)+  u + v@AstConcreteS{} = AstPlusS v u+  u + AstPlusS v@AstConcreteS{} w = AstPlusS v (AstPlusS u w)+  u + v = AstPlusS u v++  AstReplicate snat stk@STKS{} u * AstReplicate _ STKS{} v =+    AstReplicate snat stk $ u * v+  AstFromPrimal u * AstFromPrimal v = AstFromPrimal $ u * v+--  AstFromDual{} * AstFromDual{} = 0+  AstConcreteS z * _ | Just 0 <- sunReplicateScal z = AstConcreteS z+  _ * AstConcreteS z | Just 0 <- sunReplicateScal z = AstConcreteS z+  AstConcreteS s * u | Just 1 <- sunReplicateScal s = u+  u * AstConcreteS s | Just 1 <- sunReplicateScal s = u+  AstConvert c u * AstConvert _ v+    | FTKS ZSS x <- convertFTK c (ftkAst u)+    , Just Refl <- matchingFTK x (ftkAst u)+    , Just Refl <- matchingFTK x (ftkAst v) =+      AstConvert c $ u * v+  AstConcreteS n * AstConcreteS k = AstConcreteS (n * k)+  AstConcreteS n * AstTimesS (AstConcreteS k) u =+    AstTimesS (AstConcreteS (n * k)) u+  AstTimesS (AstConcreteS n) u * AstConcreteS k =+    AstTimesS (AstConcreteS (n * k)) u+  AstTimesS (AstConcreteS n) u * AstTimesS (AstConcreteS k) v =+    AstTimesS (AstConcreteS (n * k)) (AstTimesS u v)++  u@AstConcreteS{} * AstPlusS v w = AstPlusS (u * v) (u * w)+  AstTimesS u@AstConcreteS{} x * AstPlusS v w =+    AstTimesS x (AstPlusS (u * v) (u * w))+  AstPlusS v w * u@AstConcreteS{} = AstPlusS (v * u) (w * u)+  AstPlusS v w * AstTimesS u@AstConcreteS{} x =+    AstTimesS (AstPlusS (v * u) (w * u)) x++  AstN1S NegateOp u * AstN1S NegateOp v = AstTimesS u v++  AstTimesS u@AstConcreteS{} v * w = AstTimesS u (AstTimesS v w)+  u * v@AstConcreteS{} = AstTimesS v u+  u * AstTimesS v@AstConcreteS{} w = AstTimesS v (AstTimesS u w)+  u * v = AstTimesS u v++  negate (AstReplicate snat stk@STKS{} u) = AstReplicate snat stk (negate u)+  negate (AstCond b n k) = AstCond b (negate n) (negate k)+-- TODO: negate (AstBuild1 k stk (var, v)) = AstBuild1 k stk (var, negate v)+  negate (AstLet var n k) = AstLet var n (negate k)+  negate (AstFromPrimal n) = AstFromPrimal (negate n)+  negate (AstFromDual n) = AstFromDual (negate n)+  negate (AstPlusS u v) = AstPlusS (negate u) (negate v)+  negate (AstTimesS u v) = AstTimesS (negate u) v+  negate (AstN1S NegateOp u) = u+  negate (AstN1S SignumOp u) = AstN1S SignumOp (negate u)+  negate (AstI2S QuotOp u v) = AstI2S QuotOp (negate u) v+    -- v is likely positive and let's keep it so+  negate (AstI2S RemOp u v) = AstI2S RemOp (negate u) v+    -- v is likely positive and let's keep it so+  negate (AstConcreteS n) = AstConcreteS (negate n)+  negate (AstGatherS @shm @shn @shp shn v (vars, ix)) =+    AstGatherS @shm @shn @shp shn (negate v) (vars, ix)+  negate (AstConvert c n)+    | FTKS ZSS x <- convertFTK c (ftkAst n)+    , Just Refl <- matchingFTK x (ftkAst n) =+      AstConvert c (negate n)+  negate u = AstN1S NegateOp u+  abs (AstReplicate snat stk@STKS{} u) = AstReplicate snat stk (abs u)+  abs (AstFromPrimal n) = AstFromPrimal (abs n)+  abs (AstFromDual n) = AstFromDual (abs n)+  abs (AstN1S AbsOp u) = AstN1S AbsOp u+  abs (AstConcreteS u) = AstConcreteS (abs u)+  abs (AstConvert c n)+    | FTKS ZSS x <- convertFTK c (ftkAst n)+    , Just Refl <- matchingFTK x (ftkAst n) =+      AstConvert c (abs n)+  abs (AstN1S NegateOp u) = abs u+  abs u = AstN1S AbsOp u+  signum (AstReplicate snat stk@STKS{} u) = AstReplicate snat stk (signum u)+  signum (AstFromPrimal n) = AstFromPrimal (signum n)+  signum (AstFromDual n) = AstFromDual (signum n)+  signum (AstN1S SignumOp u) = AstN1S SignumOp u+  signum (AstConcreteS u) = AstConcreteS (signum u)+  signum (AstConvert c n)+    | FTKS ZSS x <- convertFTK c (ftkAst n)+    , Just Refl <- matchingFTK x (ftkAst n) =+      AstConvert c (signum n)+  signum u = AstN1S SignumOp u+  fromInteger i = error $ "fromInteger is not defined for shaped tensors: "+                          ++ show i++instance (GoodScalar r, IntegralH r, Nested.IntElt r, AstSpan s)+         => IntegralH (AstTensor ms s (TKS sh r)) where+  quotH (AstReplicate snat stk@STKS{} u) (AstReplicate _ STKS{} v) =+    AstReplicate snat stk $ quotH u v+  quotH (AstFromPrimal n) (AstFromPrimal k) = AstFromPrimal (quotH n k)+  quotH (AstConvert c n) (AstConvert _ k)+    | FTKS ZSS x <- convertFTK c (ftkAst n)+    , Just Refl <- matchingFTK x (ftkAst n)+    , Just Refl <- matchingFTK x (ftkAst k) =+      AstConvert c (quotH n k)+  quotH (AstConcreteS n) (AstConcreteS k) = AstConcreteS (quotH n k)+  quotH (AstConcreteS z) _ | Just 0 <- sunReplicateScal z = AstConcreteS z+  quotH u (AstConcreteS s) | Just 1 <- sunReplicateScal s = u+  quotH (AstI2S QuotOp u v) w = quotH u (v * w)+  quotH u v = AstI2S QuotOp u v++  remH (AstReplicate snat stk@STKS{} u) (AstReplicate _ STKS{} v) =+    AstReplicate snat stk $ remH u v+  remH (AstFromPrimal n) (AstFromPrimal k) = AstFromPrimal (remH n k)+  remH (AstConvert c n) (AstConvert _ k)+    | FTKS ZSS x <- convertFTK c (ftkAst n)+    , Just Refl <- matchingFTK x (ftkAst n)+    , Just Refl <- matchingFTK x (ftkAst k) =+      AstConvert c (remH n k)+  remH (AstConcreteS n) (AstConcreteS k) = AstConcreteS (remH n k)+  remH (AstConcreteS z) _ | Just 0 <- sunReplicateScal z = AstConcreteS z+--  remH _ (AstConcreteS s) | Just 1 <- sunReplicateScal s = AstConcreteS 0+  remH u v = AstI2S RemOp u v++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => Fractional (AstTensor ms s (TKS sh r)) where+  AstFromPrimal u / AstFromPrimal v = AstFromPrimal $ u / v+  AstConcreteS u / AstConcreteS v = AstConcreteS $ u / v+  u / v = AstR2S DivideOp u v+  recip (AstFromPrimal u) = AstFromPrimal (recip u)+  recip (AstConcreteS u) = AstConcreteS (recip u)+  recip u = AstR1S RecipOp u+  fromRational r = error $ "fromRational is not defined for shaped tensors: "+                           ++ show r++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => Floating (AstTensor ms s (TKS sh r)) where+  pi = error "pi is not defined for tensors"+  exp (AstFromPrimal u) = AstFromPrimal $ exp u+  exp (AstConcreteS u) = AstConcreteS $ exp u+  exp u = AstR1S ExpOp u+  log (AstFromPrimal u) = AstFromPrimal $ log u+  log (AstConcreteS u) = AstConcreteS $ log u+  log u = AstR1S LogOp u+  sqrt (AstFromPrimal u) = AstFromPrimal $ sqrt u+  sqrt (AstConcreteS u) = AstConcreteS $ sqrt u+  sqrt u = AstR1S SqrtOp u+  (AstFromPrimal u) ** (AstFromPrimal v) = AstFromPrimal $ u ** v+  (AstConcreteS u) ** (AstConcreteS v) = AstConcreteS $ u ** v+  u ** v = AstR2S PowerOp u v+  logBase (AstFromPrimal u) (AstFromPrimal v) = AstFromPrimal $ logBase u v+  logBase (AstConcreteS u) (AstConcreteS v) = AstConcreteS $ logBase u v+  logBase u v = AstR2S LogBaseOp u v+  sin (AstFromPrimal u) = AstFromPrimal $ sin u+  sin (AstConcreteS u) = AstConcreteS $ sin u+  sin u = AstR1S SinOp u+  cos (AstFromPrimal u) = AstFromPrimal $ cos u+  cos (AstConcreteS u) = AstConcreteS $ cos u+  cos u = AstR1S CosOp u+  tan (AstFromPrimal u) = AstFromPrimal $ tan u+  tan (AstConcreteS u) = AstConcreteS $ tan u+  tan u = AstR1S TanOp u+  asin (AstFromPrimal u) = AstFromPrimal $ asin u+  asin (AstConcreteS u) = AstConcreteS $ asin u+  asin u = AstR1S AsinOp u+  acos (AstFromPrimal u) = AstFromPrimal $ acos u+  acos (AstConcreteS u) = AstConcreteS $ acos u+  acos u = AstR1S AcosOp u+  atan (AstFromPrimal u) = AstFromPrimal $ atan u+  atan (AstConcreteS u) = AstConcreteS $ atan u+  atan u = AstR1S AtanOp u+  sinh (AstFromPrimal u) = AstFromPrimal $ sinh u+  sinh (AstConcreteS u) = AstConcreteS $ sinh u+  sinh u = AstR1S SinhOp u+  cosh (AstFromPrimal u) = AstFromPrimal $ cosh u+  cosh (AstConcreteS u) = AstConcreteS $ cosh u+  cosh u = AstR1S CoshOp u+  tanh (AstFromPrimal u) = AstFromPrimal $ tanh u+  tanh (AstConcreteS u) = AstConcreteS $ tanh u+  tanh u = AstR1S TanhOp u+  asinh (AstFromPrimal u) = AstFromPrimal $ asinh u+  asinh (AstConcreteS u) = AstConcreteS $ asinh u+  asinh u = AstR1S AsinhOp u+  acosh (AstFromPrimal u) = AstFromPrimal $ acosh u+  acosh (AstConcreteS u) = AstConcreteS $ acosh u+  acosh u = AstR1S AcoshOp u+  atanh (AstFromPrimal u) = AstFromPrimal $ atanh u+  atanh (AstConcreteS u) = AstConcreteS $ atanh u+  atanh u = AstR1S AtanhOp u++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => RealFloatH (AstTensor ms s (TKS sh r)) where+  atan2H (AstFromPrimal u) (AstFromPrimal v) = AstFromPrimal $ atan2H u v+  atan2H (AstConcreteS u) (AstConcreteS v) = AstConcreteS $ atan2H u v+  atan2H u v = AstR2S Atan2Op u v+++-- * Unlawful numeric instances for mixed AST; lawful modulo evaluation++instance (GoodScalar r, AstSpan s)+         => Num (AstTensor ms s (TKX sh r)) where+  (+) = liftXFromS2 (+)+  (-) = liftXFromS2 (-)+  (*) = liftXFromS2 (*)+  negate = liftXFromS1 negate+  abs = liftXFromS1 abs+  signum = liftXFromS1 signum+  fromInteger i = error $ "fromInteger is not defined for mixed tensors: "+                          ++ show i++instance (GoodScalar r, IntegralH r, Nested.IntElt r, AstSpan s)+         => IntegralH (AstTensor ms s (TKX sh r)) where+  quotH = liftXFromS2 quotH+  remH = liftXFromS2 remH++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => Fractional (AstTensor ms s (TKX sh r)) where+  (/) = liftXFromS2 (/)+  recip = liftXFromS1 recip+  fromRational r = error $ "fromRational is not defined for mixed tensors: "+                           ++ show r++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => Floating (AstTensor ms s (TKX sh r)) where+  pi = error "pi is not defined for tensors"+  exp = liftXFromS1 exp+  log = liftXFromS1 log+  sqrt = liftXFromS1 sqrt+  (**) = liftXFromS2 (**)+  logBase = liftXFromS2 logBase+  sin = liftXFromS1 sin+  cos = liftXFromS1 cos+  tan = liftXFromS1 tan+  asin = liftXFromS1 asin+  acos = liftXFromS1 acos+  atan = liftXFromS1 atan+  sinh = liftXFromS1 sinh+  cosh = liftXFromS1 cosh+  tanh = liftXFromS1 tanh+  asinh = liftXFromS1 asinh+  acosh = liftXFromS1 acosh+  atanh = liftXFromS1 atanh++instance (GoodScalar r, RealFloatH r, Nested.FloatElt r, AstSpan s)+         => RealFloatH (AstTensor ms s (TKX sh r)) where+  atan2H = liftXFromS2 atan2H+++-- * Unlawful instances of AST for bool; they are lawful modulo evaluation++-- Simple variable comparisons, if any, come first.+instance Boolean (AstBool ms) where+  true = AstBoolConst True+  false = AstBoolConst False+  notB (AstBoolConst b) = AstBoolConst $ not b+  notB (AstBoolNot b) = b+  notB b = AstBoolNot b+  AstBoolConst True &&* b = b+  AstBoolConst False &&* _b = AstBoolConst False+  b &&* AstBoolConst True = b+  _b &&* AstBoolConst False = AstBoolConst False+  AstBoolAnd b c &&* d = b &&* (c &&* d)+  b@(AstLeqK AstConcreteK{} AstVar{}) &&* c = AstBoolAnd b c+  b@(AstLeqK AstConcreteK{} (AstN1K NegateOp+                                    AstVar{})) &&* c = AstBoolAnd b c+  b@(AstBoolNot+       (AstLeqK AstConcreteK{} AstVar{})) &&* c = AstBoolAnd b c+  b@(AstBoolNot+       (AstLeqK AstConcreteK{} (AstN1K NegateOp+                                       AstVar{}))) &&* c = AstBoolAnd b c+  b@(AstBoolNot+       (AstBoolAnd (AstLeqK AstConcreteK{} AstVar{}) _)) &&* c = AstBoolAnd b c+  b@(AstBoolNot+       (AstBoolAnd+          (AstLeqK AstConcreteK{}+                   (AstN1K NegateOp AstVar{})) _)) &&* c = AstBoolAnd b c+  b@(AstBoolNot+       (AstBoolAnd (AstBoolNot (AstLeqK AstConcreteK{}+                                        AstVar{})) _)) &&* c = AstBoolAnd b c+  b@(AstBoolNot+       (AstBoolAnd+          (AstBoolNot+             (AstLeqK AstConcreteK{}+                      (AstN1K NegateOp AstVar{}))) _)) &&* c = AstBoolAnd b c+  b &&* c@(AstLeqK AstConcreteK{} AstVar{}) = AstBoolAnd c b+  b &&* c@(AstLeqK AstConcreteK{} (AstN1K NegateOp+                                          AstVar{})) = AstBoolAnd c b+  b &&* c@(AstBoolNot+             (AstLeqK AstConcreteK{} AstVar{})) = AstBoolAnd c b+  b &&* c@(AstBoolNot+             (AstLeqK AstConcreteK{} (AstN1K NegateOp+                                             AstVar{}))) = AstBoolAnd c b+  b &&* c@(AstBoolNot+             (AstBoolAnd (AstLeqK AstConcreteK{} AstVar{}) _)) = AstBoolAnd c b+  b &&* c@(AstBoolNot+             (AstBoolAnd+                (AstLeqK AstConcreteK{}+                         (AstN1K NegateOp AstVar{})) _)) = AstBoolAnd c b+  b &&* c@(AstBoolNot+             (AstBoolAnd (AstBoolNot (AstLeqK AstConcreteK{}+                                              AstVar{})) _)) = AstBoolAnd c b+  b &&* c@(AstBoolNot+             (AstBoolAnd+                (AstBoolNot+                   (AstLeqK AstConcreteK{}+                            (AstN1K NegateOp AstVar{}))) _)) = AstBoolAnd c b+  b &&* AstBoolAnd+          c@(AstLeqK AstConcreteK{} AstVar{}) d = AstBoolAnd c (b &&* d)+  b &&* AstBoolAnd+          c@(AstLeqK AstConcreteK{}+                     (AstN1K NegateOp AstVar{})) d = AstBoolAnd c (b &&* d)+  b &&* AstBoolAnd+          c@(AstBoolNot (AstLeqK AstConcreteK{}+                                 AstVar{})) d = AstBoolAnd c (b &&* d)+  b &&* AstBoolAnd+          c@(AstBoolNot+               (AstLeqK AstConcreteK{}+                        (AstN1K NegateOp AstVar{}))) d = AstBoolAnd c (b &&* d)+  b &&* AstBoolAnd+          c@(AstBoolNot+               (AstBoolAnd (AstLeqK AstConcreteK{}+                                    AstVar{}) _)) d = AstBoolAnd c (b &&* d)+  b &&* AstBoolAnd+          c@(AstBoolNot+               (AstBoolAnd+                  (AstLeqK AstConcreteK{}+                           (AstN1K NegateOp+                                   AstVar{})) _)) d = AstBoolAnd c (b &&* d)+  b &&* AstBoolAnd+          c@(AstBoolNot+               (AstBoolAnd+                  (AstBoolNot+                     (AstLeqK AstConcreteK{}+                              AstVar{})) _)) d = AstBoolAnd c (b &&* d)+  b &&* AstBoolAnd+          c@(AstBoolNot+               (AstBoolAnd+                  (AstBoolNot+                     (AstLeqK AstConcreteK{}+                              (AstN1K+                                 NegateOp+                                 AstVar{}))) _)) d = AstBoolAnd c (b &&* d)+  b &&* c = AstBoolAnd b c+  b ||* c = notB (notB b &&* notB c)++-- TODO: refactor with something like liftRFromS2+instance (AstSpan s, GoodScalar r) => EqH (AstTensor ms s) (TKR n r) where+  v ==. u = case ftkAst v of+    FTKR shv' _ -> case ftkAst u of+      FTKR shu' _ ->+        withShsFromShR shv' $ \shv ->+          withShsFromShR shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                cAstSFromR shu (primalPart v)+                ==. cAstSFromR shv (primalPart u)+              _ -> error $ "(==.): shapes don't match: "+                           ++ show (shu, shv)++instance (AstSpan s, GoodScalar r) => EqH (AstTensor ms s) (TKX sh r) where+  v ==. u = case ftkAst v of+    FTKX shv' _ -> case ftkAst u of+      FTKX shu' _ ->+        withShsFromShX shv' $ \shv ->+          withShsFromShX shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                cAstSFromX shu (primalPart v)+                ==. cAstSFromX shv (primalPart u)+              _ -> error $ "(==.): shapes don't match: "+                           ++ show (shu, shv)++instance (AstSpan s, GoodScalar r) => OrdH (AstTensor ms s) (TKR n r) where+  v <=. u = case ftkAst v of+    FTKR shv' _ -> case ftkAst u of+      FTKR shu' _ ->+        withShsFromShR shv' $ \shv ->+          withShsFromShR shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                cAstSFromR shu (primalPart v)+                <=. cAstSFromR shv (primalPart u)+              _ -> error $ "(<=.): shapes don't match: "+                           ++ show (shu, shv)++instance (AstSpan s, GoodScalar r) => OrdH (AstTensor ms s) (TKX sh r) where+  v <=. u = case ftkAst v of+    FTKX shv' _ -> case ftkAst u of+      FTKX shu' _ ->+        withShsFromShX shv' $ \shv ->+          withShsFromShX shu' $ \shu ->+            case testEquality shv shu of+              Just Refl ->+                cAstSFromX shu (primalPart v)+                <=. cAstSFromX shv (primalPart u)+              _ -> error $ "(<=.): shapes don't match: "+                           ++ show (shu, shv)++-- TODO: share u and v, since they are duplicated here+instance (AstSpan s, GoodScalar r)+         => EqH (AstTensor ms s) (TKScalar r) where+  v ==. u = v <=. u &&* u <=. v+  {- TODO: for this to work, booleans have to be first-class:+  vUnshared ==. uUnshared = astLetFunNoSimplify vUnshared $ \v ->+                            astLetFunNoSimplify uUnshared $ \u ->+    v <=. u &&* u <=. v -}++instance (AstSpan s, GoodScalar r)+         => EqH (AstTensor ms s) (TKS sh r) where+  v ==. u = v <=. u &&* u <=. v++-- These are common in indexing, so worth optimizing early via AstConcrete.+-- We keep AstConcrete on the left, as with AstPlusK and others.+instance (AstSpan s, GoodScalar r)+         => OrdH (AstTensor ms s) (TKScalar r) where+  u <=. v | let (u1, u2) = bounds u+                (v1, v2) = bounds v+          , u2 <= v1 || u1 > v2 = AstBoolConst (u2 <= v1)+  AstFromPrimal u <=. AstFromPrimal v = u <=. v+  AstPrimalPart u <=. AstPrimalPart v = u <=. v+  AstFromDual{} <=. AstFromDual{} = AstBoolConst True+  AstConvert _ u <=. AstConvert _ v+    | FTKS ZSS (FTKScalar @ru) <- ftkAst u+    , FTKS ZSS (FTKScalar @rv) <- ftkAst v+    , Just Refl <- testEquality (typeRep @ru) (typeRep @rv)+    = u <=. v+  AstConcreteK u <=. AstConvert _ v+    | FTKS ZSS (FTKScalar @rv) <- ftkAst v+    , Just Refl <- testEquality (typeRep @rv) (typeRep @r)+    = AstConcreteS (unConcrete $ sfromK $ Concrete u) <=. v+  u <=. AstPlusK (AstConcreteK v) w =+    u - AstConcreteK v <=. w+  AstPlusK (AstConcreteK u) w <=. v =+    AstConcreteK u <=. v - w+  u <=. AstConcreteK v =+    AstConcreteK (negate v) <=. negate u+  AstConcreteK u <=. AstTimesK (AstConcreteK v) w+    | v > 0 && u >= 0+    , Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =+      AstConcreteK ((u + v - 1) `quotH` v) <=. w -- 10 == 5 * 2, 11 > 5 * 2+  AstConcreteK u <=. AstTimesK (AstConcreteK v) w+    | v > 0 && u < 0+    , Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =+      AstConcreteK (u `quotH` v) <=. w  -- -10 == 5 * -2, -9 > 5 * -2+  AstConcreteK u <=. AstTimesK (AstConcreteK v) w+    | v < 0+    , Just Refl <- testEquality (typeRep @r) (typeRep @Int64) =+      AstConcreteK u <=. AstTimesK (AstConcreteK $ negate v) (AstN1K NegateOp w)+  v@AstConcreteK{} <=. u =+    AstLeqK (primalPart v) (primalPart u)+  v <=. u =+    AstConcreteK 0 <=. primalPart u - primalPart v++instance (AstSpan s, GoodScalar r)+         => OrdH (AstTensor ms s) (TKS sh r) where+  AstFromPrimal u <=. AstFromPrimal v = u <=. v+  AstFromDual{} <=. AstFromDual{} = AstBoolConst True+  AstPrimalPart u <=. AstPrimalPart v = u <=. v+  AstConvert c u <=. AstConvert _ v+    | FTKS ZSS x <- convertFTK c (ftkAst u)+    , Just Refl <- matchingFTK x (ftkAst u)+    , Just Refl <- matchingFTK x (ftkAst v) = u <=. v+  AstConcreteS u <=. AstConvert c v+    | FTKS ZSS (FTKScalar @rz) <- convertFTK c (ftkAst v)+    , FTKScalar @ry <- ftkAst v+    , Just Refl <- testEquality (typeRep @ry) (typeRep @rz) =+      AstConcreteK (unConcrete $ kfromS $ Concrete u) <=. v+  AstConcreteS u <=. AstConcreteS v =+    AstBoolConst $ Concrete @(TKS sh r) u <=. Concrete v+  u <=. AstPlusS (AstConcreteS v) w =+    u - AstConcreteS v <=. w+  AstPlusS (AstConcreteS u) w <=. v =+    AstConcreteS u <=. v - w+  u <=. AstConcreteS v =+    AstConcreteS (negate v) <=. negate u+  AstVar u <=. AstVar v | u == v =+    AstBoolConst True+  AstConvert _ (AstVar u) <=. AstConvert _ (AstVar v)+    | varNameToAstVarId u == varNameToAstVarId v =+      AstBoolConst True+  v <=. u = AstLeqS (primalPart v) (primalPart u)+++-- * AstRaw, AstNoVectorize and AstNoSimplify definitions++-- | An AST variant that doesn't vectorize terms and also builds them+-- with ordinary, non-simplifying constructors. It's based on sharing+-- rather than lets and commonly used as the instance for primals+-- inside ADVal and, consequently, used for evaluating delta expressions.+type role AstRaw nominal nominal+newtype AstRaw s y =+  AstRaw {unAstRaw :: AstTensor AstMethodShare s y}+ deriving Show++-- | An AST variant for testing that doesn't vectorize terms, but still+-- builds them using simplifying smart constructors.+type role AstNoVectorize nominal nominal+newtype AstNoVectorize s y =+  AstNoVectorize {unAstNoVectorize :: AstTensor AstMethodLet s y}+ deriving Show++-- | An AST variant for testing that vectorizes terms, but builds them+-- with ordinary, non-simplifying constructors.+type role AstNoSimplify nominal nominal+newtype AstNoSimplify s y =+  AstNoSimplify {unAstNoSimplify :: AstTensor AstMethodLet s y}+ deriving Show+++-- * AstRaw, AstNoVectorize and AstNoSimplify type family instances++type instance PrimalOf (AstRaw s) = AstRaw PrimalSpan+type instance DualOf (AstRaw s) = AstTensor AstMethodShare DualSpan+type instance ShareOf (AstRaw s) = AstRaw s+type instance HFunOf (AstRaw s) x y = AstHFun s s x y+type instance BoolOf (AstRaw s) = AstBool AstMethodShare++type instance PrimalOf (AstNoVectorize s) = AstNoVectorize PrimalSpan+type instance DualOf (AstNoVectorize s) = AstTensor AstMethodLet DualSpan+type instance ShareOf (AstNoVectorize s) = AstRaw s+type instance HFunOf (AstNoVectorize s) x z = AstHFun s s x z+type instance BoolOf (AstNoVectorize s) = AstBool AstMethodLet++type instance PrimalOf (AstNoSimplify s) = AstNoSimplify PrimalSpan+type instance DualOf (AstNoSimplify s) = AstTensor AstMethodLet DualSpan+type instance ShareOf (AstNoSimplify s) = AstRaw s+type instance HFunOf (AstNoSimplify s) x z = AstHFun s s x z+type instance BoolOf (AstNoSimplify s) = AstBool AstMethodLet+++-- * AstRaw, AstNoVectorize and AstNoSimplify other instances++instance EqH (AstTensor AstMethodShare s) y => EqH (AstRaw s) y where+  AstRaw v ==. AstRaw u = v ==. u+instance OrdH (AstTensor AstMethodShare s) y => OrdH (AstRaw s) y where+  AstRaw v <=. AstRaw u = v <=. u++deriving instance Eq (AstRaw s y)+deriving instance Ord (AstRaw s y)+deriving instance Num (AstTensor AstMethodShare s y) => Num (AstRaw s y)+deriving instance IntegralH (AstTensor AstMethodShare s y)+                  => IntegralH (AstRaw s y)+deriving instance Fractional (AstTensor AstMethodShare s y)+                  => Fractional (AstRaw s y)+deriving instance Floating (AstTensor AstMethodShare s y)+                  => Floating (AstRaw s y)+deriving instance RealFloatH (AstTensor AstMethodShare s y)+                  => RealFloatH (AstRaw s y)++instance EqH (AstTensor AstMethodLet s) y => EqH (AstNoVectorize s) y where+  AstNoVectorize v ==. AstNoVectorize u = v ==. u+instance OrdH (AstTensor AstMethodLet s) y => OrdH (AstNoVectorize s) y where+  AstNoVectorize v <=. AstNoVectorize u = v <=. u+deriving instance Eq (AstNoVectorize s y)+deriving instance Ord (AstNoVectorize s y)+deriving instance Num (AstTensor AstMethodLet s y) => Num (AstNoVectorize s y)+deriving instance (IntegralH (AstTensor AstMethodLet s y))+                  => IntegralH (AstNoVectorize s y)+deriving instance Fractional (AstTensor AstMethodLet s y)+                  => Fractional (AstNoVectorize s y)+deriving instance Floating (AstTensor AstMethodLet s y)+                  => Floating (AstNoVectorize s y)+deriving instance (RealFloatH (AstTensor AstMethodLet s y))+                  => RealFloatH (AstNoVectorize s y)++instance EqH (AstTensor AstMethodLet s) y => EqH (AstNoSimplify s) y where+  AstNoSimplify v ==. AstNoSimplify u = v ==. u+instance OrdH (AstTensor AstMethodLet s) y => OrdH (AstNoSimplify s) y where+  AstNoSimplify v <=. AstNoSimplify u = v <=. u+deriving instance Eq (AstNoSimplify s y)+deriving instance Ord (AstNoSimplify s y)+deriving instance Num (AstTensor AstMethodLet s y) => Num (AstNoSimplify s y)+deriving instance (IntegralH (AstTensor AstMethodLet s y))+                  => IntegralH (AstNoSimplify s y)+deriving instance Fractional (AstTensor AstMethodLet s y)+                  => Fractional (AstNoSimplify s y)+deriving instance Floating (AstTensor AstMethodLet s y)+                  => Floating (AstNoSimplify s y)+deriving instance (RealFloatH (AstTensor AstMethodLet s y))+                  => RealFloatH (AstNoSimplify s y)+++-- * Misc++astLetFunNoSimplify+  :: forall y z s s2. AstSpan s+  => AstTensor AstMethodLet s y+  -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)+  -> AstTensor AstMethodLet s2 z+astLetFunNoSimplify a f | astIsSmall True a = f a+                            -- too important an optimization to skip+astLetFunNoSimplify a f = case a of+  AstFromS' @y2 ftkz v ->+    let (var, ast) = funToAst2 (ftkAst v) Nothing (f . cAstFromS @y2 ftkz)+    in AstLet var v ast+  AstFromPrimal (AstFromS' @y2 ftkz vRaw) ->+    let v = AstFromPrimal vRaw+        (var, ast) = funToAst2 (ftkAst v) Nothing (f . cAstFromS @y2 ftkz)+    in AstLet var v ast+  _ -> case ftkAst a of+    ftk@(FTKR @_ @x2 sh' x) ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        let (var, ast) =+              funToAst2 (FTKS sh x) Nothing+                        (f . cAstFromS @(TKS2 sh x2) ftk)+        in AstLet var (cAstSFromR @sh sh a) ast+             -- safe, because subsitution ruled out above+    ftk@(FTKX @_ @x sh' x) ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        let (var, ast) =+              funToAst2 (FTKS sh x) Nothing+                        (f . cAstFromS @(TKS2 sh x) ftk)+        in AstLet var (cAstSFromX @sh sh a) ast+    -- processing product recursively may be not worth it+    ftk -> let (var, ast) = funToAst2 ftk Nothing f+           in AstLet var a ast++sunReplicateScal :: Nested.Elt a+                 => Nested.Shaped sh a -> Maybe a+sunReplicateScal (Nested.Shaped arr)+  | all (all (== 0) . take (shxLength (Nested.mshape arr)))+        (Mixed.marrayStrides arr)+  , shxSize (Nested.mshape arr) /= 0 =+    Just $ Nested.mindex arr $ ixxZero' $ Nested.mshape arr+sunReplicateScal _ = Nothing++sunReplicate1 :: Nested.Elt a+              => Nested.Shaped (n ': sh) a -> Maybe (Nested.Shaped sh a)+sunReplicate1 a | (snat :$$ _) <- Nested.sshape a =+  sunReplicateN (snat :$$ ZSS) a++sunReplicateN :: Nested.Elt a+              => ShS shm -> Nested.Shaped (shm ++ shn) a+              -> Maybe (Nested.Shaped shn a)+sunReplicateN shm a@(Nested.Shaped arr)+  | all (all (== 0) . take (shsLength shm)) (Mixed.marrayStrides arr)+  , shsSize shm /= 0 =+    Just $ Nested.sindexPartial a $ ixsZero shm+sunReplicateN _ _ = Nothing
+ src/HordeAd/Core/CarriersConcrete.hs view
@@ -0,0 +1,281 @@+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Tensor operations implementation using the ox-arrays package.+-- These definitions, mostly class instances, are needed to make concrete+-- arrays a valid carrier for a tensor class algebra (instance) defined in+-- "HordeAd.Core.OpsConcrete".+module HordeAd.Core.CarriersConcrete+  ( -- * RepConcrete and its operations+    RepConcrete, tftkG, eltDictRep, showDictRep+    -- * Concrete and its operations+  , Concrete(..), rtoVector, stoVector, xtoVector+  ) where++import Prelude hiding (foldl')++import Control.DeepSeq (NFData (..))+import Data.Vector.Storable qualified as VS++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed qualified as Mixed+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked qualified as Ranked+import Data.Array.Nested.Shaped qualified as Shaped+import Data.Array.Nested.Shaped.Shape+import Data.Array.Strided.Orthotope (liftVEltwise1)++import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * Orphan ox-arrays instances++instance (Nested.IntElt r, Nested.PrimElt r, Eq r, Num r)+         => IntegralH (Nested.Ranked n r) where+  -- These can't be partial, because our conditionals are not lazy+  -- and so the counterfactual branches, with zeros, may get executed+  -- even though they are subsequently ignored.+  quotH a b = Nested.rquotArray a (Ranked.liftRanked1 mmakeNonZero b)+  remH a b = Nested.rremArray a (Ranked.liftRanked1 mmakeNonZero b)++instance (Nested.IntElt r, Nested.PrimElt r, Eq r, Num r)+         => IntegralH (Nested.Shaped sh r) where+  quotH a b = Nested.squotArray a (Shaped.liftShaped1 mmakeNonZero b)+  remH a b = Nested.sremArray a (Shaped.liftShaped1 mmakeNonZero b)++instance (Nested.IntElt r, Nested.PrimElt r, Eq r, Num r)+         => IntegralH (Nested.Mixed sh r) where+  quotH a b = Nested.mquotArray a (mmakeNonZero b)+  remH a b = Nested.mremArray a (mmakeNonZero b)++instance GoodScalar r+         => Real (Nested.Ranked n r) where+  toRational = error "toRational is not defined for tensors"++instance GoodScalar r+         => Real (Nested.Shaped sh r) where+  toRational = error "toRational is not defined for tensors"++instance GoodScalar r+         => Real (Nested.Mixed sh r) where+  toRational = error "toRational is not defined for tensors"++instance (GoodScalar r, Nested.FloatElt r)+         => RealFrac (Nested.Ranked n r) where+  properFraction = error "properFraction is not defined for tensors"++instance (GoodScalar r, RealFrac r, Nested.FloatElt r)+         => RealFrac (Nested.Shaped sh r) where+  properFraction = error "properFraction is not defined for tensors"++instance (GoodScalar r, Nested.FloatElt r)+         => RealFrac (Nested.Mixed sh r) where+  properFraction = error "properFraction is not defined for tensors"++instance (Nested.PrimElt r, Nested.FloatElt r)+         => RealFloatH (Nested.Ranked n r) where+  atan2H = Nested.ratan2Array++instance (Nested.PrimElt r, Nested.FloatElt r)+         => RealFloatH (Nested.Shaped sh r) where+  atan2H = Nested.satan2Array++instance (Nested.PrimElt r, Nested.FloatElt r)+         => RealFloatH (Nested.Mixed sh r) where+  atan2H = Nested.matan2Array++instance (GoodScalar r, Nested.PrimElt r, RealFloat r, Nested.FloatElt r)+         => RealFloat (Nested.Ranked n r) where+  atan2 = Nested.ratan2Array+  floatRadix = error "operation not defined for tensors"+  floatDigits = error "operation not defined for tensors"+  floatRange = error "operation not defined for tensors"+  decodeFloat = error "operation not defined for tensors"+  encodeFloat = error "operation not defined for tensors"+  isNaN = error "operation not defined for tensors"+  isInfinite = error "operation not defined for tensors"+  isDenormalized = error "operation not defined for tensors"+  isNegativeZero = error "operation not defined for tensors"+  isIEEE = error "operation not defined for tensors"++instance (GoodScalar r, Nested.PrimElt r, RealFloat r, Nested.FloatElt r)+         => RealFloat (Nested.Shaped sh r) where+  atan2 = Nested.satan2Array+  floatRadix = error "operation not defined for tensors"+  floatDigits = error "operation not defined for tensors"+  floatRange = error "operation not defined for tensors"+  decodeFloat = error "operation not defined for tensors"+  encodeFloat = error "operation not defined for tensors"+  isNaN = error "operation not defined for tensors"+  isInfinite = error "operation not defined for tensors"+  isDenormalized = error "operation not defined for tensors"+  isNegativeZero = error "operation not defined for tensors"+  isIEEE = error "operation not defined for tensors"++instance (GoodScalar r, Nested.PrimElt r, RealFloat r, Nested.FloatElt r)+         => RealFloat (Nested.Mixed sh r) where+  atan2 = Nested.matan2Array+  floatRadix = error "operation not defined for tensors"+  floatDigits = error "operation not defined for tensors"+  floatRange = error "operation not defined for tensors"+  decodeFloat = error "operation not defined for tensors"+  encodeFloat = error "operation not defined for tensors"+  isNaN = error "operation not defined for tensors"+  isInfinite = error "operation not defined for tensors"+  isDenormalized = error "operation not defined for tensors"+  isNegativeZero = error "operation not defined for tensors"+  isIEEE = error "operation not defined for tensors"++-- TODO: make more efficient somehow?+mmakeNonZero :: (Nested.PrimElt r, Eq r, Num r)+             => Nested.Mixed sh r -> Nested.Mixed sh r+mmakeNonZero =+  Mixed.mliftNumElt1+    (`liftVEltwise1` (VS.map (\x -> if x == 0 then 1 else x)))+++-- * RepConcrete and its operations++-- | The type family that represents tensor kinds in concrete arrays.+type family RepConcrete (y :: TK) where+  RepConcrete (TKScalar r) = r+  RepConcrete (TKR2 n x) = Nested.Ranked n (RepConcrete x)+  RepConcrete (TKS2 sh x) = Nested.Shaped sh (RepConcrete x)+  RepConcrete (TKX2 sh x) = Nested.Mixed sh (RepConcrete x)+  RepConcrete (TKProduct x z) = (RepConcrete x, RepConcrete z)++-- | Computing full shape tensor kinds for concrete arrays.+tftkG :: SingletonTK y -> RepConcrete y -> FullShapeTK y+tftkG stk t =+  let repackShapeTree :: SingletonTK y+                      -> Mixed.ShapeTree (RepConcrete y)+                      -> FullShapeTK y+      repackShapeTree stk0 tree = case stk0 of+        STKScalar -> FTKScalar+        STKR _ stk1 -> let (sh, rest) = tree+                       in FTKR sh $ repackShapeTree stk1 rest+        STKS _ stk1 -> let (sh, rest) = tree+                       in FTKS sh $ repackShapeTree stk1 rest+        STKX _ stk1 -> let (sh, rest) = tree+                       in FTKX sh $ repackShapeTree stk1 rest+        STKProduct stk1 stk2 ->+                       let (tree1, tree2) = tree+                       in FTKProduct (repackShapeTree stk1 tree1)+                                     (repackShapeTree stk2 tree2)+  in case stk of+    STKScalar -> FTKScalar+    STKR _ stk1 | Dict <- eltDictRep stk1 ->+      FTKR (Nested.rshape t) $ repackShapeTree stk1+      $ snd $ Mixed.mshapeTree t+    STKS sh stk1 | Dict <- eltDictRep stk1 ->+      FTKS sh $ repackShapeTree stk1+      $ snd $ Mixed.mshapeTree t+    STKX _ stk1 | Dict <- eltDictRep stk1 ->+      FTKX (Nested.mshape t) $ repackShapeTree stk1+      $ snd $ Mixed.mshapeTree t+    STKProduct stk1 stk2 ->+      FTKProduct (tftkG stk1 (fst t))+                 (tftkG stk2 (snd t))++eltDictRep :: SingletonTK y -> Dict Nested.KnownElt (RepConcrete y)+eltDictRep = \case+    STKScalar -> Dict+    STKR SNat x | Dict <- eltDictRep x -> Dict+    STKS sh x | Dict <- eltDictRep x -> withKnownShS sh Dict+    STKX sh x | Dict <- eltDictRep x -> withKnownShX sh Dict+    STKProduct stk1 stk2 | Dict <- eltDictRep stk1+                         , Dict <- eltDictRep stk2 -> Dict++showDictRep :: SingletonTK y -> Dict Show (RepConcrete y)+showDictRep = \case+    STKScalar -> Dict+    STKR _ x | Dict <- showDictRep x+             , Dict <- eltDictRep x -> Dict+    STKS _ x | Dict <- showDictRep x+             , Dict <- eltDictRep x -> Dict+    STKX _ x | Dict <- showDictRep x+             , Dict <- eltDictRep x -> Dict+    STKProduct stk1 stk2 | Dict <- showDictRep stk1+                         , Dict <- showDictRep stk2 -> Dict++nfdataDictRep :: SingletonTK y -> Dict NFData (RepConcrete y)+nfdataDictRep = \case+    STKScalar -> Dict+    STKR _ x | Dict <- nfdataDictRep x+             , Dict <- eltDictRep x -> Dict+    STKS _ x | Dict <- nfdataDictRep x+             , Dict <- eltDictRep x -> Dict+    STKX _ x | Dict <- nfdataDictRep x+             , Dict <- eltDictRep x -> Dict+    STKProduct stk1 stk2 | Dict <- nfdataDictRep stk1+                         , Dict <- nfdataDictRep stk2 -> Dict+++-- * Concrete and its instances++-- | A newtype wrapper over 'RepConcrete'.+-- It's needed because @RepConcrete@ can't be partially applied.+-- This type also lets us work around the woes with defining 'Show'+-- for the @RepConcrete@ type family. It gives us a concrete thing+-- to attach a @Show@ instance to.+type role Concrete nominal+newtype Concrete y = Concrete {unConcrete :: RepConcrete y}++instance KnownSTK y => Show (Concrete y) where+  showsPrec d (Concrete t) | Dict <- showDictRep (knownSTK @y) = showsPrec d t++instance KnownSTK y => NFData (Concrete y) where+  rnf (Concrete t) | Dict <- nfdataDictRep (knownSTK @y) = rnf t++type instance BoolOf Concrete = Bool++type instance HFunOf Concrete x z = RepConcrete x -> RepConcrete z++type instance PrimalOf Concrete = Concrete++type instance DualOf Concrete = DummyDualTarget++type instance ShareOf Concrete = Concrete++instance GoodScalar r => EqH Concrete (TKScalar r) where+  Concrete u ==. Concrete v = u == v++instance GoodScalar r => OrdH Concrete (TKScalar r) where+  Concrete u <=. Concrete v = u <= v++instance GoodScalar r => EqH Concrete (TKR n r) where+  Concrete u ==. Concrete v = u == v++instance GoodScalar r => OrdH Concrete (TKR n r) where+  Concrete u <=. Concrete v = u <= v++instance GoodScalar r => EqH Concrete (TKS sh r) where+  Concrete u ==. Concrete v = u == v++instance GoodScalar r => OrdH Concrete (TKS sh r) where+  Concrete u <=. Concrete v = u <= v++instance GoodScalar r => EqH Concrete (TKX sh r) where+  Concrete u ==. Concrete v = u == v++instance GoodScalar r => OrdH Concrete (TKX sh r) where+  Concrete u <=. Concrete v = u <= v++deriving instance Eq (RepConcrete y) => Eq (Concrete y)+deriving instance Ord (RepConcrete y) => Ord (Concrete y)+deriving instance Num (RepConcrete y) => Num (Concrete y)+deriving instance IntegralH (RepConcrete y) => IntegralH (Concrete y)+deriving instance Real (RepConcrete y) => Real (Concrete y)+deriving instance Fractional (RepConcrete y) => Fractional (Concrete y)+deriving instance Floating (RepConcrete y) => Floating (Concrete y)+deriving instance RealFrac (RepConcrete y) => RealFrac (Concrete y)+deriving instance RealFloatH (RepConcrete y) => RealFloatH (Concrete y)+deriving instance RealFloat (RepConcrete y) => RealFloat (Concrete y)++rtoVector :: GoodScalar r => Concrete (TKR n r) -> VS.Vector r+rtoVector = Nested.rtoVector . unConcrete++stoVector :: GoodScalar r => Concrete (TKS sh r) -> VS.Vector r+stoVector = Nested.stoVector . unConcrete++xtoVector :: GoodScalar r => Concrete (TKX sh r) -> VS.Vector r+xtoVector = Nested.mtoVector . unConcrete
+ src/HordeAd/Core/ConvertTensor.hs view
@@ -0,0 +1,234 @@+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+-- | A class for converting tensors between different forms that contain+-- the same data but varying amounts of shape information.+module HordeAd.Core.ConvertTensor+  ( ConvertTensor(..)+  ) where++import Prelude++import Data.Type.Equality (gcastWith, (:~:) (Refl))+import GHC.TypeLits (KnownNat, Nat, type (+))++import Data.Array.Nested (MapJust, Replicate, type (++))+import Data.Array.Nested.Convert (shxFromShS)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.TensorKind+import HordeAd.Core.Types++class ConvertTensor (target :: Target) where+  -- | The universal conversion function that can emulate all conversion+  -- methods below, but requires an explicit recipe in the first argument.+  --+  -- All conversion operations below could be defined in terms of @tconvert@,+  -- but they'd need additional singleton arguments or constraints+  -- or we'd need to depend on 'HordeAd.Core.Ops.BaseTensor'+  -- to use 'HordeAd.Core.Ops.rshape', etc.+  tconvert :: TKConversion a b -> SingletonTK a -> target a -> target b++  -- | The conversion from a rank 0 ranked tensor to a scalar.+  kfromR :: GoodScalar r => target (TKR 0 r) -> target (TKScalar r)+  kfromR = kfromS . sfromR+  -- | The conversion from an empty shape shaped tensor to a scalar.+  kfromS :: GoodScalar r => target (TKS '[] r) -> target (TKScalar r)+  kfromS = let c = ConvCmp ConvX0 ConvSX+           in tconvert c (STKS ZSS STKScalar)+  kfromX :: GoodScalar r => target (TKX '[] r) -> target (TKScalar r)+  kfromX = kfromS . sfromX+  rfromK :: GoodScalar r => target (TKScalar r) -> target (TKR 0 r)+  rfromK = rfromS . sfromK+  -- | The conversion from a shaped tensor to the corresponding ranked tensor+  -- of the same rank.+  rfromS :: forall sh x. (KnownShS sh, KnownSTK x)+         => target (TKS2 sh x) -> target (TKR2 (Rank sh) x)+  rfromS | Refl <- lemRankMapJust (knownShS @sh) =+    let c = ConvCmp (ConvXR knownSTK) ConvSX+    in tconvert c (STKS knownShS knownSTK)+  rfromX :: forall sh x. KnownSTK x+         => target (TKX2 sh x) -> target (TKR2 (Rank sh) x)+  sfromK :: GoodScalar r => target (TKScalar r) -> target (TKS '[] r)+  sfromK = let c = ConvCmp ConvXS (Conv0X STKScalar)+           in tconvert c STKScalar+  -- | The conversion from a ranked tensor to the corresponding shaped tensor+  -- of the same rank.+  sfromR :: (KnownShS sh, KnownSTK x)+         => target (TKR2 (Rank sh) x) -> target (TKS2 sh x)+  sfromX :: (KnownShS sh, Rank sh ~ Rank sh', KnownSTK x)+         => target (TKX2 sh' x) -> target (TKS2 sh x)+  xfromK :: GoodScalar r => target (TKScalar r) -> target (TKX '[] r)+  xfromK = xfromS . sfromK+  xfromR :: (KnownShX sh', KnownSTK x)+         => target (TKR2 (Rank sh') x) -> target (TKX2 sh' x)+  xfromS :: (KnownShS sh, KnownShX sh', Rank sh ~ Rank sh', KnownSTK x)+         => target (TKS2 sh x) -> target (TKX2 sh' x)++  rzip :: forall y z n. (KnownSTK y, KnownSTK z)+       => target (TKProduct (TKR2 n y) (TKR2 n z))+       -> target (TKR2 n (TKProduct y z))+  runzip :: forall y z n.+            target (TKR2 n (TKProduct y z))+         -> target (TKProduct (TKR2 n y) (TKR2 n z))+  szip :: forall y z sh. (KnownSTK y, KnownSTK z)+       => target (TKProduct (TKS2 sh y) (TKS2 sh z))+       -> target (TKS2 sh (TKProduct y z))+  sunzip :: forall y z sh.+            target (TKS2 sh (TKProduct y z))+         -> target (TKProduct (TKS2 sh y) (TKS2 sh z))+  xzip :: forall y z sh. (KnownSTK y, KnownSTK z)+       => target (TKProduct (TKX2 sh y) (TKX2 sh z))+       -> target (TKX2 sh (TKProduct y z))+  xunzip :: forall y z sh.+            target (TKX2 sh (TKProduct y z))+         -> target (TKProduct (TKX2 sh y) (TKX2 sh z))++  rnest :: forall n m x.+           (KnownNat m, KnownSTK x)+        => SNat n -> target (TKR2 (n + m) x)+        -> target (TKR2 n (TKR2 m x))+  rnest n@SNat =+    gcastWith (unsafeCoerceRefl :: Rank (Replicate n (Nothing @Nat)) :~: n) $+    gcastWith (unsafeCoerceRefl :: Rank (Replicate n (Nothing @Nat)+                                         ++ Replicate m Nothing) :~: n + m) $+    gcastWith (unsafeCoerceRefl :: Replicate (n + m) (Nothing @Nat)+                                    :~: Replicate n (Nothing @Nat)+                                        ++ Replicate m Nothing) $+    withKnownShX (ssxReplicate n) $+    withKnownShX (ssxReplicate (SNat @(n + m))) $+    rfromX . xnestR (ssxReplicate n) . xfromR @_ @(Replicate (n + m) Nothing)+  -- Some of these operations have awkward type signatures, but the signatures+  -- express the most type-safe, or in other words the strongest versions+  -- of the typing possible.+  rnestS :: forall n sh2 x.+            (KnownShS sh2, KnownSTK x)+         => SNat n -> target (TKX2 (Replicate n Nothing ++ MapJust sh2) x)+         -> target (TKR2 n (TKS2 sh2 x))+  rnestS n@SNat =+    gcastWith (unsafeCoerceRefl :: Rank (Replicate n (Nothing @Nat)) :~: n) $+    withKnownShX (ssxReplicate n) $+    rfromX . xnestS (ssxReplicate n)+  rnestX :: forall n sh2 x.+            (KnownShX sh2, KnownSTK x)+         => SNat n -> target (TKX2 (Replicate n Nothing ++ sh2) x)+         -> target (TKR2 n (TKX2 sh2 x))+  rnestX n@SNat =+    gcastWith (unsafeCoerceRefl :: Rank (Replicate n (Nothing @Nat)) :~: n) $+    withKnownShX (ssxReplicate n) $+    rfromX . xnest (ssxReplicate n)+  snestR :: forall sh1 m x.+            (KnownNat m, KnownSTK x)+         => ShS sh1 -> target (TKX2 (MapJust sh1 ++ Replicate m Nothing) x)+         -> target (TKS2 sh1 (TKR2 m x))+  snestR sh1 =+    gcastWith (lemRankMapJust sh1) $+    withKnownShS sh1 $+    withKnownShX (ssxFromShX (shxFromShS sh1)) $+    sfromX . xnestR (ssxFromShX (shxFromShS sh1))+  snest :: forall sh1 sh2 x.+           (KnownShS sh2, KnownSTK x)+        => ShS sh1 -> target (TKS2 (sh1 ++ sh2) x)+        -> target (TKS2 sh1 (TKS2 sh2 x))+  snest sh1 =+    gcastWith (lemRankMapJust sh1) $+    gcastWith (unsafeCoerceRefl :: Rank (MapJust sh1 ++ MapJust sh2)+                                   :~: Rank (sh1 ++ sh2)) $+    withKnownShS sh1 $+    withKnownShX (ssxFromShX (shxFromShS sh1)) $+    withKnownShS (sh1 `shsAppend` knownShS @sh2) $+    withKnownShX (ssxFromShX (shxFromShS sh1)+                  `ssxAppend` ssxFromShX (shxFromShS (knownShS @sh2))) $+    sfromX . xnestS (ssxFromShX (shxFromShS sh1)) . xfromS+  snestX :: forall sh1 sh2 x.+            (KnownShX sh2, KnownSTK x)+         => ShS sh1 -> target (TKX2 (MapJust sh1 ++ sh2) x)+         -> target (TKS2 sh1 (TKX2 sh2 x))+  snestX sh1 =+    gcastWith (lemRankMapJust sh1) $+    withKnownShS sh1 $+    withKnownShX (ssxFromShX (shxFromShS sh1)) $+    sfromX . xnest (ssxFromShX (shxFromShS sh1))+  -- These three are primitives; the others are defined from them.+  xnestR :: forall sh1 m x.+            (KnownNat m, KnownSTK x)+         => StaticShX sh1 -> target (TKX2 (sh1 ++ Replicate m Nothing) x)+         -> target (TKX2 sh1 (TKR2 m x))+  xnestS :: forall sh1 sh2 x.+            (KnownShS sh2, KnownSTK x)+         => StaticShX sh1 -> target (TKX2 (sh1 ++ MapJust sh2) x)+         -> target (TKX2 sh1 (TKS2 sh2 x))+  xnest :: forall sh1 sh2 x.+           (KnownShX sh2, KnownSTK x)+        => StaticShX sh1 -> target (TKX2 (sh1 ++ sh2) x)+        -> target (TKX2 sh1 (TKX2 sh2 x))++  runNest :: (KnownNat n, KnownNat m, KnownSTK x)+          => target (TKR2 n (TKR2 m x)) -> target (TKR2 (n + m) x)+  runNest @n @m =+    gcastWith (unsafeCoerceRefl :: Rank (Replicate n (Nothing @Nat)) :~: n) $+    gcastWith (unsafeCoerceRefl :: Rank (Replicate n (Nothing @Nat)+                                          ++ Replicate m Nothing) :~: n + m) $+    withKnownShX (ssxReplicate (SNat @n)) $+    withKnownShX (ssxReplicate (SNat @n) `ssxAppend` ssxReplicate (SNat @m)) $+    rfromX . xunNestR . xfromR @_ @(Replicate n Nothing)+  runNestS :: (KnownNat n, KnownShS sh2, KnownSTK x)+           => target (TKR2 n (TKS2 sh2 x))+           -> target (TKX2 (Replicate n Nothing ++ MapJust sh2) x)+  runNestS @n =+    gcastWith (unsafeCoerceRefl :: Rank (Replicate n (Nothing @Nat)) :~: n) $+    withKnownShX (ssxReplicate (SNat @n)) $+    xunNestS . xfromR @_ @(Replicate n Nothing)+  runNestX :: (KnownNat n, KnownShX sh2, KnownSTK x)+           => target (TKR2 n (TKX2 sh2 x))+           -> target (TKX2 (Replicate n Nothing ++ sh2) x)+  runNestX @n @sh2=+    gcastWith (unsafeCoerceRefl :: Rank (Replicate n (Nothing @Nat)) :~: n) $+    withKnownShX (ssxReplicate (SNat @n)) $+    withKnownShX (ssxReplicate (SNat @n) `ssxAppend` knownShX @sh2) $+    xunNest . xfromR @_ @(Replicate n Nothing)+  sunNestR :: (KnownShS sh1, KnownNat m, KnownSTK x)+           => target (TKS2 sh1 (TKR2 m x))+           -> target (TKX2 (MapJust sh1 ++ Replicate m Nothing) x)+  sunNestR @sh1 =+    gcastWith (lemRankMapJust (knownShS @sh1)) $+    withKnownShX (ssxFromShX (shxFromShS (knownShS @sh1))) $+    xunNestR . xfromS @_ @_ @(MapJust sh1)+  sunNest :: (KnownShS sh1, KnownShS sh2, KnownSTK x)+          => target (TKS2 sh1 (TKS2 sh2 x)) -> target (TKS2 (sh1 ++ sh2) x)+  sunNest @sh1 @sh2 =+    gcastWith (lemRankMapJust (knownShS @sh1)) $+    gcastWith (unsafeCoerceRefl+               :: Rank (MapJust sh1 ++ MapJust sh2) :~: Rank (sh1 ++ sh2)) $+    withKnownShS (knownShS @sh1 `shsAppend` knownShS @sh2) $+    withKnownShX (ssxFromShX (shxFromShS (knownShS @sh1))) $+    withKnownShX (ssxFromShX (shxFromShS (knownShS @sh1))+                  `ssxAppend` ssxFromShX (shxFromShS (knownShS @sh2))) $+    sfromX . xunNestS . xfromS @_ @_ @(MapJust sh1)+  sunNestX :: (KnownShS sh1, KnownShX sh2, KnownSTK x)+           => target (TKS2 sh1 (TKX2 sh2 x))+           -> target (TKX2 (MapJust sh1 ++ sh2) x)+  sunNestX @sh1 @sh2 =+    gcastWith (lemRankMapJust (knownShS @sh1)) $+    withKnownShX (ssxFromShX (shxFromShS (knownShS @sh1))) $+    withKnownShX (ssxFromShX (shxFromShS (knownShS @sh1))+                  `ssxAppend` knownShX @sh2) $+    xunNest . xfromS @_ @_ @(MapJust sh1)+  -- These three are primitives; the others are defined from them.+  xunNestR :: (KnownShX sh1, KnownNat m, KnownSTK x)+           => target (TKX2 sh1 (TKR2 m x))+           -> target (TKX2 (sh1 ++ Replicate m Nothing) x)+  xunNestS :: (KnownShX sh1, KnownShS sh2, KnownSTK x)+           => target (TKX2 sh1 (TKS2 sh2 x))+           -> target (TKX2 (sh1 ++ MapJust sh2) x)+  xunNest :: (KnownShX sh1, KnownShX sh2, KnownSTK x)+          => target (TKX2 sh1 (TKX2 sh2 x)) -> target (TKX2 (sh1 ++ sh2) x)++  -- Two aliases to make the class sufficient for Unwind.+  -- | A clone of tpair, to make this class independent of @BaseTensor@+  -- but sufficient for "Unwind".+  tpairConv :: target x -> target z -> target (TKProduct x z)+  -- | A clone of tunpair, if @ShareTensor@ is available, or an implementation+  -- that duplicates the argument, otherwise.+  tunpairConv :: target (TKProduct x z) -> (target x, target z)
+ src/HordeAd/Core/Delta.hs view
@@ -0,0 +1,459 @@+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | The grammar of delta expressions.+--+-- A delta expression can be viewed as a concise representation+-- of a linear map (which is the derivative of the objective function)+-- and its evaluation on a given argument (in module "HordeAd.Core.DeltaEval")+-- as an adjoint (in the algebraic sense) of the linear map+-- applied to that argument. Since linear maps can be represented+-- as matrices, this operation corresponds to a transposition+-- of the matrix. However, the matrix is not constructed,+-- but is represented and transposed preserving the sparsity+-- of the representation.+--+-- The \'sparsity\' is less obvious when a delta expression+-- contains big concrete tensors, e.g., via the `DeltaScale` constructor.+-- However, via 'DeltaReplicate' and other constructors, the tensors+-- can be enlarged much beyond what's embedded in the delta term.+-- Also, if the expression refers to unknown inputs ('DeltaInput')+-- it may denote, after evaluation, a still larger tensor.+--+-- The algebraic structure here is an extension of vector space+-- with some additional constructors. The crucial extra constructor+-- 'DeltaInput' replaces the usual one-hot access to parameters+-- with something cheaper and more uniform.+-- A lot of the remaining additional constructors is for introducing+-- and reducing dimensions of tensors and it mimics many of the operations+-- available for the primal value arrays.+module HordeAd.Core.Delta+  ( -- * Delta identifiers+    NodeId, mkNodeId, nodeIdToFTK+  , InputId, mkInputId, inputIdToFTK+    -- * The grammar of delta expressions+  , Delta(..), NestedTarget(..)+    -- * Full tensor kind derivation for delta expressions+  , ftkDelta+  ) where++import Prelude++import Control.Exception.Assert.Sugar+import Data.Dependent.EnumMap.Strict qualified as DMap+import Data.Kind (Type)+import Data.Some+import Data.Type.Equality (TestEquality (..), gcastWith, testEquality, (:~:))+import Data.Vector.Generic qualified as V+import Data.Vector.Strict qualified as Data.Vector+import GHC.TypeLits (type (+), type (<=))+import Text.Show.Functions ()++import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (snatPlus, unsafeCoerceRefl)++import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- * Delta identifiers++-- | The identifiers for nodes of delta expression trees.+type role NodeId nominal nominal+data NodeId :: Target -> TK -> Type where+  NodeId :: forall target y. FullShapeTK y -> Int -> NodeId target y++-- No Eq instance to limit hacks outside this module.++instance Show (NodeId target y) where+  showsPrec d (NodeId _ n) =+    showsPrec d n  -- less verbose, more readable++instance DMap.Enum1 (NodeId target) where+  type Enum1Info (NodeId target) = Some FullShapeTK+  fromEnum1 (NodeId ftk n) = (n, Some ftk)+  toEnum1 n (Some ftk) = Some $ NodeId ftk n++instance TestEquality (NodeId target) where+  testEquality (NodeId ftk1 _) (NodeId ftk2 _) = matchingFTK ftk1 ftk2++-- | Wrap non-negative (only!) integers in the t'NodeId' newtype.+mkNodeId :: FullShapeTK y -> Int -> NodeId f y+mkNodeId ftk i = assert (i >= 0) $ NodeId ftk i++nodeIdToFTK :: NodeId f y -> FullShapeTK y+nodeIdToFTK (NodeId ftk _) = ftk++-- | The identifiers for input leaves of delta expressions.+type role InputId nominal nominal+data InputId :: Target -> TK -> Type where+  InputId :: forall target y. FullShapeTK y -> Int -> InputId target y++-- No Eq instance to limit hacks outside this module.++instance Show (InputId target y) where  -- backward compatibility+  showsPrec _ (InputId _ n) =+    showParen True+    $ showString "InputId "+      . shows n++instance DMap.Enum1 (InputId target) where+  type Enum1Info (InputId target) = Some FullShapeTK+  fromEnum1 (InputId ftk n) = (n, Some ftk)+  toEnum1 n (Some ftk) = Some $ InputId ftk n++instance TestEquality (InputId target) where+  testEquality (InputId ftk1 _) (InputId ftk2 _) = matchingFTK ftk1 ftk2++-- | Wrap non-negative (only!) integers in the t'InputId' newtype.+mkInputId :: FullShapeTK y -> Int -> InputId f y+mkInputId ftk i = assert (i >= 0) $ InputId ftk i++inputIdToFTK :: InputId f y -> FullShapeTK y+inputIdToFTK (InputId ftk _) = ftk+++-- * The grammar of delta expressions++-- | The grammar of delta expressions.+--+-- The t`NodeId` identifier that appears in a @DeltaShare n d@ expression+-- is the unique identity stamp of subterm @d@, that is, there is+-- no different term @e@ such that @DeltaShare n e@ appears in any delta+-- expression term in memory during the same run of an executable.+-- The subterm identity is used to avoid evaluating shared+-- subterms repeatedly in gradient and derivative computations.+-- The identifiers also represent data dependencies among terms+-- for the purpose of gradient and derivative computation. Computation for+-- a term may depend only on data obtained from terms with lower value+-- of their node identifiers. Such data dependency determination+-- agrees with the subterm relation, but is faster than traversing+-- the term tree in order to determine the relation of terms.+--+-- When computing gradients, node identifiers are also used to index,+-- directly or indirectly, the data accumulated for each node,+-- in the form of cotangents, that is partial derivatives+-- of the objective function with respect to the position(s)+-- of the node in the whole objective function dual number term+-- (or, more precisely, with respect to the single node in the term DAG,+-- in which subterms with the same node identifier are collapsed).+-- Only the @DeltaInput@ nodes have a separate data storage.+-- The t`InputId` identifiers in the @DeltaInput@ term constructors+-- are indexes into a contiguous vector of cotangents of @DeltaInput@+-- subterms of the whole term. The value at that index is the partial+-- derivative of the objective function (represented by the whole term,+-- or more precisely by (the data flow graph of) its particular+-- evaluation from which the delta expression originates)+-- with respect to the input parameter component at that index+-- in the objective function domain.+type role Delta nominal nominal+data Delta :: Target -> Target where+  -- Sharing-related operations+  DeltaShare :: NodeId target y -> Delta target y -> Delta target y+  DeltaInput :: InputId target y -> Delta target y++  -- General operations+  DeltaPair :: forall y z target.+               Delta target y -> Delta target z+            -> Delta target (TKProduct y z)+  DeltaProject1 :: forall y z target.+                   Delta target (TKProduct y z) -> Delta target y+  DeltaProject2 :: forall y z target.+                   Delta target (TKProduct y z) -> Delta target z+  DeltaFromVector :: forall y k target.+                     SNat k -> SingletonTK y+                  -> Data.Vector.Vector (Delta target y)+                  -> Delta target (BuildTensorKind k y)+  DeltaSum :: forall y k target.+              SNat k -> SingletonTK y+           -> Delta target (BuildTensorKind k y)+           -> Delta target y+  DeltaReplicate :: forall y k target.+                    SNat k -> SingletonTK y+                 -> Delta target y+                 -> Delta target (BuildTensorKind k y)+  DeltaMapAccumR+    :: forall target k accy by ey.+       ( Show (target (BuildTensorKind k accy))+       , Show (target (BuildTensorKind k ey)) )+    => SNat k+    -> FullShapeTK by+    -> FullShapeTK ey+    -> target (BuildTensorKind k accy)+    -> target (BuildTensorKind k ey)+    -> HFun (TKProduct (ADTensorKind (TKProduct accy ey))+                       (TKProduct accy ey))+            (ADTensorKind (TKProduct accy by))+    -> HFun (TKProduct (ADTensorKind (TKProduct accy by))+                       (TKProduct accy ey))+            (ADTensorKind (TKProduct accy ey))+    -> Delta target accy+    -> Delta target (BuildTensorKind k ey)+    -> Delta target (TKProduct accy (BuildTensorKind k by))+  DeltaMapAccumL+    :: forall target k accy by ey.+       ( Show (target (BuildTensorKind k accy))+       , Show (target (BuildTensorKind k ey)) )+    => SNat k+    -> FullShapeTK by+    -> FullShapeTK ey+    -> target (BuildTensorKind k accy)+    -> target (BuildTensorKind k ey)+    -> HFun (TKProduct (ADTensorKind (TKProduct accy ey))+                       (TKProduct accy ey))+            (ADTensorKind (TKProduct accy by))+    -> HFun (TKProduct (ADTensorKind (TKProduct accy by))+                       (TKProduct accy ey))+            (ADTensorKind (TKProduct accy ey))+    -> Delta target accy+    -> Delta target (BuildTensorKind k ey)+    -> Delta target (TKProduct accy (BuildTensorKind k by))++  -- Vector space operations+  DeltaZero :: FullShapeTK y -> Delta target y+  DeltaScale :: Num (target y)+             => NestedTarget target y -> Delta target y -> Delta target y+  DeltaAdd :: Num (target y)+           => Delta target y -> Delta target y -> Delta target y++  -- Scalar arithmetic+  DeltaCastK :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+             => Delta target (TKScalar r1) -> Delta target (TKScalar r2)++  -- Ranked tensor operations+  DeltaCastR :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+             => Delta target (TKR n r1) -> Delta target (TKR n r2)+  DeltaSum0R :: Delta target (TKR2 n r) -> Delta target (TKR2 0 r)+  DeltaDot0R :: (GoodScalar r, Show (target (TKR n r)))+             => target (TKR n r) -> Delta target (TKR n r)+             -> Delta target (TKR 0 r)+  DeltaIndexR :: forall m n r target.+                 SNat n+              -> Delta target (TKR2 (m + n) r) -> IxROf target m+              -> Delta target (TKR2 n r)+  DeltaScatterR :: forall m n p r target.+                   SNat m -> SNat n -> SNat p+                -> IShR (p + n) -> Delta target (TKR2 (m + n) r)+                -> (IxROf target m -> IxROf target p)+                -> Delta target (TKR2 (p + n) r)+  DeltaGatherR :: forall m n p r target.+                  SNat m -> SNat n -> SNat p+               -> IShR (m + n) -> Delta target (TKR2 (p + n) r)+               -> (IxROf target m -> IxROf target p)+               -> Delta target (TKR2 (m + n) r)+  DeltaAppendR :: Delta target (TKR2 (1 + n) r)+               -> Delta target (TKR2 (1 + n) r)+               -> Delta target (TKR2 (1 + n) r)+  DeltaSliceR :: Int -> Int -> Delta target (TKR2 (1 + n) r)+              -> Delta target (TKR2 (1 + n) r)+  DeltaReverseR :: Delta target (TKR2 (1 + n) r)+                -> Delta target (TKR2 (1 + n) r)+  DeltaTransposeR :: Permutation.PermR -> Delta target (TKR2 n r)+                  -> Delta target (TKR2 n r)+  DeltaReshapeR :: IShR m -> Delta target (TKR2 n r)+                -> Delta target (TKR2 m r)++  -- Shaped tensor operations+  DeltaCastS :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+             => Delta target (TKS sh r1) -> Delta target (TKS sh r2)+  DeltaSum0S :: Delta target (TKS2 sh r) -> Delta target (TKS2 '[] r)+  DeltaDot0S :: (GoodScalar r, Show (target (TKS sh r)))+             => target (TKS sh r) -> Delta target (TKS sh r)+             -> Delta target (TKS '[] r)+  DeltaIndexS :: forall shm shn r target.+                 ShS shn+              -> Delta target (TKS2 (shm ++ shn) r) -> IxSOf target shm+              -> Delta target (TKS2 shn r)+  DeltaScatterS :: forall shm shn shp r target.+                   ShS shm -> ShS shn -> ShS shp+                -> Delta target (TKS2 (shm ++ shn) r)+                -> (IxSOf target shm -> IxSOf target shp)+                -> Delta target (TKS2 (shp ++ shn) r)+  DeltaGatherS :: forall shm shn shp r target.+                  ShS shm -> ShS shn -> ShS shp+               -> Delta target (TKS2 (shp ++ shn) r)+               -> (IxSOf target shm -> IxSOf target shp)+               -> Delta target (TKS2 (shm ++ shn) r)+  DeltaAppendS :: forall target r m n sh.+                  Delta target (TKS2 (m ': sh) r)+               -> Delta target (TKS2 (n ': sh) r)+               -> Delta target (TKS2 ((m + n) ': sh) r)+  DeltaSliceS :: SNat i -> SNat n -> SNat k+              -> Delta target (TKS2 (i + n + k ': sh) r)+              -> Delta target (TKS2 (n ': sh) r)+  DeltaReverseS :: Delta target (TKS2 (n ': sh) r)+                -> Delta target (TKS2 (n ': sh) r)+  DeltaTransposeS :: forall perm sh r target.+                     (Permutation.IsPermutation perm, Rank perm <= Rank sh)+                  => Permutation.Perm perm+                  -> Delta target (TKS2 sh r)+                  -> Delta target (TKS2 (Permutation.PermutePrefix perm sh) r)+  DeltaReshapeS :: Product sh ~ Product sh2+                => ShS sh2+                -> Delta target (TKS2 sh r)+                -> Delta target (TKS2 sh2 r)++  -- Mixed tensor operations+  DeltaCastX :: (GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2)+             => Delta target (TKX sh r1) -> Delta target (TKX sh r2)+  DeltaSum0X :: Delta target (TKX2 sh r) -> Delta target (TKX2 '[] r)+  DeltaDot0X :: (GoodScalar r, Show (target (TKX sh r)))+             => target (TKX sh r) -> Delta target (TKX sh r)+             -> Delta target (TKX '[] r)+  DeltaIndexX :: forall shm shn r target.+                 StaticShX shn+              -> Delta target (TKX2 (shm ++ shn) r) -> IxXOf target shm+              -> Delta target (TKX2 shn r)+  DeltaScatterX :: StaticShX shm -> StaticShX shn -> StaticShX shp+                -> IShX (shp ++ shn) -> Delta target (TKX2 (shm ++ shn) r)+                -> (IxXOf target shm -> IxXOf target shp)+                -> Delta target (TKX2 (shp ++ shn) r)+  DeltaGatherX :: StaticShX shm -> StaticShX shn -> StaticShX shp+               -> IShX (shm ++ shn) -> Delta target (TKX2 (shp ++ shn) r)+               -> (IxXOf target shm -> IxXOf target shp)+               -> Delta target (TKX2 (shm ++ shn) r)+  DeltaAppendX :: Delta target (TKX2 (Just m ': sh) r)+               -> Delta target (TKX2 (Just n ': sh) r)+               -> Delta target (TKX2 (Just (m + n) ': sh) r)+  DeltaSliceX :: SNat i -> SNat n -> SNat k+              -> Delta target (TKX2 (Just (i + n + k) ': sh) r)+              -> Delta target (TKX2 (Just n ': sh) r)+  DeltaReverseX :: Delta target (TKX2 (mn ': sh) r)+                -> Delta target (TKX2 (mn ': sh) r)+  DeltaTransposeX :: forall perm sh r target.+                     (Permutation.IsPermutation perm, Rank perm <= Rank sh)+                  => Permutation.Perm perm+                  -> Delta target (TKX2 sh r)+                  -> Delta target (TKX2 (Permutation.PermutePrefix perm sh) r)+  DeltaReshapeX :: IShX sh2 -> Delta target (TKX2 sh r)+                -> Delta target (TKX2 sh2 r)++  -- Conversions+  DeltaConvert :: TKConversion a b -> Delta target a -> Delta target b++deriving instance Show (IntOf target) => Show (Delta target y)++-- | A newtype defined only to cut the knot of 'Show' instances in 'DeltaScale'+-- that are problematic to pass around as dictionaries without+-- bloating each constructor. The @DeltaScale@ constructor appears+-- in delta expressions a lot and so the primal+-- subterm would bloat the pretty-printed output (though OTOH the primal+-- terms are often important).+--+-- Possibly, @Has Show (Delta target)@ is a better solution.+type NestedTarget :: Target -> Target+type role NestedTarget nominal nominal+newtype NestedTarget target y = NestedTarget (target y)++instance Show (NestedTarget target y) where+  showsPrec _ _ = showString "<primal>"+++-- * Full tensor kind derivation for delta expressions++-- | Full tensor kind derivation for delta expressions.+ftkDelta :: forall target y.+            Delta target y -> FullShapeTK y+ftkDelta = \case+  DeltaShare i _ -> nodeIdToFTK i+  DeltaInput i -> inputIdToFTK i++  DeltaPair t1 t2 -> FTKProduct (ftkDelta t1) (ftkDelta t2)+  DeltaProject1 v -> case ftkDelta v of+    FTKProduct ftk1 _ -> ftk1+  DeltaProject2 v -> case ftkDelta v of+    FTKProduct _ ftk2 -> ftk2+  DeltaFromVector snat _ l -> case V.uncons l of+    Nothing -> error "ftkDelta: empty vector"+    Just (d, _) -> buildFTK snat (ftkDelta d)+  DeltaSum snat stk d -> razeFTK snat stk (ftkDelta d)+  DeltaReplicate snat _ d -> buildFTK snat (ftkDelta d)+  DeltaMapAccumR k bftk _eftk _q _es _df _rf acc0' _es' ->+    FTKProduct (ftkDelta acc0') (buildFTK k bftk)+  DeltaMapAccumL k bftk _eftk _q _es _df _rf acc0' _es' ->+    FTKProduct (ftkDelta acc0') (buildFTK k bftk)++  DeltaZero ftk -> ftk+  DeltaScale _ d -> ftkDelta d+  DeltaAdd (DeltaShare i _) _ -> nodeIdToFTK i+  DeltaAdd _ e -> ftkDelta e++  DeltaCastK{} -> FTKScalar++  DeltaCastR d -> case ftkDelta d of+    FTKR sh _ -> FTKR sh FTKScalar+  DeltaSum0R d -> case ftkDelta d of+    FTKR _ x -> FTKR ZSR x+  DeltaDot0R{} -> FTKR ZSR FTKScalar+  DeltaIndexR SNat d ix | SNat <- ixrRank ix -> case ftkDelta d of+    FTKR sh x -> FTKR (shrDrop sh) x+  DeltaScatterR _ _ _ sh d _ -> case ftkDelta d of+    FTKR _ x -> FTKR sh x+  DeltaGatherR _ _ _ sh d _ -> case ftkDelta d of+    FTKR _ x -> FTKR sh x+  DeltaAppendR a b -> case ftkDelta a of+    FTKR ZSR _ -> error "ftkDelta: impossible pattern needlessly required"+    FTKR (ai :$: ash) x -> case ftkDelta b of+      FTKR ZSR _ -> error "ftkDelta: impossible pattern needlessly required"+      FTKR (bi :$: _) _ -> FTKR (ai + bi :$: ash) x+  DeltaSliceR _ n d -> case ftkDelta d of+    FTKR sh x -> FTKR (n :$: shrTail sh) x+  DeltaReverseR d -> ftkDelta d+  DeltaTransposeR perm d -> case ftkDelta d of+    FTKR sh x -> FTKR (shrPermutePrefix perm sh) x+  DeltaReshapeR sh d -> case ftkDelta d of+    FTKR _ x -> FTKR sh x++  DeltaCastS d -> case ftkDelta d of+    FTKS sh FTKScalar -> FTKS sh FTKScalar+  DeltaSum0S d -> case ftkDelta d of+    FTKS _ x -> FTKS ZSS x+  DeltaDot0S{} -> FTKS ZSS FTKScalar+  DeltaIndexS shn d _ix -> case ftkDelta d of+    FTKS _ x -> FTKS shn x+  DeltaScatterS _shm shn shp d _ -> case ftkDelta d of+    FTKS _ x -> FTKS (shp `shsAppend` shn) x+  DeltaGatherS shm shn _shp d _ -> case ftkDelta d of+    FTKS _ x -> FTKS (shm `shsAppend` shn) x+  DeltaAppendS a b -> case (ftkDelta a, ftkDelta b) of+    (FTKS (m :$$ sh) x, FTKS (n :$$ _) _) -> FTKS (snatPlus m n :$$ sh) x+  DeltaSliceS _ n@SNat _ d -> case ftkDelta d of+    FTKS (_ :$$ sh) x -> FTKS (n :$$ sh) x+  DeltaReverseS d -> ftkDelta d+  DeltaTransposeS perm d -> case ftkDelta d of+    FTKS sh x -> FTKS (shsPermutePrefix perm sh) x+  DeltaReshapeS sh2 d -> case ftkDelta d of+    FTKS _ x -> FTKS sh2 x++  DeltaCastX d -> case ftkDelta d of+    FTKX sh FTKScalar -> FTKX sh FTKScalar+  DeltaSum0X d -> case ftkDelta d of+    FTKX _ x -> FTKX ZSX x+  DeltaDot0X{} -> FTKX ZSX FTKScalar+  DeltaIndexX @shm @shn shn d ix -> case ftkDelta d of+    FTKX sh x | SNat @len <- ixxRank ix ->+      gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $+      withKnownShX (ssxFromShX sh) $+      withKnownShX shn $+      FTKX (shxDrop @len sh) x  -- TODO: (shxDropSSX sh (ssxFromIxX ix)) x+  DeltaScatterX _ _ _ sh d _ -> case ftkDelta d of+    FTKX _ x -> FTKX sh x+  DeltaGatherX _ _ _ sh d _ -> case ftkDelta d of+    FTKX _ x -> FTKX sh x+  DeltaAppendX a b -> case (ftkDelta a, ftkDelta b) of+    (FTKX (Nested.SKnown m :$% sh) x, FTKX (Nested.SKnown n :$% _) _) ->+      FTKX (Nested.SKnown (snatPlus m n) :$% sh) x+  DeltaSliceX _ n@SNat _ d -> case ftkDelta d of+    FTKX (_ :$% sh) x -> FTKX (Nested.SKnown n :$% sh) x+  DeltaReverseX d -> ftkDelta d+  DeltaTransposeX perm d -> case ftkDelta d of+    FTKX sh x -> FTKX (shxPermutePrefix perm sh) x+  DeltaReshapeX sh2 d -> case ftkDelta d of+    FTKX _ x -> FTKX sh2 x++  DeltaConvert c d -> convertFTK c $ ftkDelta d
+ src/HordeAd/Core/DeltaEval.hs view
@@ -0,0 +1,1098 @@+{-# LANGUAGE AllowAmbiguousTypes, QuantifiedConstraints #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Evaluation of delta expressions, that is, transpose of the linear+-- maps of which the delta expressions are sparse representations.+-- See comments in "HordeAd.Core.Delta".+module HordeAd.Core.DeltaEval+  ( -- * Delta expression evaluation+    gradientFromDelta, derivativeFromDelta+    -- * Exported to be specialized elsewhere+  , evalRev, evalRevFTK, evalRevSame, evalRevFromnMap, EvalState+  ) where++import Prelude++import Control.Arrow (second)+import Control.Exception.Assert.Sugar+import Data.Dependent.EnumMap.Strict (DEnumMap)+import Data.Dependent.EnumMap.Strict qualified as DMap+import Data.Dependent.Sum (DSum (..))+import Data.Proxy (Proxy (Proxy))+import Data.Traversable (mapAccumL)+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import Text.Show (showListWith)+import Text.Show.Functions ()+import Type.Reflection (typeRep)++import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation (permInverse)+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.ConvertTensor+import HordeAd.Core.Delta+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind++-- * Top-level functions for computing derivatives from delta expressions++-- | The top-level function for computing a gradient of an objective function.+--+-- Delta expressions naturally denote forward derivatives, as encoded+-- in function 'derivativeFromDelta'. However, we are usually more+-- interested in computing gradients, which is what @gradientFromDelta@ does.+-- The two functions are bound by the equation from Lemma 5 from the paper+-- "Provably correct, asymptotically efficient, higher-order reverse-mode+-- automatic differentiation":+--+-- > dt <.> derivativeFromDelta d ds = gradientFromDelta d dt <.> ds+--+-- where @\<.\>@ denotes generalized dot product (multiplying+-- all tensors element-wise and summing the results), @d@ is the top level+-- delta expression from translation of the objective function @f@ to dual+-- numbers, @ds@ belongs to the domain of @f@ and @dt@ to the codomain.+-- In other words, @ds@ is a perturbation (small change) of the arguments+-- of @f@, for which we compute the derivative, and @dt@ is a sensitivity+-- of the result of @f@, for which we compute the gradient.+-- Nota bene, this property is checked for many example objective functions+-- (and perturbations and sensitivities) in the horde-ad testsuite.+gradientFromDelta+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => FullShapeTK x+  -> FullShapeTK z+  -> target (ADTensorKind z)+  -> Delta target z+  -> target (ADTensorKind x)+gradientFromDelta !xftk !zftk !dt deltaTopLevel =+  let s0 = initEvalState xftk+      s1 = evalRev zftk s0 dt deltaTopLevel+      s2 = evalRevFromnMap s1+      (res, remainder) =+        rebuildInputs @(ADTensorKind x) (DMap.toAscList $ iMap s2) s2+        $ adFTK xftk+  in assert (null remainder) res++-- | The top-level function for computing a (forward) derivative+-- of an objective function.+derivativeFromDelta+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => Delta target z -> FullShapeTK (ADTensorKind x)+  -> target (ADTensorKind x)+  -> target (ADTensorKind z)+derivativeFromDelta deltaTopLevel ftk ds =+    let iMap = DMap.fromDistinctAscList $ fst $ generateDSums 0 ftk ds+        s0 = DMap.empty+        !(!_s2, !c) = evalFwd iMap s0 deltaTopLevel+    in c+++-- * Auxiliary datatypes for Delta evaluation++type ADMap target = DEnumMap (NodeId target) (Cotangent target)++type IMap target = DEnumMap (InputId target) (TensorOrZero target)++showsPrec_IMap+  :: (forall y. KnownSTK y => Show (TensorOrZero target y))+  => Int -> IMap target -> ShowS+showsPrec_IMap d demap =+  showParen (d > 10) $+    showString "fromList "+    . showListWith+        (\(k :=> target) ->+          withKnownSTK (ftkToSTK $ inputIdToFTK k) $+          showsPrec 2 k . showString " :=> " . showsPrec 1 target)+        (DMap.toList demap)++show_IMap+  :: (forall y. KnownSTK y => Show (TensorOrZero target y))+  => IMap target -> String+show_IMap iMap = showsPrec_IMap 0 iMap ""++type role Cotangent nominal nominal+newtype Cotangent target y =+  Cotangent {unCotangent :: target (ADTensorKind y)}++-- | This is a tensor representation where zero tensors are marked specially+-- at construction, when it's cheap to do so (as opposed to, later on.+-- requiring a traversal of large terms or checking that each cell+-- of a huge concrete array is zero). It also makes the computation+-- of a special case of addTensorOrZero cheaper.+type role TensorOrZero nominal nominal+data TensorOrZero target y =+    TOTensor (target y)+  | TOZero (FullShapeTK y)+  deriving Show++evalTensorOrZero :: forall target x. ADReadyNoLet target+                 => TensorOrZero target x -> target x+evalTensorOrZero = \case+  TOTensor t -> t+  TOZero ftk -> tdefTarget ftk++-- The ShareTensor constraint is needed, despite what GHC says,+-- in order not to require duplicable arguments.+addTensorOrZero :: forall target y. (ADReadyNoLet target, ShareTensor target)+                => SingletonTK y+                -> TensorOrZero target y -> TensorOrZero target y+                -> TensorOrZero target y+addTensorOrZero stk a b = case (a, b) of+  (TOTensor ta, TOTensor tb) -> TOTensor $ taddTarget stk ta tb+    -- target has a ShareTensor instance, so ta and tb don't need+    -- to be duplicable+  (TOZero{}, _) -> b+  (_, TOZero{}) -> a++-- Matches generateDSumsDummy.+rebuildInputs :: forall ady target. ADReadyNoLet target+              => [DSum (InputId target) (TensorOrZero target)]+              -> EvalState target  -- original state; only for error messages+              -> FullShapeTK ady+              -> (target ady, [DSum (InputId target) (TensorOrZero target)])+rebuildInputs els s2 ftk = case ftk of+  FTKProduct ftk1 ftk2 ->+      let (t1, rest1) = rebuildInputs els s2 ftk1+          (t2, rest2) = rebuildInputs rest1 s2 ftk2+          !t = tpair t1 t2+      in (t, rest2)+  _ | differentiableFTK ftk -> case els of+    (n :=> tz@(TOTensor t)) : rest ->+      case matchingFTK (inputIdToFTK n) ftk of+        Just Refl ->+          (t, rest)+        _ | Dict <- lemKnownSTK (ftkToSTK $ inputIdToFTK n) ->+          error $ "rebuildInputs: wrong Tensor type: "+                  ++ show (n, tz, show_IMap (iMap s2))+    (n :=> tz@(TOZero ftk2)) : rest ->+      case matchingFTK ftk2 ftk of+        Just Refl ->+          let !zero = tdefTarget ftk+          in (zero, rest)+          -- TODO: actually pass this ZERO through to optimizers+          -- and use there to avoid updating the gradient+          -- and maybe use elsewhere, too, to manage sparsity a bit.+          -- We'd probably need TKProduct to take TensorOrZero.+        _ | Dict <- lemKnownSTK (ftkToSTK $ inputIdToFTK n) ->+          error $ "rebuildInputs: wrong Zero type: "+                  ++ show (n, tz, show_IMap (iMap s2))+    _ -> error $ "rebuildInputs: illegal TensorOrZero: "+                 ++ show_IMap (iMap s2)+  _ -> (tdefTarget ftk, els)++-- Matches generateDeltaInputs.+generateDSumsDummy :: Int -> FullShapeTK y+                   -> ([DSum (InputId target) (TensorOrZero target)], Int)+generateDSumsDummy j ftk  = case ftk of+  FTKProduct ftk1 ftk2 ->+    let (ds1, j1) = generateDSumsDummy j ftk1+        (ds2, j2) = generateDSumsDummy j1 ftk2+    in (ds1 ++ ds2, j2)+  _ | differentiableFTK ftk -> ([mkInputId ftk j :=> TOZero ftk], j + 1)+  _ -> ([], j)++-- Matches generateDeltaInputs.+generateDSums :: ShareTensor target+              => Int -> FullShapeTK y -> target y+              -> ([DSum (InputId target) (TensorOrZero target)], Int)+generateDSums j ftk t = case ftk of+  FTKProduct ftk1 ftk2 ->+    let (t1, t2) = tunpair t+        (ds1, j1) = generateDSums j ftk1 t1+        (ds2, j2) = generateDSums j1 ftk2 t2+    in (ds1 ++ ds2, j2)+  _ | differentiableFTK ftk -> ([mkInputId ftk j :=> TOTensor t], j + 1)+  _ -> ([], j)+++-- * Delta evaluation state++-- | The state of evaluation. It consists of several maps.+-- The maps indexed by input identifiers and node identifiers+-- eventually store cotangents for their respective nodes.+-- The cotangents are built gradually during the evaluation,+-- by summing cotangent contributions.+--+-- Data invariant: keys nMap == keys dMap.+type role EvalState nominal+data EvalState target = EvalState+  { iMap :: IMap target+      -- ^ eventually, cotangents of objective function inputs+      -- (eventually copied to the vector representing the gradient+      -- of the objective function);+      -- the identifiers need to be contiguous and start at 0+  , dMap :: ADMap target+      -- ^ eventually, cotangents of non-input subterms indexed+      -- by their node identifiers+  , nMap :: DEnumMap (NodeId target) (Delta target)+      -- ^ nodes left to be evaluated;+      -- we can't evaluate them at once, because their other shared copies+      -- may still not be processed, so we'd not take advantage of the sharing+      -- and/or not take into account the whole summed context when finally+      -- evaluating+  }++-- | Initialization of the evalutation state, which consists of+-- creating the finite maps that hold values associated with inputs+-- and with (possibly shared) term tree nodes.+-- The former are usually initialized with dummy values so that it's cheap+-- to check if any update has already been performed to a cell+-- (allocating big vectors filled with zeros is too costly,+-- especially if never used in an iteration, and adding to such vectors+-- and especially using them as cotangent accumulators is wasteful).+initEvalState :: FullShapeTK x -> EvalState target+initEvalState ftk0 =+  let iMap = DMap.fromDistinctAscList $ fst $ generateDSumsDummy 0 $ adFTK ftk0+      dMap = DMap.empty+      nMap = DMap.empty+  in EvalState {..}+++-- * Reverse pass, transpose/evaluation of the delta expressions++evalRevScalarRuntimeSpecialized+  :: forall r target.+     (GoodScalar r, ADReadyNoLet target, ShareTensor target)+  => EvalState target -> target (ADTensorKind (TKScalar r))+  -> Delta target (TKScalar r)+  -> EvalState target+{-# INLINE evalRevScalarRuntimeSpecialized #-}+evalRevScalarRuntimeSpecialized !s !c =+  case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> evalRevSame @(TKScalar Double) s c+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> evalRevSame @(TKScalar Float) s c+      _ -> const s++evalRevRRuntimeSpecialized+  :: forall n r target.+     (GoodScalar r, ADReadyNoLet target, ShareTensor target)+  => EvalState target -> target (ADTensorKind (TKR n r))+  -> Delta target (TKR n r)+  -> EvalState target+{-# INLINE evalRevRRuntimeSpecialized #-}+evalRevRRuntimeSpecialized !s !c =+  -- We dispatch on all expected underyling scalar types, which is+  -- necessary to run the correct specialization when unpacking+  -- an existential type. All IfDifferentiable instances should+  -- be included in the list of expected underlying scalar types.+  -- If the scalar type is not on the list, performance suffers greatly.+  case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> evalRevSame @(TKR n Double) s c+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> evalRevSame @(TKR n Float) s c+      _ -> const s++evalSRuntimeSpecialized+  :: forall sh r target.+     (GoodScalar r, ADReadyNoLet target, ShareTensor target)+  => EvalState target -> target (ADTensorKind (TKS sh r))+  -> Delta target (TKS sh r)+  -> EvalState target+{-# INLINE evalSRuntimeSpecialized #-}+evalSRuntimeSpecialized !s !c =+  case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> evalRevSame @(TKS sh Double) s c+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> evalRevSame @(TKS sh Float) s c+      _ -> const s++evalXRuntimeSpecialized+  :: forall sh r target.+     (GoodScalar r, ADReadyNoLet target, ShareTensor target)+  => EvalState target -> target (ADTensorKind (TKX sh r))+  -> Delta target (TKX sh r)+  -> EvalState target+{-# INLINE evalXRuntimeSpecialized #-}+evalXRuntimeSpecialized !s !c =+  case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> evalRevSame @(TKX sh Double) s c+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> evalRevSame @(TKX sh Float) s c+      _ -> const s++-- | Reverse pass, that is, transpose/evaluation of the delta expressions+-- in order to produce the gradient for the objective function runtime+-- trace represented by the delta expression.+--+-- The first argument is the tensor kind that constrains the shapes+-- of the contangent accumulator and the delta expression arguments.+-- The second is the evaluation state being modified.+-- The third is the cotangent accumulator that will become an actual+-- cotangent contribution when complete (see below for an explanation).+-- The fourth is the delta expression node to evaluate.+--+-- Obtaining the gradient amounts to transposing the linear map+-- that is straightforwardly represented by the delta expression.+-- The @evalRev@ function transposes the linear map and,+-- at the same time, evaluates the transposed map on the cotangent accumulator+-- value contained in the third argument. If the cotangent and the tensor+-- operations are symbolic, the resulting value represents the transposed+-- map itself, if its free variables are treated as the map's inputs.+evalRev+  :: forall y target.+     (ADReadyNoLet target, ShareTensor target)+  => FullShapeTK y+  -> EvalState target -> target (ADTensorKind y) -> Delta target y+  -> EvalState target+evalRev ftk !s !c d = case ftk of+  FTKScalar @r -> evalRevScalarRuntimeSpecialized @r s c d+  FTKR @n _ (FTKScalar @r) -> evalRevRRuntimeSpecialized @n @r s c d+  FTKS @sh _ (FTKScalar @r) -> evalSRuntimeSpecialized @sh @r s c d+  FTKX @sh _ (FTKScalar @r) -> evalXRuntimeSpecialized @sh @r s c d+  _ -> evalRevFTK s c d++-- | A helper function to `evalRev`. The @FTK@ suffix denotes it doesn't get+-- an FTK as an argument but reconstructs it as needed.+--+-- All constructors that can have a type with TKProduct kind+-- need to be handled here,+-- as opposed to in 'evalRevSame', except for DeltaInput that is always+-- constructed only in basic kinds even though its type permits others.+evalRevFTK+  :: forall y target.+     (ADReadyNoLet target, ShareTensor target)+  => EvalState target -> target (ADTensorKind y) -> Delta target y+  -> EvalState target+evalRevFTK !s !c d0 = case d0 of+  DeltaShare n d ->+    -- In this context, by construction, @d@ is the dual component+    -- of a dual number term. Let's say that, at this point, evaluation+    -- considers position (node) p out of possibly multiple positions+    -- at which that dual number resides in the whole term tree+    -- of the dual number representation of the objective function.+    -- (Equivalently, considers edge p, one of many leading to the only+    -- node with identifier @n@ in the DAG representing the term).+    -- If so, the @c@ argument of @eval0@ is the cotangent+    -- contribution for position p, that is, the partial derivative+    -- of the objective function with respect to position p.+    --+    -- If there are indeed multiple such positions+    -- (the term is non-trivially shared) then,+    -- over the course of evaluation, cotangent contributions+    -- of them all are gradually accumulated in the finite+    -- maps and eventually their total sum represents the total+    -- influence of the objective function's subcomputation+    -- (more precisely, subgraph of the data flow graph in question)+    -- corresponding to the shared term @DeltaShare n d@. This total+    -- influence over the objective function's behaviour is called+    -- in short the cotangent of the node identifier @n@.+    -- In other words, the cotangent of @n@ is the sum,+    -- over all positions (edges) q in the global delta-expression DAG+    -- that are a reference to node @n@, of the partial derivative+    -- of the objective function with respect to the subcomputation+    -- corresponding to @q@ (meaning, subcomputations denoted by+    -- Haskell terms whose dual components are @Share n ...@).+    --+    -- For @Input@ terms, the eventual lists of cotangents end up+    -- in the cells of the gradient vectors that are the final+    -- result of the evaluation.+    assert (case d of  -- should match shareDelta+              DeltaZero{} -> False+              DeltaPair{} -> False+              DeltaInput{} -> False+              DeltaShare{} -> False+              _ -> True)+    $ if DMap.member n $ nMap s+      then let addc x =+                 Cotangent $ taddTarget (adSTK $ ftkToSTK $ nodeIdToFTK n)+                                        c (unCotangent x)+             -- target has a ShareTensor instance, so taddTarget arguments+             -- don't need to be duplicable+           in s {dMap = DMap.adjust addc n $ dMap s}+      else let cd = Cotangent c+           in s { nMap = DMap.insert n d $ nMap s+                , dMap = DMap.insert n cd $ dMap s }++  DeltaPair d1 d2 ->+    let (c1, c2) = tunpair c+    in evalRevFTK (evalRevFTK s c1 d1) c2 d2+  DeltaProject1 d -> case ftkDelta d of+    FTKProduct _ ftk2 ->+      let zero = tdefTarget $ adFTK ftk2+      in evalRevFTK s (tpair c zero) d+    -- if y is, e.g., TKR 0 Int64, we eval this delta anyway, even though+    -- we could ignore it at the price of complicating or duplicating+    -- the code slightly more+  DeltaProject2 d -> case ftkDelta d of+    FTKProduct ftk1 _ ->+      let zero = tdefTarget $ adFTK ftk1+      in evalRevFTK s (tpair zero c) d+  DeltaFromVector snat stk ld | Refl <- lemBuildOfAD snat stk ->+    let cxs = tunravelToListShare snat (adSTK stk) c+    in foldl' (\ !s2 (cx, d2) -> evalRevFTK s2 cx d2) s+       $ zip cxs (V.toList ld)+  DeltaSum snat stk d | Refl <- lemBuildOfAD snat stk ->+    evalRevFTK s (treplicate snat (adSTK stk) c) d+  DeltaReplicate snat stk d | Refl <- lemBuildOfAD snat stk ->+    evalRevFTK s (tsum snat (adSTK stk) c) d+  DeltaMapAccumR k bftk eftk q es _df rf acc0' es'+   | Refl <- lemBuildOfAD k (ftkToSTK bftk)+   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->+    let accftk = ftkDelta acc0'+        accftkAD = adFTK accftk+        bftkAD = adFTK bftk+        eftkAD = adFTK eftk+        (c0, crest) = tunpair c+        dacc_des =+          tmapAccumL (Proxy @target)+                     k accftkAD eftkAD (FTKProduct bftkAD+                                                   (FTKProduct accftk eftk))+                     (\dx db_acc_e ->+                        ttlet db_acc_e $ \ !db_acc_e1 ->+                          unHFun rf (tpair (tpair dx (tproject1 db_acc_e1))+                                           (tproject2 db_acc_e1)))+                     c0+                     (tpair crest (tpair q es))+        (dacc, des) = tunpair dacc_des+        s2 = evalRevFTK s dacc acc0'+    in evalRevFTK s2 des es'+  DeltaMapAccumL k bftk eftk q es _df rf acc0' es'+   | Refl <- lemBuildOfAD k (ftkToSTK bftk)+   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->+    let accftk = ftkDelta acc0'+        accftkAD = adFTK accftk+        bftkAD = adFTK bftk+        eftkAD = adFTK eftk+        (c0, crest) = tunpair c+        dacc_des =+          tmapAccumR (Proxy @target)+                     k accftkAD eftkAD (FTKProduct bftkAD+                                                   (FTKProduct accftk eftk))+                     (\dx db_acc_e ->+                        ttlet db_acc_e $ \ !db_acc_e1 ->+                          unHFun rf (tpair (tpair dx (tproject1 db_acc_e1))+                                           (tproject2 db_acc_e1)))+                     c0+                     (tpair crest (tpair q es))+        (dacc, des) = tunpair dacc_des+        s2 = evalRevFTK s dacc acc0'+    in evalRevFTK s2 des es'++  _ -> let y = ftkDelta d0+       in case matchingFTK y (adFTK y) of+         Just Refl -> evalRevSame s c d0+         _ -> s  -- the constructors remaining here have y that is+                 -- a non-TKProduct so if y is equal to ADTensorKind y,+                 -- the latter has the Z1 scalar type and so no influence+                 -- on the derivative.++-- | A helper function to `evalRev`. It assumes the scalar underlying+-- the tensor kind of its arguments is differentiable.+--+-- All constructors that can only have types with non-TKProduct kinds+-- (and the DeltaInput constructor and the vector space constructors)+-- can be handled here, where the extra equality constraint makes it easier.+evalRevSame+  :: forall y target.+     (ADReadyNoLet target, ShareTensor target, y ~ ADTensorKind y)+  => EvalState target -> target (ADTensorKind y) -> Delta target y+  -> EvalState target+evalRevSame !s !c = \case+  DeltaInput i ->+    let cs = TOTensor c+    in s {iMap = DMap.adjust (addTensorOrZero (ftkToSTK $ inputIdToFTK i) cs) i+                 $ iMap s}+    -- This and similar don't need to be runtime-specialized,+    -- because the type of c determines the Num instance for (+).+    -- Note that we can't express sharing by inserting DeltaShare constructors+    -- into iMap, because often sharing needs to work across many+    -- iMap keys. That's why global sharing is used.++  -- By placing these here, we force their derivatives to be zeroed+  -- whenever they are called on non-base types, which they should not ever be.+  -- This is ensured by the types of the three constructors, assuming that+  -- no Num instances are defined for the non-base type tensors.+  DeltaZero{} -> s+  DeltaScale (NestedTarget k) d -> evalRevSame s (k * c) d+  DeltaAdd d e ->+    let cShared = tshare c+    in evalRevSame (evalRevSame s cShared d) cShared e++  DeltaCastK @r1 d ->+    evalRevScalarRuntimeSpecialized+    s (toADTensorKindShared (FTKScalar @r1) $ tkcast c) d+  DeltaCastR d -> case ftkDelta d of+    y ->+      evalRevRRuntimeSpecialized+      s (toADTensorKindShared y $ trcast c) d+  DeltaSum0R d -> case ftkDelta d of+    FTKR sh x | SNat <- shrRank sh ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (trreplicate0N sh c) d+  DeltaDot0R v d -> case ftkDelta d of+    FTKR sh x | SNat <- shrRank sh ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (v * trreplicate0N (rshape v) c) d+        -- too slow: evalRevSame s (rmap0N (* (tscalar c)) v) vd+  DeltaIndexR SNat d ix -> case ftkDelta d of+    FTKR sh x | SNat <- ixrRank ix ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (troneHot (shrTake sh) c ix) d  -- TODO: ixrToShR ix+  DeltaScatterR SNat SNat SNat _sh d f -> case ftkDelta d of+    FTKR sh x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (trgather sh c f) d+  DeltaGatherR SNat SNat SNat _sh d f -> case ftkDelta d of+    FTKR sh x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (trscatter sh c f) d+  DeltaAppendR d e -> case (ftkDelta d, ftkDelta e) of+    (FTKR (m :$: _) x, FTKR (n :$: _) _) ->+      withKnownSTK (ftkToSTK x) $+      let cShared = tshare c+          s2 = evalRevSame s (trslice 0 m cShared) d+      in evalRevSame s2 (trslice m n cShared) e+    _ -> error "evalRevSame: impossible pattern needlessly required"+  DeltaSliceR i n d -> case ftkDelta d of+    FTKR (l :$: rest) x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (trappend+                       (tdefTarget (FTKR (i :$: rest) x))+                       (trappend c+                          (tdefTarget (FTKR (l - i - n :$: rest) x)))) d+    FTKR ZSR _ -> error "evalRevSame: impossible pattern needlessly required"+  DeltaReverseR d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (trreverse c) d+  DeltaTransposeR perm d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (ftkToSTK x) $+      let permR = permRInverse perm+      in evalRevSame s (trtranspose permR c) d+  DeltaReshapeR _sh2 d -> case ftkDelta d of+    FTKR sh x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (trreshape sh c) d++  DeltaCastS d -> case ftkDelta d of+    y ->+      evalSRuntimeSpecialized+      s (toADTensorKindShared y $ tscast c) d+  DeltaSum0S d -> case ftkDelta d of+    FTKS sh x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (tsreplicate0N sh c) d+  DeltaDot0S v d -> case ftkDelta d of+    FTKS sh FTKScalar ->+      evalRevSame s (v * tsreplicate0N sh c) d+  DeltaIndexS shn d ix -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShS shn $+      withKnownShS (shsFromIxS ix) $+      evalRevSame s (tsoneHot c ix) d+  DeltaScatterS @shm @shn shm shn shp d f -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShS shm $+      withKnownShS shn $+      withKnownShS shp $+      evalRevSame s (tsgather @_ @shm @shn c f) d+  DeltaGatherS @shm @shn shm shn shp d f -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShS shm $+      withKnownShS shn $+      withKnownShS shp $+      evalRevSame s (tsscatter @_ @shm @shn c f) d+  DeltaAppendS d e -> case (ftkDelta d, ftkDelta e) of+    (FTKS (msnat :$$ _) x, FTKS (_ :$$ _) _) ->+      withKnownSTK (ftkToSTK x) $+      let cShared = tshare c+          s2 = evalRevSame s (tsslice (SNat @0) SNat SNat cShared) d+      in evalRevSame s2 (tsslice msnat SNat SNat cShared) e+  DeltaSliceS i@SNat _ k@SNat d -> case ftkDelta d of+    FTKS (_ :$$ sh) x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (tsappend+                       (tdefTarget (FTKS (i :$$ sh) x))+                          (tsappend+                             c (tdefTarget (FTKS (k :$$ sh) x)))) d+  DeltaReverseS d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (tsreverse c) d+  DeltaTransposeS @perm @sh2 perm d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      permInverse perm $ \(permRev :: Permutation.Perm permR) _ ->+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix+                        permR (Permutation.PermutePrefix perm sh2) :~: sh2)+        $ gcastWith (unsafeCoerceRefl+                     :: Rank (Permutation.PermutePrefix perm sh2) :~: Rank sh2)+        $ gcastWith (unsafeCoerceRefl+                     :: Rank permR :~: Rank perm)+        $ evalRevSame s (tstranspose permRev c) d+  DeltaReshapeS _sh2 d -> case ftkDelta d of+    FTKS sh x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (tsreshape sh c) d++  DeltaCastX d -> case ftkDelta d of+    y ->+      evalXRuntimeSpecialized+      s (toADTensorKindShared y $ txcast c) d+  DeltaSum0X d -> case ftkDelta d of+    FTKX sh x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShX (ssxFromShX sh) $+      evalRevSame s (txreplicate0N sh c) d+  DeltaDot0X v d -> case ftkDelta d of+    FTKX sh FTKScalar ->+      withKnownShX (ssxFromShX sh) $+      evalRevSame s (v * txreplicate0N (xshape v) c) d+  DeltaIndexX @shm @shn shn d ix -> case ftkDelta d of+    FTKX sh x | SNat @len <- ixxRank ix ->+      withKnownSTK (ftkToSTK x) $+      withKnownShX shn $+      withKnownShX (ssxFromShX sh) $+      withKnownShX (ssxTakeIx @shm @shn (ssxFromShX sh) ix) $+      gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $+      evalRevSame s (txoneHot (shxTake @len sh) c ix) d+--TODO      evalRevSame s (xoneHot (shxTakeSSX (Proxy @shn) sh+--                                         (ssxFromIxX ix)) c ix) d+  DeltaScatterX @shm @shn shm shn shp _sh d f -> case ftkDelta d of+    FTKX sh x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShX shm $+      withKnownShX shn $+      withKnownShX shp $+      evalRevSame s (txgather @_ @shm @shn sh c f) d+  DeltaGatherX @shm @shn shm shn shp _sh d f -> case ftkDelta d of+    FTKX sh x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShX shm $+      withKnownShX shn $+      withKnownShX shp $+      evalRevSame s (txscatter @_ @shm @shn sh c f) d+  DeltaAppendX d e -> case (ftkDelta d, ftkDelta e) of+    (FTKX (Nested.SKnown m@SNat :$% _) x, FTKX (Nested.SKnown SNat :$% _) _) ->+      withKnownSTK (ftkToSTK x) $+      let cShared = tshare c+          s2 = evalRevSame s (txslice (SNat @0) SNat SNat cShared) d+      in evalRevSame s2 (txslice m SNat SNat cShared) e+  DeltaSliceX i@SNat _ k@SNat d -> case ftkDelta d of+    FTKX (_ :$% sh) x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (txappend+                       (tdefTarget (FTKX (Nested.SKnown i :$% sh) x))+                          (txappend+                             c (tdefTarget+                                  (FTKX (Nested.SKnown k :$% sh) x)))) d+  DeltaReverseX d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (txreverse c) d+  DeltaTransposeX @perm @sh2 perm d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (ftkToSTK x) $+      permInverse perm $ \(permR :: Permutation.Perm permR) _ ->+        gcastWith (unsafeCoerceRefl+                   :: Permutation.PermutePrefix+                        permR (Permutation.PermutePrefix perm sh2) :~: sh2) $+        gcastWith (unsafeCoerceRefl+                   :: Rank (Permutation.PermutePrefix perm sh2) :~: Rank sh2) $+        gcastWith (unsafeCoerceRefl+                   :: Rank permR :~: Rank perm) $+        evalRevSame s (txtranspose permR c) d+  DeltaReshapeX _sh2 d -> case ftkDelta d of+    FTKX sh x ->+      withKnownSTK (ftkToSTK x) $+      evalRevSame s (txreshape sh c) d++  DeltaConvert @a c1 d -> case ftkDelta d of+    aftk ->+      -- This follows from the same property for @b@ and from @c1@+      -- not changing the underlying scalar types.+      gcastWith (unsafeCoerceRefl :: ADTensorKind a :~: a) $+      evalRevSame+        s (tconvert (transposeTKConversion aftk c1)+                    (convertSTK c1 $ ftkToSTK $ ftkDelta d) c) d++  d -> evalRevFTK s c d+    -- the remaining constructors are already handled in evalRevFTK++transposeTKConversion :: FullShapeTK a -> TKConversion a b -> TKConversion b a+transposeTKConversion aftk c0 = case c0 of+  ConvId -> ConvId+  ConvCmp c1 c2 -> ConvCmp (transposeTKConversion aftk c2)+                           (transposeTKConversion (convertFTK c2 aftk) c1)+  ConvRX | FTKR @n _ x <- aftk+         , Refl <- lemRankReplicate (Proxy @n) ->+    ConvXR (ftkToSTK x)+  ConvSX -> ConvXS+  ConvXR @_ @sh _stk | Refl <- lemRankReplicate (Proxy @(Rank sh)) ->+    ConvCmp (ConvXX' aftk) ConvRX+  ConvXS -> ConvSX+  ConvXS' (FTKS sh _) | Refl <- lemRankMapJust sh ->+    ConvCmp (ConvXX' aftk) ConvSX+  ConvXX' _ftk -> ConvXX' aftk+  ConvRR c | FTKR _ x <- aftk -> ConvRR (transposeTKConversion x c)+  ConvSS c | FTKS _ x <- aftk -> ConvSS (transposeTKConversion x c)+  ConvXX c | FTKX _ x <- aftk -> ConvXX (transposeTKConversion x c)+  ConvT2 c1 c2 | FTKProduct x1 x2 <- aftk ->+    ConvT2 (transposeTKConversion x1 c1) (transposeTKConversion x2 c2)+  Conv0X _stk -> ConvX0+  ConvX0 | FTKX ZSX x <- aftk -> Conv0X (ftkToSTK x)+  ConvNest _stk -> ConvUnnest+  ConvUnnest | (FTKX shx (FTKX _ x)) <- aftk ->+    ConvNest (STKX (ssxFromShX shx) (ftkToSTK x))+  ConvZip stk1 stk2 -> ConvUnzip stk1 stk2+  ConvUnzip stk1 stk2 -> ConvZip stk1 stk2++evalRevFromnMap :: forall target. (ADReadyNoLet target, ShareTensor target)+                => EvalState target -> EvalState target+evalRevFromnMap s@EvalState{nMap, dMap} =+  case DMap.maxViewWithKey nMap of+    Just (n :=> d, nMap2) ->+      let s2 = s {nMap = nMap2}+          s3 = case DMap.lookup n dMap of+            Just (Cotangent c) -> evalRev (nodeIdToFTK n) s2 c d+            Nothing -> error $ "evalRevFromnMap: missing cotangent " ++ show n+      in evalRevFromnMap s3+    Nothing -> s  -- loop ends++{- TODO: optimize similarly?+        -- The general case is given as the last one below,+        -- but for a few constructors it's faster to inline @evalRev@ instead.+        -- BTW, such an optimization doesn't really belong in the simplified+        -- horde-ad and no consistent benefit should be expected here.+        DeltaIndex0 ZeroR{} _ _ -> s  -- shortcut+        DeltaIndex0 (InputR i) ixs' sh ->+          let ixs = indexToList ixs'+              f v = if isTensorDummy v+                    then treplicate0ND sh 0 `OD.update` [(ixs, c)]+                    else v `OD.update` [(ixs, v `rindex0D` ixs + c)]+          in s {iMap = DMap.adjust f i $ iMap s}+        DeltaIndex0 (ShareR n d) ixs' sh ->+          let ixs = indexToList ixs'+          in case DMap.lookup n $ nMap s of+            Just (DynamicRanked _) ->+              let f v = v `OD.update` [(ixs, v `rindex0D` ixs + c)]+              in s {dMap = DMap.adjust f n $ dMap s}+                -- This would be an asymptotic optimization compared to+                -- the general case below, if not for the non-mutable update,+                -- which implies copying the whole @v@ vector,+                -- so it's only several times faster (same allocation,+                -- but not adding to each cell of @v@).+            Nothing ->+              let v = treplicate0ND sh 0 `OD.update` [(ixs, c)]+              in s { nMap = DMap.insert n (DynamicRanked d) $ nMap s+                   , dMap = DMap.insert n v $ dMap s }+            _ -> error "evalRevFromnMap: corrupted nMap"+-}+++-- * Forward derivative computation from the delta expressions++-- | Forward derivative computation via forward-evaluation of delta-expressions+-- (which is surprisingly competitive to the direct forward method,+-- until the allocation of deltas gets large enough to affect cache hits).+-- This is the directional derivative, calculated for the point,+-- at which the delta expression was computed (which is the point+-- represented by the parameters of the objective function and used+-- to compute it's dual number result) and along the direction vector+-- given by the parameters in the arguments.+--+-- This mimics the reverse derivative code, but in reverse. Perhaps this can be+-- simplified, but the obvious simplest formulation does not honour sharing+-- and evaluates shared subexpressions repeatedly, so this state-passing+-- formulation is adopted.+evalFwd+  :: forall target y. (ADReadyNoLet target, ShareTensor target)+  => IMap target -> ADMap target -> Delta target y+  -> (ADMap target, target (ADTensorKind y))+evalFwd params s d0 = case d0 of+  DeltaShare n d ->+    case DMap.lookup n s of+      Just e1 -> (s, unCotangent e1)+      Nothing ->+        let (s2, cRaw) = evalFwd params s d+            cShared = tshare cRaw+            cd = Cotangent cShared+              -- cRaw is shared, because it's put into the map and then+              -- potentially looked up many times, so it'd get duplicated+            s3 = DMap.insert n cd s2+        in (s3, cShared)+  DeltaInput inputId ->+    case DMap.lookup inputId params of+      Just dtk -> (s, toADTensorKindShared (inputIdToFTK inputId)+                      $ evalTensorOrZero dtk)+      Nothing -> error "evalFwd: missing input"++  DeltaPair d1 d2 ->+    let (s2, t) = evalFwd params s d1+        (s3, u) = evalFwd params s2 d2+    in (s3, tpair t u)+  DeltaProject1 d ->+    let (s2, v) = evalFwd params s d+    in (s2, tproject1 v)+  DeltaProject2 d ->+    let (s2, v) = evalFwd params s d+    in (s2, tproject2 v)+  DeltaFromVector snat stk lsd | Refl <- lemBuildOfAD snat stk ->+    let (s2, l) = mapAccumL (evalFwd params) s lsd+    in (s2, tfromVector snat (adSTK stk) l)+  DeltaSum snat stk d | Refl <- lemBuildOfAD snat stk ->+    let (s2, t) = evalFwd params s d+    in (s2, tsum snat (adSTK stk) t)+  DeltaReplicate snat stk d | Refl <- lemBuildOfAD snat stk ->+    let (s2, t) = evalFwd params s d+    in (s2, treplicate snat (adSTK stk) t)+  DeltaMapAccumR k bftk eftk q es df _rf acc0' es'+   | Refl <- lemBuildOfAD k (ftkToSTK bftk)+   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->+    let accftk = ftkDelta acc0'+        accftkAD = adFTK accftk+        bftkAD = adFTK bftk+        eftkAD = adFTK eftk+        (s2, cacc0) = evalFwd params s acc0'+        (s3, ces) = evalFwd params s2 es'+    in (s3, tmapAccumR (Proxy @target)+                       k accftkAD bftkAD (FTKProduct eftkAD+                                                     (FTKProduct accftk eftk))+                       (\dacc de_acc_e ->+                        ttlet de_acc_e $ \ !de_acc_e1 ->+                          unHFun df (tpair (tpair dacc (tproject1 de_acc_e1))+                                           (tproject2 de_acc_e1)))+                       cacc0+                       (tpair ces (tpair q es)))+  DeltaMapAccumL k bftk eftk q es df _rf acc0' es'+   | Refl <- lemBuildOfAD k (ftkToSTK bftk)+   , Refl <- lemBuildOfAD k (ftkToSTK eftk) ->+    let accftk = ftkDelta acc0'+        accftkAD = adFTK accftk+        bftkAD = adFTK bftk+        eftkAD = adFTK eftk+        (s2, cacc0) = evalFwd params s acc0'+        (s3, ces) = evalFwd params s2 es'+    in (s3, tmapAccumL (Proxy @target)+                       k accftkAD bftkAD (FTKProduct eftkAD+                                                     (FTKProduct accftk eftk))+                       (\dacc de_acc_e ->+                        ttlet de_acc_e $ \ !de_acc_e1 ->+                          unHFun df (tpair (tpair dacc (tproject1 de_acc_e1))+                                           (tproject2 de_acc_e1)))+                       cacc0+                       (tpair ces (tpair q es)))++  _ -> let y = ftkDelta d0+           ay = adFTK y+       in case matchingFTK y ay of+         Just Refl -> evalFwdSame params s d0+         _ -> (s, tdefTarget ay)++evalFwdSame+  :: forall target y.+     (ADReadyNoLet target, ShareTensor target, y ~ ADTensorKind y)+  => IMap target -> ADMap target -> Delta target y+  -> (ADMap target, target (ADTensorKind y))+evalFwdSame params s = \case+  DeltaInput inputId ->+    case DMap.lookup inputId params of+      Just dtk -> (s, evalTensorOrZero dtk)+      Nothing -> error "evalFwdSame: missing input"++  -- See the comment about these three in evalRevSame.+  DeltaZero ftk -> (s, tdefTarget $ adFTK ftk)+  DeltaScale (NestedTarget k) d -> second (* k) $ evalFwdSame params s d+  DeltaAdd d e -> let (s2, t) = evalFwdSame params s d+                      (s3, u) = evalFwdSame params s2 e+                  in (s3, t + u)++  d0@(DeltaCastK @r1 d) ->+    case sameSTK (STKScalar @r1) (adSTK (STKScalar @r1)) of+      Just Refl -> second tkcast $ evalFwdSame params s d+      _ -> (s, tdefTarget $ adFTK $ ftkDelta d0)++  d0@(DeltaCastR d) -> case ftkDelta d of+    y -> case matchingFTK y (adFTK y) of+      Just Refl -> second trcast $ evalFwdSame params s d+      _ -> (s, tdefTarget $ adFTK $ ftkDelta d0)+  DeltaSum0R (DeltaZero (FTKR _ x)) -> (s, tdefTarget (FTKR ZSR x))+  DeltaSum0R d -> case ftkDelta d of+    FTKR sh x | SNat <- shrRank sh ->+      withKnownSTK (ftkToSTK x) $+      second trsum0 $ evalFwdSame params s d+  DeltaDot0R _ DeltaZero{} -> (s, trconcrete $ Nested.rscalar 0)+  DeltaDot0R v d -> case ftkDelta d of+    FTKR sh x | SNat <- shrRank sh ->+      withKnownSTK (ftkToSTK x) $+      second (trdot0 v) $ evalFwdSame params s d+  DeltaIndexR SNat d ix -> case ftkDelta d of+    FTKR _ x | SNat <- ixrRank ix ->+      withKnownSTK (ftkToSTK x) $+      second (`trindex` ix) $ evalFwdSame params s d+  DeltaScatterR SNat SNat SNat sh d f -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (ftkToSTK x) $+      let (s2, t) = evalFwdSame params s d+      in (s2, trscatter sh t f)+  DeltaGatherR SNat SNat SNat sh d f -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (ftkToSTK x) $+      let (s2, t) = evalFwdSame params s d+      in (s2, trgather sh t f)+  DeltaAppendR d e -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (ftkToSTK x) $+      let (s2, t) = evalFwdSame params s d+          (s3, u) = evalFwdSame params s2 e+      in (s3, trappend t u)+  DeltaSliceR i n d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (ftkToSTK x) $+      second (trslice i n) $ evalFwdSame params s d+  DeltaReverseR d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (ftkToSTK x) $+      second trreverse $ evalFwdSame params s d+  DeltaTransposeR perm d -> case ftkDelta d of+    FTKR _ x ->+      withKnownSTK (ftkToSTK x) $+      second (trtranspose perm) $ evalFwdSame params s d+  DeltaReshapeR sh2 d -> case ftkDelta d of+    FTKR _sh x ->+      withKnownSTK (ftkToSTK x) $+      second (trreshape sh2) $ evalFwdSame params s d++  d0@(DeltaCastS d) -> case ftkDelta d of+    y -> case matchingFTK y (adFTK y) of+      Just Refl -> second tscast $ evalFwdSame params s d+      _ -> (s, tdefTarget $ adFTK $ ftkDelta d0)+  DeltaSum0S (DeltaZero (FTKS _ x)) -> (s, tdefTarget (FTKS ZSS x))+  DeltaSum0S d -> case ftkDelta d of+    FTKS sh x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShS sh $+      second tssum0 $ evalFwdSame params s d+  DeltaDot0S _ DeltaZero{} -> (s, tsconcrete $ Nested.sscalar 0)+  DeltaDot0S v d -> case ftkDelta d of+    FTKS sh FTKScalar ->+      withKnownShS sh $+      second (tsdot0 v) $ evalFwdSame params s d+  DeltaIndexS shn d ix -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShS shn $+      withKnownShS (shsFromIxS ix) $+      second (`tsindex` ix) $ evalFwdSame params s d+  DeltaScatterS @shm @shn shm shn shp d f -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShS shm $+      withKnownShS shn $+      withKnownShS shp $+      let (s2, t) = evalFwdSame params s d+      in (s2, tsscatter @_ @shm @shn t f)+  DeltaGatherS @shm @shn shm shn shp d f -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShS shm $+      withKnownShS shn $+      withKnownShS shp $+      let (s2, t) = evalFwdSame params s d+      in (s2, tsgather @_ @shm @shn t f)+  DeltaAppendS d e -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      let (s2, t) = evalFwdSame params s d+          (s3, u) = evalFwdSame params s2 e+      in (s3, tsappend t u)+  DeltaSliceS i n k d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      second (tsslice i n k) $ evalFwdSame params s d+  DeltaReverseS d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      second tsreverse $ evalFwdSame params s d+  DeltaTransposeS perm d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      second (tstranspose perm) $ evalFwdSame params s d+  DeltaReshapeS sh2 d -> case ftkDelta d of+    FTKS _ x ->+      withKnownSTK (ftkToSTK x) $+      second (tsreshape sh2) $ evalFwdSame params s d++  d0@(DeltaCastX d) -> case ftkDelta d of+    y -> case matchingFTK y (adFTK y) of+      Just Refl -> second txcast $ evalFwdSame params s d+      _ -> (s, tdefTarget $ adFTK $ ftkDelta d0)+  DeltaSum0X (DeltaZero (FTKX _ x)) -> (s, tdefTarget (FTKX ZSX x))+  DeltaSum0X d -> case ftkDelta d of+    FTKX sh x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShX (ssxFromShX sh) $+      second txsum0 $ evalFwdSame params s d+  DeltaDot0X _ DeltaZero{} -> (s, txconcrete $ Nested.mscalar 0)+  DeltaDot0X v d -> case ftkDelta d of+    FTKX sh FTKScalar ->+      withKnownShX (ssxFromShX sh) $+      second (txdot0 v) $ evalFwdSame params s d+  DeltaIndexX @shm @shn shn d ix -> case ftkDelta d of+    FTKX sh x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShX shn $+-- TODO      withKnownShX (ssxFromIxX ix) $+      withKnownShX (ssxTakeIx @shm @shn (ssxFromShX sh) ix) $+      second (`txindex` ix) $ evalFwdSame params s d+  DeltaScatterX @shm @shn shm shn shp sh d f -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShX shm $+      withKnownShX shn $+      withKnownShX shp $+      let (s2, t) = evalFwdSame params s d+      in (s2, txscatter @_ @shm @shn sh t f)+  DeltaGatherX @shm @shn shm shn shp sh d f -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (ftkToSTK x) $+      withKnownShX shm $+      withKnownShX shn $+      withKnownShX shp $+      let (s2, t) = evalFwdSame params s d+      in (s2, txgather @_ @shm @shn sh t f)+  DeltaAppendX d e -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (ftkToSTK x) $+      let (s2, t) = evalFwdSame params s d+          (s3, u) = evalFwdSame params s2 e+      in (s3, txappend t u)+  DeltaSliceX i n k d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (ftkToSTK x) $+      second (txslice i n k) $ evalFwdSame params s d+  DeltaReverseX d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (ftkToSTK x) $+      second txreverse $ evalFwdSame params s d+  DeltaTransposeX perm d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (ftkToSTK x) $+      second (txtranspose perm) $ evalFwdSame params s d+  DeltaReshapeX sh2 d -> case ftkDelta d of+    FTKX _ x ->+      withKnownSTK (ftkToSTK x) $+      second (txreshape sh2) $ evalFwdSame params s d++  DeltaConvert @a c1 d ->+    -- This follows from the same property for @b@ and from @c1@+    -- not changing the underlying scalar types.+    gcastWith (unsafeCoerceRefl :: ADTensorKind a :~: a) $+    second (tconvert c1 (ftkToSTK (ftkDelta d)))+           (evalFwdSame params s d)++  d -> evalFwd params s d
+ src/HordeAd/Core/DeltaFreshId.hs view
@@ -0,0 +1,68 @@+-- | This module uses and rather safely encapsulates impure side-effects.+-- The impurity produces pure data with a particular property.+-- The property is an order of per-node integer identifiers that represents+-- data dependencies and sharing between delta expressions. The low-level API+-- of AD depends on this property, but is completely isolated from the impurity.+-- The high-level API of AD triggers the impurity but can't observe+-- any impure behaviour. Neither can any other module in the package,+-- except for the testing modules that import testing-exclusive+-- low level operations on the impure counter.+module HordeAd.Core.DeltaFreshId+  ( shareDelta+    -- * Low level counter manipulation to be used only in sequential tests+  , resetIdCounter+  ) where++import Prelude++import Control.Concurrent.Counter (Counter, add, new, set)+import System.IO.Unsafe (unsafePerformIO)++import HordeAd.Core.Delta++unsafeGlobalCounter :: Counter+{-# NOINLINE unsafeGlobalCounter #-}+unsafeGlobalCounter = unsafePerformIO (new 100000001)++-- | Do not use; this is exposed only for special low level tests.+-- e.g., to ensure @show@ applied to terms has a stable length.+-- Tests using this need to be run with -ftest_seq to avoid variable confusion.+resetIdCounter :: IO ()+resetIdCounter = set unsafeGlobalCounter 100000001++-- This is the only operation directly touching the single impure counter+-- that holds fresh and continuously incremented integer identifiers,+--+-- We start at a large number to make tests measuring the size of pretty+-- printed terms less fragile. @Counter@ datatype is just as safe,+-- but faster than an @MVar@ or an atomic @IORef@ (and even non-atomic @IORef@).+-- The operation is manually inlined to prevent GHCs deciding otherwise+-- and causing performance anomalies in benchmarks.+unsafeGetFreshId :: IO Int+{-# INLINE unsafeGetFreshId #-}+unsafeGetFreshId = add unsafeGlobalCounter 1++-- Tests don't show a speedup from `unsafeDupablePerformIO`,+-- perhaps due to counter gaps that it may introduce.+--+-- | The impurity exported from this module by @shareDelta@,+-- stemming from the use of @unsafeGetFreshId@ under @unsafePerformIO@,+-- is thread-safe, admits parallel tests+-- and does not require @-fno-full-laziness@ nor @-fno-cse@.+--+-- The pattern-matching in @shareDelta@ is a crucial optimization+-- and it could be extended to limit which terms get an identifier,+-- trading off sharing for reducing direct memory usage.+shareDelta :: forall y target.+              Delta target y -> Delta target y+{-# NOINLINE shareDelta #-}+shareDelta d = unsafePerformIO $ do+  n <- unsafeGetFreshId+  return $! case d of+    DeltaShare{} -> d  -- should not happen, but older/lower id is safer anyway+    DeltaInput{} -> d+    DeltaPair DeltaShare{} DeltaShare{} -> d  -- all work done inside+    DeltaProject1 DeltaShare{} -> d+    DeltaProject2 DeltaShare{} -> d+    DeltaZero{} -> d+    _ -> DeltaShare (mkNodeId (ftkDelta d) n) d
+ src/HordeAd/Core/Ops.hs view
@@ -0,0 +1,1206 @@+{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists, QuantifiedConstraints,+             UndecidableInstances #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | A collection of classes containing array operations,+-- with some extra algebraic operations and dual numbers+-- operations added in.+--+-- Note that @Ast*@ modules rarely depend on @Ops*@ and @Carriers*@ modules+-- (except for "HordeAd.Core.AstInterpret" and "HordeAd.Core.AstEnv"+-- that describe how to go from @Ast*@ to @Ops*@). Similarly, @Ops*@+-- and @Carriers*@ modules rarely depend on @Ast*@ modules+-- (except for "HordeAd.Core.OpsAst" and "HordeAd.Core.CarriersAst"+-- that describe how to define @Ops*@ in terms of @Ast*@).+-- Syntax is relatively separated from semantics and they meet+-- in the interpreter ("HordeAd.Core.AstInterpret")+-- and in the semantic model constructed from syntax ("HordeAd.Core.OpsAst").+--+-- (A copy of the text above is in "HordeAd.Core.Ast".)+module HordeAd.Core.Ops+  ( -- * The tensor classes and support datatypes+    LetTensor(..), ShareTensor(..), BaseTensor(..), HFun(..)+    -- * The giga-constraint+  , ADReady, ADReadyNoLet, ADReadyEqs, ADReadyClasses, ADReadyEqsClasses+  , AllTargetShow, CommonTargetEqOrd+    -- * Helper functions+  , rtr, rflatten, str, sflatten, xtr, xflatten+  , tmapAccumR, tmapAccumL+  , rbuild, sbuild, xbuild+    -- * Helper classes and types+  , IntegralHAndIntElt, RealFloatAndFloatElt+  , TensorSupportsX, TensorSupportsS, TensorSupportsR, TensorSupports+  ) where++import Prelude++import Data.Foldable qualified as Foldable+import Data.Int (Int64)+import Data.Kind (Constraint, Type)+import Data.Maybe (fromMaybe)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import Data.Vector.Strict qualified as Data.Vector+import GHC.Exts (IsList (..))+import GHC.TypeLits (KnownNat, type (+), type (<=), type (<=?))+import Type.Reflection (typeRep)++import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (withShsFromShX)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (Init, unsafeCoerceRefl)++import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++-- These user API functions are used in default definitions of methods,+-- so they have to be defined already here:++rtr :: forall n x target. (KnownSTK x, BaseTensor target)+    => target (TKR2 (2 + n) x) -> target (TKR2 (2 + n) x)+rtr = trtranspose [1, 0]+rflatten :: forall n x target. (KnownSTK x, BaseTensor target)+         => target (TKR2 n x) -> target (TKR2 1 x)+rflatten u = trreshape (rsize u :$: ZSR) u+str :: forall n m sh x target. (KnownSTK x, BaseTensor target)+    => target (TKS2 (n ': m ': sh) x) -> target (TKS2 (m ': n ': sh) x)+str = gcastWith (unsafeCoerceRefl :: (2 <=? Rank (n ': m ': sh)) :~: True) $+      tstranspose (Permutation.makePerm @'[1, 0])+sflatten :: (KnownShS sh, KnownSTK x, BaseTensor target )+         => target (TKS2 sh x) -> target (TKS2 '[Product sh] x)+sflatten @sh | SNat <- shsProduct (knownShS @sh) = tsreshape knownShS+xtr :: forall n m sh x target. (KnownSTK x, BaseTensor target)+    => target (TKX2 (Just n ': Just m ': sh) x)+    -> target (TKX2 (Just m ': Just n ': sh) x)+xtr = gcastWith (unsafeCoerceRefl+                 :: (2 <=? Rank (Just n ': Just m ': sh)) :~: True) $+      txtranspose (Permutation.makePerm @'[1, 0])+xflatten :: forall sh x target. (KnownSTK x, BaseTensor target)+         => target (TKX2 sh x) -> target (TKX2 '[Nothing] x)+xflatten u = txreshape (Nested.SUnknown (xsize u) :$% ZSX) u++rbuild :: (KnownNat m, KnownNat n, KnownSTK x, BaseTensor target)+       => IShR (m + n)  -- ^ the shape of the resulting tensor+       -> (IxROf target m -> target (TKR2 n x))+            -- ^ the function to build with+       -> target (TKR2 (m + n) x)+rbuild @m @n @x @target sh0 f0 =+  let buildSh :: IShR m1 -> (IxROf target m1 -> target (TKR2 n x))+              -> target (TKR2 (m1 + n) x)+      buildSh ZSR f = f ZIR+      buildSh (k :$: sh) f | SNat <- shrRank sh =+        let g i = buildSh sh (\ix -> f (i :.: ix))+        in trbuild1 k g+  in buildSh (shrTake @m @n sh0) f0+sbuild :: (KnownShS (Take m sh), KnownShS sh, KnownSTK x, BaseTensor target)+       => (IxSOf target (Take m sh) -> target (TKS2 (Drop m sh) x))+            -- ^ the function to build with+       -> target (TKS2 sh x)+sbuild @m @sh @x @target =+  let buildSh+        :: forall sh1.+           ShS sh1 -> ShS (sh1 ++ Drop m sh)+        -> (IxSOf target sh1 -> target (TKS2 (Drop m sh) x))+        -> target (TKS2 (sh1 ++ Drop m sh) x)+      buildSh sh1 sh1m f = case (sh1, sh1m) of+        (ZSS, _) -> f ZIS+        (SNat :$$ sh2, _ :$$ sh2m) ->+          withKnownShS sh2m $+          let g i = buildSh sh2 sh2m (f . (i :.$))+          in tsbuild1 g+  in gcastWith (unsafeCoerceRefl :: sh :~: Take m sh ++ Drop m sh)+     $ buildSh (knownShS @(Take m sh)) (knownShS @sh)+xbuild :: ( KnownShX (Take m sh), KnownSTK x+          , BaseTensor target, ConvertTensor target )+       => IShX sh  -- ^ the shape of the resulting tensor+       -> (IxXOf target (Take m sh) -> target (TKX2 (Drop m sh) x))+            -- ^ the function to build with+       -> target (TKX2 sh x)+xbuild @m @sh @x @target sh0 f0 =+  let buildSh :: IShX sh1 -> IShX (sh1 ++ Drop m sh)+              -> (IxXOf target sh1 -> target (TKX2 (Drop m sh) x))+              -> target (TKX2 (sh1 ++ Drop m sh) x)+      buildSh sh1 sh1m f = case (sh1, sh1m) of+        (ZSX, _) -> f ZIX+        (k :$% sh2, _ :$% sh2m) ->+          withKnownShX (ssxFromShX sh2m) $+          let g i = buildSh sh2 sh2m (f . (i :.%))+          in withSNat (fromSMayNat' k) $ \(SNat @n) ->+               xmcast (ssxFromShX sh1m) $ txbuild1 @_ @n g+  in gcastWith (unsafeCoerceRefl :: sh :~: Take m sh ++ Drop m sh)+     $ buildSh (shxTakeSSX (Proxy @(Drop m sh)) (knownShX @(Take m sh)) sh0) sh0 f0++-- | A strict right mapAccum.+tmapAccumR+  :: forall accy by ey k target. BaseTensor target+  => Proxy target+  -> SNat k  -- ^ length of the input+  -> FullShapeTK accy  -- ^ shape of the accumulator+  -> FullShapeTK by  -- ^ shape of the output+  -> FullShapeTK ey  -- ^ shape of an individual input+  -> (forall f. ADReady f+      => f accy -> f ey -> f (TKProduct accy by))+       -- ^ the function to mapAccum with+  -> target accy  -- ^ the initial accumulator+  -> target (BuildTensorKind k ey)  -- ^ the inputs+  -> target (TKProduct accy (BuildTensorKind k by))+{-# INLINE tmapAccumR #-}  -- this doesn't want to specialize+tmapAccumR proxy !k !accftk !bftk !eftk f acc0 es =+  let xftk = FTKProduct accftk eftk+      fl :: forall f. ADReady f+         => f (TKProduct accy ey)+         -> f (TKProduct accy by)+      fl !args = ttlet args $ \ !args1 ->+                   f (tproject1 args1) (tproject2 args1)+  in tmapAccumRDer proxy k accftk bftk eftk+                   (tlambda @target xftk (HFun fl))+                   (tjvp @target xftk $ HFun fl)+                   (tvjp @target xftk $ HFun fl)+                   acc0 es+-- | A strict left mapAccum.+tmapAccumL+  :: forall accy by ey k target. BaseTensor target+  => Proxy target+  -> SNat k  -- ^ length of the input+  -> FullShapeTK accy  -- ^ shape of the accumulator+  -> FullShapeTK by  -- ^ shape of the output+  -> FullShapeTK ey  -- ^ shape of an individual input+  -> (forall f. ADReady f+      => f accy -> f ey -> f (TKProduct accy by))+       -- ^ the function to mapAccum with+  -> target accy  -- ^ the initial accumulator+  -> target (BuildTensorKind k ey)  -- ^ the inputs+  -> target (TKProduct accy (BuildTensorKind k by))+{-# INLINE tmapAccumL #-}  -- this doesn't want to specialize+tmapAccumL proxy !k !accftk !bftk !eftk f acc0 es =+  let xftk = FTKProduct accftk eftk+      fl :: forall f. ADReady f+         => f (TKProduct accy ey)+         -> f (TKProduct accy by)+      fl !args = ttlet args $ \ !args1 ->+                   f (tproject1 args1) (tproject2 args1)+  in tmapAccumLDer proxy k accftk bftk eftk+                   (tlambda @target xftk (HFun fl))+                   (tjvp @target xftk $ HFun fl)+                   (tvjp @target xftk $ HFun fl)+                   acc0 es++type TensorSupports :: (Type -> Constraint) -> (Type -> Constraint)+                    -> Target -> Constraint+type TensorSupports c1 c2 f =+  forall r. GoodScalar r+            => c1 r => c2 (f (TKScalar r))++type TensorSupportsR :: (Type -> Constraint) -> (Type -> Constraint)+                     -> Target -> Constraint+type TensorSupportsR c1 c2 f =+  forall r n. GoodScalar r+              => c1 r => c2 (f (TKR n r))++type TensorSupportsS :: (Type -> Constraint) -> (Type -> Constraint)+                     -> Target -> Constraint+type TensorSupportsS c1 c2 f =+  forall r sh. GoodScalar r+               => c1 r => c2 (f (TKS sh r))++type TensorSupportsX :: (Type -> Constraint) -> (Type -> Constraint)+                     -> Target -> Constraint+type TensorSupportsX c1 c2 f =+  forall r sh. GoodScalar r+               => c1 r => c2 (f (TKX sh r))++class (RealFloatH r, Nested.FloatElt r)+      => RealFloatAndFloatElt r+instance (RealFloatH r, Nested.FloatElt r)+         => RealFloatAndFloatElt r++class (IntegralH r, Nested.IntElt r)+      => IntegralHAndIntElt r+instance (IntegralH r, Nested.IntElt r)+      => IntegralHAndIntElt r++class LetTensor (target :: Target) where+  ttlet :: target x -> (target x -> target z) -> target z+  ttletPrimal :: PrimalOf target x -> (PrimalOf target x -> target z)+              -> target z+  toShare :: target y -> ShareOf target y+  tunshare :: ShareOf target y -> target y+  tunshare = error "tunshare: this instance should never be used"+  tappend :: forall m n y. BaseTensor target+          => SNat m -> SNat n -> SingletonTK y+          -> target (BuildTensorKind m y) -> target (BuildTensorKind n y)+          -> target (BuildTensorKind (m + n) y)+  tappend msnat@SNat nsnat@SNat stk a b = case stk of+    STKScalar -> tsappend a b+    STKR _ x | Dict <- lemKnownSTK x -> trappend a b+    STKS _ x | Dict <- lemKnownSTK x -> tsappend a b+    STKX _ x | Dict <- lemKnownSTK x -> txappend a b+    STKProduct stk1 stk2 ->+      ttlet a $ \ !aShared -> ttlet b $ \ !bShared ->+        tpair (tappend msnat nsnat stk1 (tproject1 aShared) (tproject1 bShared))+              (tappend msnat nsnat stk2 (tproject2 aShared) (tproject2 bShared))+  tD :: BaseTensor target+     => SingletonTK y -> PrimalOf target y -> DualOf target y+     -> target y+  tD stk p d =+    -- Lets needed, because taddTarget requires duplicable arguments.+    ttletPrimal p $ \pShared ->+    ttlet (tfromDual d) $ \dShared ->+      taddTarget stk (tfromPrimal stk pShared) dShared+  -- | A strict left fold.+  tfold+    :: forall yn ym k. BaseTensor target+    => SNat k  -- ^ length of the input+    -> SingletonTK yn  -- ^ partial shape of the accumulator+    -> SingletonTK ym  -- ^ partial shape of an individual input+    -> (forall f. ADReady f => f yn -> f ym -> f yn)+         -- ^ the function to fold with+    -> target yn  -- ^ the initial accumulator+    -> target (BuildTensorKind k ym)  -- ^ the inputs+    -> target yn+  {-# INLINE tfold #-}  -- this doesn't want to specialize+  tfold k nstk mstk f acc0 es =+    tproject1+    $ tmapAccumL (Proxy @target)+       k+       (tftk nstk acc0)+       (FTKScalar @Z1)+       (razeFTK k mstk (tftk (buildSTK k mstk) es))+       (let g :: forall f. ADReady f+              => f yn -> f ym -> f (TKProduct yn TKUnit)+            g !acc !e = tpair (f acc e) (tkconcrete Z1)+        in g)+       acc0+       es+  -- | A strict left scan.+  tscan+    :: forall yn ym k. BaseTensor target+    => SNat k  -- ^ length of the input+    -> SingletonTK yn  -- ^ partial shape of the accumulator+    -> SingletonTK ym  -- ^ partial shape of an individual input+    -> (forall f. ADReady f => f yn -> f ym -> f yn)+         -- ^ the function to scan with+    -> target yn  -- ^ the initial accumulator+    -> target (BuildTensorKind k ym)  -- ^ the inputs+    -> target (BuildTensorKind (1 + k) yn)+  {-# INLINE tscan #-}  -- this doesn't want to specialize+  tscan k nstk mstk f acc0 es =+    let bs :: target (BuildTensorKind k yn)+        bs = tproject2+             $ tmapAccumL (Proxy @target)+                k+                (tftk nstk acc0)+                (tftk nstk acc0)+                (razeFTK k mstk (tftk (buildSTK k mstk) es))+              (let g :: forall f. ADReady f+                     => f yn -> f ym -> f (TKProduct yn yn)+                   g !acc !e = ttlet (f acc e) $ \ !res -> tpair res res+               in g)+              acc0+              es+    in tappend (SNat @1) k nstk+               (tfromVector (SNat @1) nstk (V.fromList [acc0])) bs++class ShareTensor (target :: Target) where+  tshare :: target y -> target y+  tunpair :: target (TKProduct x z) -> (target x, target z)+  -- This would suffer from lack of sharing with LetTensor, because+  -- ttlet doesn't work over a list. With sharing it's fine.+  tunravelToListShare :: forall y k. (BaseTensor target, ConvertTensor target)+                      => SNat k -> SingletonTK y+                      -> target (BuildTensorKind k y)+                      -> [target y]+  tunravelToListShare snat@SNat stk u = case stk of+    STKScalar -> let !uShared = tshare u+                 in map kfromS $ tsunravelToList uShared+    STKR SNat x | Dict <- lemKnownSTK x -> let !uShared = tshare u+                                           in trunravelToList uShared+    STKS sh x | Dict <- lemKnownSTK x -> let !uShared = tshare u+                                         in withKnownShS sh+                                            $ tsunravelToList uShared+    STKX sh x | Dict <- lemKnownSTK x -> let !uShared = tshare u+                                         in withKnownShX sh+                                            $ txunravelToList uShared+    STKProduct stk1 stk2 ->+      let (!u1, !u2) = tunpair u+      in zipWith tpair (tunravelToListShare snat stk1 u1)+                       (tunravelToListShare snat stk2 u2)++-- | The superclasses indicate that it's not only a container array,+-- but also a mathematical tensor, sporting numeric operations.+class ( Num (IntOf target)+      , IntegralH (IntOf target)+      , TensorSupports Num Num target+      , TensorSupports RealFloatAndFloatElt Floating target+      , TensorSupports RealFloatAndFloatElt RealFloatH target+      , TensorSupports IntegralHAndIntElt IntegralH target+      , TensorSupportsR Num Num target+      , TensorSupportsR RealFloatAndFloatElt Floating target+      , TensorSupportsR RealFloatAndFloatElt RealFloatH target+      , TensorSupportsR IntegralHAndIntElt IntegralH target+      , TensorSupportsS Num Num target+      , TensorSupportsS RealFloatAndFloatElt Floating target+      , TensorSupportsS RealFloatAndFloatElt RealFloatH target+      , TensorSupportsS IntegralHAndIntElt IntegralH target+      , TensorSupportsX Num Num target+      , TensorSupportsX RealFloatAndFloatElt Floating target+      , TensorSupportsX RealFloatAndFloatElt RealFloatH target+      , TensorSupportsX IntegralHAndIntElt IntegralH target )+      => BaseTensor (target :: Target) where++  -- First type argument being @target@ is acceptable here, since these+  -- operations are mostly used when the shape is not known at the type level,+  -- so it can't be used as an explicit type argument.+  rshape :: forall n x. KnownSTK x+         => target (TKR2 n x) -> IShR n+  rlength :: forall n x. KnownSTK x+          => target (TKR2 n x) -> Int+  rlength = shrLength . rshape+  rsize :: forall n x. KnownSTK x+        => target (TKR2 n x) -> Int+  rsize = shrSize . rshape+  rwidth :: forall n x. KnownSTK x+          => target (TKR2 (1 + n) x) -> Int+  rwidth a = case rshape a of+    k :$: _ -> k++  sshape :: forall sh x. KnownSTK x+         => target (TKS2 sh x) -> ShS sh+  slength :: forall sh x. KnownSTK x+          => target (TKS2 sh x) -> Int+  slength = shsLength . sshape+  ssize :: forall sh x. KnownSTK x+        => target (TKS2 sh x) -> Int+  ssize = shsSize . sshape+  swidth :: forall n sh x. KnownSTK x+          => target (TKS2 (n ': sh) x) -> Int+  swidth a = case sshape a of+    n :$$ _ -> sNatValue n++  xshape :: forall sh x. KnownSTK x+         => target (TKX2 sh x) -> IShX sh+  xlength :: forall sh x. KnownSTK x+          => target (TKX2 sh x) -> Int+  xlength = shxLength . xshape+  xsize :: forall sh x. KnownSTK x+        => target (TKX2 sh x) -> Int+  xsize = shxSize . xshape+  xwidth :: forall mn sh x. KnownSTK x+          => target (TKX2 (mn ': sh) x) -> Int+  xwidth a = case xshape a of+    mn :$% _ -> fromSMayNat' mn++  tsize :: SingletonTK y -> target y -> Int+  tsize stk a = case stk of+    STKScalar @r -> case testEquality (typeRep @r) (typeRep @Z1) of+      Just Refl -> 0+      _ -> 1+    STKR _ x | Dict <- lemKnownSTK x -> rsize a+    STKS _ x | Dict <- lemKnownSTK x -> ssize a+    STKX _ x | Dict <- lemKnownSTK x -> xsize a+    STKProduct stk1 stk2 ->+      tsize stk1 (tproject1 a) + tsize stk2 (tproject2 a)+  tftk :: SingletonTK y -> target y -> FullShapeTK y++  -- Unlikely to require type applications at all+  tpair :: target x -> target z -> target (TKProduct x z)+  tproject1 :: target (TKProduct x z) -> target x+  tproject2 :: target (TKProduct x z) -> target z++  -----------+  -- Everything below is indended to be rarely used and usually there are+  -- more specific and/or more convienient functions that do the same job+  -- in other modules.+  -----------------++  -- | The operation is potentially strict in all arguments.+  tcond :: Boolean (BoolOf target)+        => SingletonTK y+        -> BoolOf target -> target y -> target y -> target y++  -- A more precise type would have `PrimalOf target`, but it's require+  -- the user to convert, so we leave that precision to the AST only+  -- which means the AST instance will automatically insert such+  -- conversions as needed. The same holds for trfloor and many others.+  trconcrete :: GoodScalar r+             => Nested.Ranked n r -> target (TKR n r)+  tsconcrete :: GoodScalar r+             => Nested.Shaped sh r -> target (TKS sh r)+  txconcrete :: GoodScalar r+             => Nested.Mixed sh r -> target (TKX sh r)+  tkconcrete :: GoodScalar r => r -> target (TKScalar r)+  tconcrete :: FullShapeTK y -> Concrete y -> target y++  -- These nine methods can't be replaced by tfromVector, because the concrete+  -- instance has much faster implementations.+  --+  -- This is morally non-empty strict vectors:+  trfromVector :: (KnownNat n, KnownSTK x)+               => Data.Vector.Vector (target (TKR2 n x))+               -> target (TKR2 (1 + n) x)+  trfromVector v = withSNat (V.length v) $ \k ->+    tfromVector k (STKR SNat knownSTK) v+  trfromVector0N :: forall n x. KnownSTK x+                 => IShR n -> Data.Vector.Vector (target (TKR2 0 x))+                 -> target (TKR2 n x)+  trfromVector0N sh v | Dict <- eltDictRep (knownSTK @x) =+    if V.null v+    then let arr = Nested.remptyArray+         in trreshape sh $ tconcrete (tftkG knownSTK arr) (Concrete arr)+    else trreshape sh $ trfromVector v+  trunravelToList :: (KnownNat n, KnownSTK x)+                  => target (TKR2 (1 + n) x) -> [target (TKR2 n x)]+  trunravelToList @n @x t =+    let f :: Int -> target (TKR2 n x)+        f i = trindex t (fromIntegral i :.: ZIR)+    in map f [0 .. rwidth t - 1]++  tsfromVector :: (KnownNat n, KnownShS sh, KnownSTK x)+               => Data.Vector.Vector (target (TKS2 sh x))+               -> target (TKS2 (n ': sh) x)+  tsfromVector = tfromVector SNat (STKS knownShS knownSTK)+  tsfromVector0N :: (KnownShS sh, KnownSTK x)+                 => Data.Vector.Vector (target (TKS2 '[] x))+                 -> target (TKS2 sh x)+  tsfromVector0N @sh @x v | Dict <- eltDictRep (knownSTK @x)+                          , SNat <- shsProduct (knownShS @sh) =+    if V.null v+    then gcastWith (unsafeCoerceRefl :: Product sh :~: 0) $+         let arr = Nested.semptyArray ZSS+         in tsreshape knownShS $ tconcrete (tftkG knownSTK arr) (Concrete arr)+    else tsreshape (knownShS @sh) $ tsfromVector v+  tsunravelToList :: (KnownNat n, KnownShS sh, KnownSTK x)+                  => target (TKS2 (n ': sh) x) -> [target (TKS2 sh x)]+  tsunravelToList @_ @sh @x t =+    let f :: Int -> target (TKS2 sh x)+        f i = tsindex t (fromIntegral i :.$ ZIS)+    in map f [0 .. swidth t - 1]++  txfromVector :: (KnownNat n, KnownShX sh, KnownSTK x)+               => Data.Vector.Vector (target (TKX2 sh x))+               -> target (TKX2 (Just n ': sh) x)+  txfromVector = tfromVector SNat (STKX knownShX knownSTK)+  txfromVector0N :: forall sh x. KnownSTK x+                 => IShX sh -> Data.Vector.Vector (target (TKX2 '[] x))+                 -> target (TKX2 sh x)+  txfromVector0N sh v | Dict <- eltDictRep (knownSTK @x) =+    if V.null v+    then let arr = Nested.memptyArray ZSX+         in txreshape sh $ tconcrete (tftkG knownSTK arr) (Concrete arr)+    else withSNat (shxSize sh) $ \(SNat @n) ->+           txreshape @_ @'[Just n] sh $ txfromVector v+  txunravelToList :: (KnownNat n, KnownShX sh, KnownSTK x)+                  => target (TKX2 (Just n ': sh) x) -> [target (TKX2 sh x)]+  txunravelToList @_ @sh @x t =+    let f :: Int -> target (TKX2 sh x)+        f i = txindex t (fromIntegral i :.% ZIX)+    in map f [0 .. xwidth t - 1]++  tfromVector+    :: forall y k.+       SNat k -> SingletonTK y -> Data.Vector.Vector (target y)+    -> target (BuildTensorKind k y)+  tfromListR :: SingletonTK y -> ListR k (target y)+             -> target (BuildTensorKind k y)+  tfromListR stk l =+    tfromVector (listrRank l) stk . V.fromList . Foldable.toList $ l++  -- A number suffix in the name may indicate the rank of the codomain,+  -- if bounded. Suffix 1 may also mean the operations builds up codomain+  -- by 1 dimension.+  trsum :: (KnownNat n, KnownSTK x)+        => target (TKR2 (1 + n) x) -> target (TKR2 n x)+  -- This op (and it's Delta constructor) is worthwhile, because flattening+  -- is O(n) sometimes, unlike transpose, etc.+  trsum0 :: (KnownNat n, KnownSTK x)+         => target (TKR2 n x) -> target (TKR2 0 x)+  trsum0 = trsum . rflatten+  trdot0 :: (KnownNat n, GoodScalar r)+         => target (TKR n r) -> target (TKR n r) -> target (TKR 0 r)+  trdot0 t u = trsum (rflatten (t * u))+  trdot1In :: (KnownNat n, GoodScalar r)+           => target (TKR (1 + n) r) -> target (TKR (1 + n) r)+           -> target (TKR n r)+  trdot1In @n t u = trsum $ trtranspose (permCycle $ 1 + valueOf @n) (t * u)+  trmatvecmul :: GoodScalar r+              => target (TKR 2 r) -> target (TKR 1 r) -> target (TKR 1 r)+-- How to generalize (#69)? The few straightforward generalizations+-- differ in types but all are far from matmul2.+-- rmatvecmul m v = rflatten $ rmap1 (rreplicate 1 . rdot0 v) m+  trmatvecmul m v = trbuild1 (rwidth m) (\i -> trdot0 v (m `trindex` [i]))+  trmatmul2 :: GoodScalar r+            => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)+-- How to generalize to tmatmul (#69)?+-- Just rmatmul2 the two outermost dimensions?+-- rmatmul2 m1 m2 = rmap1 (rmatvecmul (rtr m2)) m1+  trmatmul2 m1 m2 =+    trbuild1 (rwidth m1) (\i -> trmatvecmul (rtr m2) (m1 `trindex` [i]))+  trreplicate :: (KnownNat n, KnownSTK x)+              => Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)+  trreplicate0N :: (KnownNat n, KnownSTK x)+                => IShR n -> target (TKR2 0 x) -> target (TKR2 n x)+  trreplicate0N sh = trreshape sh . trreplicate (shrSize sh)++  tssum :: (KnownNat n, KnownShS sh, KnownSTK x)+        => target (TKS2 (n ': sh) x) -> target (TKS2 sh x)+  tssum0 :: (KnownShS sh, KnownSTK x)+         => target (TKS2 sh x) -> target (TKS2 '[] x)+  tssum0 @sh | SNat <- shsProduct (knownShS @sh) = tssum . sflatten+  tsdot0 :: (KnownShS sh, GoodScalar r)+         => target (TKS sh r) -> target (TKS sh r) -> target (TKS '[] r)+  tsdot0 @sh t u | SNat <- shsProduct (knownShS @sh) = tssum (sflatten (t * u))+  tsdot1In :: (KnownShS sh, GoodScalar r)+           => SNat n -> target (TKS (sh ++ '[n]) r)+           -> target (TKS (sh ++ '[n]) r)+           -> target (TKS sh r)+  tsdot1In @sh (SNat @n) t u =+    let cpermR = permCycle $ 1 + shsLength (knownShS @sh)+    in Permutation.permFromList cpermR $ \(cperm :: Permutation.Perm cperm) ->+         gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (sh ++ '[n])) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix cperm (sh ++ '[n])+                       :~: n : sh) $+         fromMaybe (error "tsdot1In: impossible non-permutation")+         $ Permutation.permCheckPermutation cperm+         $ tssum $ tstranspose cperm (t * u)+  tsmatvecmul :: (KnownNat m, KnownNat n, GoodScalar r)+              => target (TKS '[m, n] r) -> target (TKS '[n] r)+              -> target (TKS '[m] r)+  tsmatvecmul @m m v = tsbuild1 @_ @m (\i -> tsdot0 v (m `tsindex` (i :.$ ZIS)))+  tsmatmul2 :: (KnownNat m, KnownNat n, KnownNat p, GoodScalar r)+            => target (TKS '[m, n] r) -> target (TKS '[n, p] r)+            -> target (TKS '[m, p] r)+  tsmatmul2 @m m1 m2 =+    tsbuild1 @_ @m (\i -> tsmatvecmul (str m2) (m1 `tsindex` (i :.$ ZIS)))+  tsreplicate :: forall sh k x. KnownSTK x+              => SNat k -> ShS sh -> target (TKS2 sh x)+              -> target (TKS2 (k ': sh) x)+  tsreplicate0N :: forall sh x. KnownSTK x+                => ShS sh -> target (TKS2 '[] x)+                -> target (TKS2 sh x)+  tsreplicate0N sh = tsreshape sh . tsreplicate (shsProduct sh) ZSS++  -- The choice in BuildTensorKind makes it hard to support this one,+  -- due to DeltaSum and AstSum being typed with BuildTensorKind:+  -- xsum :: (KnownShX sh, KnownShX (mn ': sh), KnownSTK x)+  --     => target (TKX2 (mn ': sh) x) -> target (TKX2 sh x)+  txsum :: (KnownNat n, KnownShX sh, KnownSTK x)+        => target (TKX2 (Just n ': sh) x) -> target (TKX2 sh x)+  txsum0 :: (KnownShX sh, KnownSTK x, ConvertTensor target)+         => target (TKX2 sh x) -> target (TKX2 '[] x)+  txsum0 t = withSNat (shxSize $ xshape t) $ \snat ->+    txsum (xmcast (Nested.SKnown snat :!% ZKX) $ xflatten t)+  txdot0 :: (KnownShX sh, GoodScalar r, ConvertTensor target)+         => target (TKX sh r) -> target (TKX sh r) -> target (TKX '[] r)+  txdot0 t u = withSNat (shxSize $ xshape t) $ \snat ->+    txsum (xmcast (Nested.SKnown snat :!% ZKX) $ xflatten (t * u))+  txdot1In :: (KnownShX sh, GoodScalar r)+           => SNat n -> target (TKX (sh ++ '[Just n]) r)+           -> target (TKX (sh ++ '[Just n]) r)+           -> target (TKX sh r)+  txdot1In @sh (SNat @n) t u =+    let cpermR = permCycle $ 1 + sNatValue (ssxRank (knownShX @sh))+    in Permutation.permFromList cpermR $ \(cperm :: Permutation.Perm cperm) ->+         gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (sh ++ '[Just n])) $+         gcastWith (unsafeCoerceRefl+                    :: Permutation.PermutePrefix cperm (sh ++ '[Just n])+                       :~: Just n : sh) $+         fromMaybe (error "txdot1In: impossible non-permutation")+         $ Permutation.permCheckPermutation cperm+         $ txsum $ txtranspose cperm (t * u)+  txmatvecmul :: forall mm mn r. (GoodScalar r, ConvertTensor target)+              => Nested.SMayNat Int SNat mm -> Nested.SMayNat Int SNat mn+              -> target (TKX '[mm, mn] r) -> target (TKX '[mn] r)+              -> target (TKX '[mm] r)+  txmatvecmul mm mn m v =+    withKnownShX (ssxFromShX $ mm :$% ZSX) $+    withKnownShX (ssxFromShX $ mn :$% ZSX) $+    withSNat (fromSMayNat' mm) $ \(SNat @k) ->+      xmcast (ssxFromShX $ mm :$% ZSX)+      $ txbuild1 @_ @k (\i -> txdot0 v (m `txindex` (i :.% ZIX)))+  txmatmul2 :: ( KnownNat m, KnownNat n, KnownNat p, GoodScalar r+               , ConvertTensor target )+            => target (TKX '[Just m, Just n] r)+            -> target (TKX '[Just n, Just p] r)+            -> target (TKX '[Just m, Just p] r)+  txmatmul2 @m @n @p m1 m2 =+    txbuild1 @_ @m (\i ->+      txmatvecmul (Nested.SKnown (SNat @p)) (Nested.SKnown (SNat @n))+                  (xtr m2) (m1 `txindex` (i :.% ZIX)))+  txreplicate :: forall sh k x. KnownSTK x+              => SNat k -> StaticShX sh -> target (TKX2 sh x)+              -> target (TKX2 (Just k ': sh) x)+  txreplicate0N :: (KnownShX sh, KnownSTK x)+                => IShX sh -> target (TKX2 '[] x) -> target (TKX2 sh x)+  txreplicate0N sh = withSNat (shxSize sh) $ \snat ->+    txreshape sh . txreplicate snat knownShX++  trindex :: (KnownNat m, KnownNat n, KnownSTK x)+          => target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)+  trindex0 :: (KnownNat m, KnownSTK x)+           => target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)+  trindex0 = trindex+  troneHot :: ( KnownNat m, KnownNat n, KnownSTK x+              , BoolOf (PrimalOf target) ~ BoolOf target+              , EqH (PrimalOf target) (TKScalar Int64))+           => IShR m -> target (TKR2 n x) -> IxROf target m+           -> target (TKR2 (m + n) x)+  {-# INLINE troneHot #-}+  troneHot @_ @_ @x sh v ix = case knownSTK @x of+    STKScalar ->+      trscatter @_ @0 (shrAppend sh (rshape v)) v (const ix)+    _ -> case tftk knownSTK v of+      FTKR _ ftk2 ->+        -- TODO: def at out of bounds+        let f ix2 = tcond knownSTK+                          (foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true+                           $ zip (toList ix) (toList ix2))+                          (trindex0 v (ixrDrop ix2))+                          (tdefTarget (FTKR ZSR ftk2))+        in rbuild (shrAppend sh (rshape v)) f+           -- TODO: if this is used often, maybe express this as the gather that+           -- would come out of vectorization, making sure it simplifies well+  trscatter :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x)+            => IShR (p + n) -> target (TKR2 (m + n) x)+            -> (IxROf target m -> IxROf target p)+            -> target (TKR2 (p + n) x)+  trscatter1 :: (KnownNat n, KnownNat p, KnownSTK x)+             => IShR (p + n) -> target (TKR2 (1 + n) x)+             -> (IntOf target -> IxROf target p)+             -> target (TKR2 (p + n) x)+  trscatter1 sh v f = trscatter @target @1 sh v (\(i :.: ZIR) -> f i)+  trgather :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x)+           => IShR (m + n) -> target (TKR2 (p + n) x)+           -> (IxROf target m -> IxROf target p)+           -> target (TKR2 (m + n) x)+  trgather1 :: (KnownNat n, KnownNat p, KnownSTK x)+            => Int -> target (TKR2 (p + n) x)+            -> (IntOf target -> IxROf target p)+            -> target (TKR2 (1 + n) x)+  trgather1 k v f = trgather @target @1+                             (k :$: shrDrop (rshape v)) v+                             (\(i :.: ZIR) -> f i)++  tsindex :: (KnownShS shm, KnownShS shn, KnownSTK x)+          => target (TKS2 (shm ++ shn) x) -> IxSOf target shm+          -> target (TKS2 shn x)+  tsindex0 :: (KnownShS sh1, KnownSTK x)+           => target (TKS2 sh1 x) -> IxSOf target sh1+           -> target (TKS2 '[] x)+  tsindex0 @sh1 | Refl <- lemAppNil @sh1 = tsindex+  tsoneHot :: ( KnownShS sh1, KnownShS sh2, KnownSTK x+              , BoolOf (PrimalOf target) ~ BoolOf target+              , EqH (PrimalOf target) (TKScalar Int64) )+           => target (TKS2 sh2 x) -> IxSOf target sh1+           -> target (TKS2 (sh1 ++ sh2) x)+  {-# INLINE tsoneHot #-}  -- this doesn't want to specialize+  tsoneHot @sh1 @sh2 @x v ix+   | SNat <- shsRank (knownShS @sh1) = case knownSTK @x of+    STKScalar ->+      gcastWith (unsafeCoerceRefl :: Take (Rank sh1) (sh1 ++ sh2) :~: sh1) $+      gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) (sh1 ++ sh2) :~: sh2) $+      tsscatter @_ @'[] @_ @sh1 v (const ix)+    _ -> case tftk knownSTK v of+      FTKS _ ftk2 ->+        -- TODO: def at out of bounds+        gcastWith (unsafeCoerceRefl+                   :: Drop (Rank (sh1 ++ sh2)) (sh1 ++ sh2) :~: '[]) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank (sh1 ++ sh2)) (sh1 ++ sh2) :~: (sh1 ++ sh2)) $+        gcastWith (unsafeCoerceRefl+                   :: Drop (Rank sh1) (sh1 ++ sh2) :~: sh2) $+        withKnownShS (knownShS @sh1 `shsAppend` knownShS @sh2) $+        let f ix2 = tcond knownSTK+                          (foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true+                           $ zip (Foldable.toList ix) (Foldable.toList ix2))+                          (tsindex0 v (ixsDrop @(Rank sh1) ix2))+                          (tdefTarget (FTKS ZSS ftk2))+        in sbuild @(Rank (sh1 ++ sh2)) f+  tsscatter+     :: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x)+     => target (TKS2 (shm ++ shn) x)+     -> (IxSOf target shm -> IxSOf target shp)+     -> target (TKS2 (shp ++ shn) x)+  tsscatter1+     :: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x)+     => target (TKS2 (n2 ': shn) x)+     -> (IntOf target -> IxSOf target shp)+     -> target (TKS2 (shp ++ shn) x)+  tsscatter1 @n2 v f = tsscatter @_ @'[n2] v (\(i :.$ _) -> f i)+  tsgather+     :: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x)+     => target (TKS2 (shp ++ shn) x)+     -> (IxSOf target shm -> IxSOf target shp)+     -> target (TKS2 (shm ++ shn) x)+  tsgather1+     :: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x)+     => target (TKS2 (shp ++ shn) x)+     -> (IntOf target -> IxSOf target shp)+     -> target (TKS2 (n2 ': shn) x)+  tsgather1 @n2 v f = tsgather @target @'[n2] v (\(i :.$ _) -> f i)++  txindex :: (KnownShX sh1, KnownShX sh2, KnownSTK x)+          => target (TKX2 (sh1 ++ sh2) x) -> IxXOf target sh1+          -> target (TKX2 sh2 x)+  txindex0 :: (KnownShX sh1, KnownSTK x)+           => target (TKX2 sh1 x) -> IxXOf target sh1+           -> target (TKX2 '[] x)+  txindex0 @sh1 | Refl <- lemAppNil @sh1 = txindex+  txoneHot :: ( KnownShX sh1, KnownShX sh2, KnownSTK x+              , BoolOf (PrimalOf target) ~ BoolOf target+              , EqH (PrimalOf target) (TKScalar Int64), ConvertTensor target )+           => IShX sh1 -> target (TKX2 sh2 x) -> IxXOf target sh1+           -> target (TKX2 (sh1 ++ sh2) x)+  {-# INLINE txoneHot #-}+  txoneHot @sh1 @sh2 @x sh1 v ix+   | SNat <- ssxRank (knownShX @sh1) = case knownSTK @x of+    STKScalar ->+      gcastWith (unsafeCoerceRefl :: Take (Rank sh1) (sh1 ++ sh2) :~: sh1) $+      gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) (sh1 ++ sh2) :~: sh2) $+      txscatter @_ @'[] @_ @sh1 (shxAppend sh1 (xshape v)) v (const ix)+    _ -> case tftk knownSTK v of+      FTKX _ ftk2 ->+        -- TODO: def at out of bounds+        gcastWith (unsafeCoerceRefl+                   :: Drop (Rank (sh1 ++ sh2)) (sh1 ++ sh2) :~: '[]) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank (sh1 ++ sh2)) (sh1 ++ sh2) :~: (sh1 ++ sh2)) $+        gcastWith (unsafeCoerceRefl+                   :: Drop (Rank sh1) (sh1 ++ sh2) :~: sh2) $+        withKnownShX (knownShX @sh1 `ssxAppend` knownShX @sh2) $+        let f ix2 = tcond knownSTK+                          (foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true+                           $ zip (Foldable.toList ix) (Foldable.toList ix2))+                          (txindex0 v (ixxDrop' @(Rank sh1) ix2))+                          (tdefTarget (FTKX ZSX ftk2))+        in xbuild @(Rank (sh1 ++ sh2)) (shxAppend sh1 (xshape v)) f+  txscatter :: (KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x)+            => IShX (shp ++ shn) -> target (TKX2 (shm ++ shn) x)+            -> (IxXOf target shm -> IxXOf target shp)+            -> target (TKX2 (shp ++ shn) x)+  -- TODO: generalize this to non-Just types.+  txscatter1 :: (KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x)+             => IShX (shp ++ shn) -> target (TKX2 (Just n2 ': shn) x)+             -> (IntOf target -> IxXOf target shp)+             -> target (TKX2 (shp ++ shn) x)+  txscatter1 @n2 @_ @shp @x sh v f = txscatter @_ @'[Just n2] @_ @shp @x sh v+                                               (\(i :.% _) -> f i)+  txgather :: (KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x)+           => IShX (shm ++ shn)+           -> target (TKX2 (shp ++ shn) x)+           -> (IxXOf target shm -> IxXOf target shp)+           -> target (TKX2 (shm ++ shn) x)+  txgather1 :: (KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x)+            => SNat n2 -> target (TKX2 (shp ++ shn) x)+            -> (IntOf target -> IxXOf target shp)+            -> target (TKX2 (Just n2 ': shn) x)+  txgather1 @n2 @_ @shp k v f =+    txgather @target @'[Just n2]+             (Nested.SKnown k :$% shxDropSSX (knownShX @shp) (xshape v)) v+             (\(i :.% ZIX) -> f i)++  trfloor :: (GoodScalar r, RealFrac r, GoodScalar r2, Integral r2)+          => target (TKR n r) -> target (TKR n r2)+  trfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2)+                 => target (TKR n r1) -> target (TKR n r2)+  trcast :: (RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2)+         => target (TKR n r1) -> target (TKR n r2)+  trminIndex, trmaxIndex  -- partial+    :: forall n r r2. (GoodScalar r, GoodScalar r2)+    => target (TKR (1 + n) r) -> target (TKR n r2)+  triota :: GoodScalar r => Int -> target (TKR 1 r)  -- from 0 to n - 1++  tsfloor :: (GoodScalar r, RealFrac r, GoodScalar r2, Integral r2)+          => target (TKS sh r) -> target (TKS sh r2)+  tsfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2)+                 => target (TKS sh r1) -> target (TKS sh r2)+  tscast :: (RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2)+         => target (TKS sh r1) -> target (TKS sh r2)+  tsminIndex, tsmaxIndex  -- partial+    :: forall n sh r r2. (GoodScalar r, GoodScalar r2)+    => target (TKS (n ': sh) r) -> target (TKS (Init (n ': sh)) r2)+  tsiota :: (KnownNat n, GoodScalar r)+         => target (TKS '[n] r)  -- from 0 to n - 1++  txfloor :: (GoodScalar r, RealFrac r, GoodScalar r2, Integral r2)+          => target (TKX sh r) -> target (TKX sh r2)+  txfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2)+                 => target (TKX sh r1) -> target (TKX sh r2)+  txcast :: (RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2)+         => target (TKX sh r1) -> target (TKX sh r2)+  txminIndex, txmaxIndex  -- partial+    :: forall mn sh r r2. (GoodScalar r, GoodScalar r2)+    => target (TKX (mn ': sh) r) -> target (TKX (Init (mn ': sh)) r2)+  txiota :: (KnownNat n, GoodScalar r)+         => target (TKX '[Just n] r)  -- from 0 to n - 1++  tkfloor :: (GoodScalar r, RealFrac r, GoodScalar r2, Integral r2)+          => target (TKScalar r) -> target (TKScalar r2)+  tkfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2)+                 => target (TKScalar r1) -> target (TKScalar r2)+  tkcast :: (RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2)+         => target (TKScalar r1) -> target (TKScalar r2)++  trappend :: forall n x. KnownSTK x+           => target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)+           -> target (TKR2 (1 + n) x)+  trslice :: forall n x. KnownSTK x+          => Int -> Int -> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)+  trreverse :: forall n x. KnownSTK x+            => target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)+  trtranspose :: forall n x. KnownSTK x+              => Permutation.PermR -> target (TKR2 n x) -> target (TKR2 n x)+  trreshape :: forall n m x. KnownSTK x+            => IShR m -> target (TKR2 n x) -> target (TKR2 m x)++  tsappend :: forall m n sh x. KnownSTK x+           => target (TKS2 (m ': sh) x) -> target (TKS2 (n ': sh) x)+           -> target (TKS2 ((m + n) ': sh) x)+  tsslice :: forall i n k sh x. KnownSTK x+          => SNat i -> SNat n -> SNat k+          -> target (TKS2 (i + n + k ': sh) x) -> target (TKS2 (n ': sh) x)+  tsreverse :: forall n sh x. KnownSTK x+            => target (TKS2 (n ': sh) x) -> target (TKS2 (n ': sh) x)+  tstranspose :: ( Permutation.IsPermutation perm+                 , Rank perm <= Rank sh, KnownSTK x )+              => Permutation.Perm perm -> target (TKS2 sh x)+              -> target (TKS2 (Permutation.PermutePrefix perm sh) x)+  tsreshape :: (Product sh ~ Product sh2, KnownSTK x)+            => ShS sh2 -> target (TKS2 sh x) -> target (TKS2 sh2 x)++  txappend :: forall m n sh x. KnownSTK x+           => target (TKX2 (Just m ': sh) x) -> target (TKX2 (Just n ': sh) x)+           -> target (TKX2 (Just (m + n) ': sh) x)+  txslice :: forall i n k sh x. KnownSTK x+          => SNat i -> SNat n -> SNat k+          -> target (TKX2 (Just (i + n + k) ': sh) x)+          -> target (TKX2 (Just n ': sh) x)+  txreverse :: forall mn sh x. KnownSTK x+            => target (TKX2 (mn ': sh) x) -> target (TKX2 (mn ': sh) x)+  txtranspose :: ( Permutation.IsPermutation perm+                 , Rank perm <= Rank sh, KnownSTK x )+              => Permutation.Perm perm -> target (TKX2 sh x)+              -> target (TKX2 (Permutation.PermutePrefix perm sh) x)+  txreshape :: forall sh sh2 x. KnownSTK x+            => IShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)++  trbuild1 :: (KnownNat n, KnownSTK x)+           => Int -> (IntOf target -> target (TKR2 n x))+           -> target (TKR2 (1 + n) x)+  trmap0N :: (KnownNat n, KnownSTK x, KnownSTK x1)+          => (target (TKR2 0 x1) -> target (TKR2 0 x)) -> target (TKR2 n x1)+          -> target (TKR2 n x)+  trmap0N f v = rbuild (rshape v) (f . trindex0 v)+  trzipWith0N :: (KnownNat n, KnownSTK x, KnownSTK x1, KnownSTK x2)+              => (target (TKR2 0 x1) -> target (TKR2 0 x2) -> target (TKR2 0 x))+              -> target (TKR2 n x1) -> target (TKR2 n x2) -> target (TKR2 n x)+  trzipWith0N f u v =+    rbuild (rshape v) (\ix -> f (trindex0 u ix) (trindex0 v ix))++  tsbuild1 :: (KnownNat k, KnownShS sh, KnownSTK x)+           => (IntOf target -> target (TKS2 sh x))+           -> target (TKS2 (k ': sh) x)+  tsmap0N :: (KnownShS sh, KnownSTK x, KnownSTK x1)+          => (target (TKS2 '[] x1) -> target (TKS2 '[] x))+          -> target (TKS2 sh x1)+          -> target (TKS2 sh x)+  tsmap0N @sh f v =+    gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])+    $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)+    $ sbuild @(Rank sh) (f . tsindex0 v)+  tszipWith0N :: (KnownShS sh, KnownSTK x, KnownSTK x1, KnownSTK x2)+              => (target (TKS2 '[] x1) -> target (TKS2 '[] x2)+                  -> target (TKS2 '[] x))+              -> target (TKS2 sh x1) -> target (TKS2 sh x2)+              -> target (TKS2 sh x)+  tszipWith0N @sh f u v =+    gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])+    $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)+    $ sbuild @(Rank sh) (\ix -> f (tsindex0 u ix) (tsindex0 v ix))++  txbuild1 :: (KnownNat k, KnownShX sh, KnownSTK x)+           => (IntOf target -> target (TKX2 sh x))+           -> target (TKX2 (Just k ': sh) x)++  tbuild1 :: forall y k. ConvertTensor target+               -- y comes first, because k easy to set via SNat+          => SNat k -> SingletonTK y -> (IntOf target -> target y)+          -> target (BuildTensorKind k y)+  tbuild1 snat@SNat stk0 f =+    let replSTK :: SingletonTK z -> (IntOf target -> target z)+                -> target (BuildTensorKind k z)+        replSTK stk g = case stk of+          STKScalar -> tsbuild1 (sfromK . g)+          STKR SNat x | Dict <- lemKnownSTK x -> trbuild1 (sNatValue snat) g+          STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsbuild1 g+          STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txbuild1 g+          STKProduct @z1 @z2 stk1 stk2 ->+              let f1 i = tproject1 @_ @z1 @z2 $ g i+                  f2 i = tproject2 @_ @z1 @z2 $ g i+                    -- TODO: looks expensive, but hard to do better,+                    -- so let's hope g is full of variables+              in tpair (replSTK stk1 f1) (replSTK stk2 f2)+    in replSTK stk0 f++  -- | A strict right mapAccum.+  --+  -- The applications of 'tjvp' and 'tvjp' performed already at this point+  -- ensure that the computation of a derivative is not repeated+  -- and that its result is shared. However, most of the time+  -- the computation is unnneeded, so the AST instance uses a non-strict+  -- constructor 'HordeAd.Core.Ast.AstLambda' for it's instance of 'HFunOf'.+  --+  -- If the same argument functions are passed to many mapAccum calls, as in+  -- > let f = ... in ... (tmapAccumR ... f ...) ... (tmapAccumL ... f ...)+  -- extra care is needed to prevent double derivative computation.+  -- One needs to use tmapAccumRDer manually as in (simplified)+  -- > let f = ...; df = tjvp f; rf = tgrad f+  -- > in ... (tmapAccumRDer f df rf ...) ... (tmapAccumLDer f df rf ...)+  tmapAccumRDer+    :: forall accy by ey k.+       Proxy target+    -> SNat k  -- ^ length of the input+    -> FullShapeTK accy  -- ^ shape of the accumulator+    -> FullShapeTK by  -- ^ shape of the output+    -> FullShapeTK ey  -- ^ shape of an individual input+    -> HFunOf target (TKProduct accy ey) (TKProduct accy by)+         -- ^ the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy by))+         -- ^ the derivative of the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy ey))+         -- ^ the reverse derivative of the function to mapAccum with+    -> target accy  -- ^ the initial accumulator+    -> target (BuildTensorKind k ey)  -- ^ the inputs+    -> target (TKProduct accy (BuildTensorKind k by))+  -- | A strict left mapAccum.+  tmapAccumLDer+    :: forall accy by ey k.+       Proxy target+    -> SNat k  -- ^ length of the input+    -> FullShapeTK accy  -- ^ shape of the accumulator+    -> FullShapeTK by  -- ^ shape of the output+    -> FullShapeTK ey  -- ^ shape of an individual input+    -> HFunOf target (TKProduct accy ey) (TKProduct accy by)+         -- ^ the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy by))+         -- ^ the derivative of the function to mapAccum with+    -> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))+                                (TKProduct accy ey))+                     (ADTensorKind (TKProduct accy ey))+         -- ^ the reverse derivative of the function to mapAccum with+    -> target accy  -- ^ the initial accumulator+    -> target (BuildTensorKind k ey)  -- ^ the inputs+    -> target (TKProduct accy (BuildTensorKind k by))+  tApply :: HFunOf target x z -> target x -> target z+  tlambda :: FullShapeTK x -> HFun x z -> HFunOf target x z++  -- | Reverse derivative.+  --+  -- The followign methods (and tlambda) are exactly what is needed as arguments+  -- of tmapAccumRDer.+  tgrad+    :: FullShapeTK x  -- ^ shape of x and dx+    -> HFun x (TKScalar r)  -- ^ x |-> TKScalar r+    -> HFunOf target x (ADTensorKind x)  -- ^ x |-> dx+  tvjp+    :: FullShapeTK x  -- ^ shape of x and dx+    -> HFun x z  -- ^ x |-> z+    -> HFunOf target (TKProduct (ADTensorKind z) x) (ADTensorKind x)+         -- ^ (dz, x) |-> dx+  tjvp+    :: FullShapeTK x  -- ^ shape of x and dx+    -> HFun x z  -- ^ x |-> z+    -> HFunOf target (TKProduct (ADTensorKind x) x) (ADTensorKind z)+         -- ^ (dx, x) |-> dz++  tprimalPart :: target y -> PrimalOf target y+  tdualPart :: SingletonTK y -> target y -> DualOf target y+  tfromPrimal :: SingletonTK y -> PrimalOf target y -> target y+  tfromDual :: DualOf target y -> target y+  tScale :: (Num (target y), Num (PrimalOf target y))+         => SingletonTK y -> PrimalOf target y -> DualOf target y+         -> DualOf target y+  tScale stk s t =+    tdualPart stk $ tfromPrimal @target stk s * tfromDual t++  -- General operations that use ShareTensor if available, LetTensor otherwise+  tsum+    :: forall z k. ConvertTensor target+    => SNat k -> SingletonTK z -> target (BuildTensorKind k z)+    -> target z+  default tsum+    :: forall z k. (ShareTensor target, ConvertTensor target)+    => SNat k -> SingletonTK z -> target (BuildTensorKind k z)+    -> target z+  tsum snat@SNat stk u = case stk of+    STKScalar -> kfromS $ tssum u+    STKR SNat x | Dict <- lemKnownSTK x -> trsum u+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u+    STKProduct stk1 stk2 ->+      let (u1, u2) = tunpair u+      in tpair (tsum snat stk1 u1)+               (tsum snat stk2 u2)+  treplicate+    :: forall z k. ConvertTensor target+    => SNat k -> SingletonTK z -> target z+    -> target (BuildTensorKind k z)+  default treplicate+    :: forall z k. (ShareTensor target, ConvertTensor target)+    => SNat k -> SingletonTK z -> target z+    -> target (BuildTensorKind k z)+  treplicate snat@SNat stk u = case stk of+    STKScalar -> tsreplicate snat ZSS $ sfromK u+    STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (sNatValue snat) u+    STKS sh x | Dict <- lemKnownSTK x -> tsreplicate snat sh u+    STKX sh x | Dict <- lemKnownSTK x -> txreplicate snat sh u+    STKProduct stk1 stk2 ->+      let (u1, u2) = tunpair u+      in tpair (treplicate snat stk1 u1)+               (treplicate snat stk2 u2)+  tindexBuild+    :: forall z k. ConvertTensor target+    => SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> IntOf target+    -> target z+  default tindexBuild+    :: forall z k. (ShareTensor target, ConvertTensor target)+    => SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> IntOf target+    -> target z+  tindexBuild snat@SNat stk u i = case stk of+    STKScalar -> kfromS $ tsindex u (i :.$ ZIS)+    STKR SNat x | Dict <- lemKnownSTK x -> trindex u (i :.: ZIR)+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsindex u (i :.$ ZIS)+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txindex u (i :.% ZIX)+    STKProduct stk1 stk2 ->+      let (u1, u2) = tunpair u+      in tpair (tindexBuild snat stk1 u1 i)+               (tindexBuild snat stk2 u2 i)++  -- Unwinding methods, needed mostly to split off the Unwind module.+  -- | Construct tensors with the given constant in each cell.+  treplTarget :: (forall r. GoodScalar r => r) -> FullShapeTK y -> target y+  -- | Construct tensors with @def@ in each cell.+  tdefTarget :: FullShapeTK y -> target y+  -- | Add pointwise all corresponding tensors within nested product, if any.+  --+  -- Requires duplicable arguments or a ShareTensor instance.+  taddTarget :: SingletonTK y -> target y -> target y -> target y+  -- | Multiply pointwise all corresponding tensors within nested products,+  -- if any.+  --+  -- Requires duplicable arguments or a ShareTensor instance.+  tmultTarget :: SingletonTK y -> target y -> target y -> target y+  -- | Sum all dimensions of each component and then sum it all. Ignore all+  -- tensors with non-differentiable elements.+  --+  -- Requires duplicable arguments or a ShareTensor instance.+  tsum0Target :: FullShapeTK y -> target y+              -> target (TKScalar Double)+  -- | Dot product each component and then sum it all. Ignore all+  -- tensors with non-differentiable elements.+  --+  -- Requires duplicable arguments or a ShareTensor instance.+  tdot0Target :: FullShapeTK y -> target y -> target y+              -> target (TKScalar Double)++  -- TODO: express without ConvertTensor or move there+  xmcast :: (KnownSTK x, KnownShX sh, Rank sh ~ Rank sh2, ConvertTensor target)+         => StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)+  xmcast sh2 a = case tftk knownSTK a of+    FTKX sh' _ ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        withKnownShX sh2 $+        withKnownShS sh $+        xfromS $ sfromX @_ @sh a++-- These are user-accessible, so the constraint is `ADReady`, which means+-- lets, but no shares.+type role HFun nominal nominal+newtype HFun (x :: TK) (z :: TK) =+  HFun {unHFun :: forall f. ADReady f+               => f x -> f z}++instance Show (HFun x y) where+  show _ = "<lambda>"+++-- * The mega-constraint++type ADReady target =+  ( ADReadyNoLet target+  , LetTensor target+-- The following can't be added, because we have instances like ADVal (AstRaw),+-- so AstRaw would need to have a LetTensor instance:+--  , LetTensor (PrimalOf target)+  )++type ADReadyNoLet target =+  ( ADReadyEqsClasses target+  , ADReadyEqsClasses (ShareOf target)+  , ShareTensor (ShareOf target)+  , ShareTensor (PrimalOf (ShareOf target))+  , ShareOf (ShareOf target) ~ ShareOf target+  )++type ADReadyEqsClasses target =+  ( ADReadyEqs target+  , ADReadyClasses target+  , ADReadyClasses (PrimalOf target)+  )++type ADReadyEqs target =+  ( BoolOf (PrimalOf target) ~ BoolOf target+  )++type ADReadyClasses target =+  ( BaseTensor target+  , ConvertTensor target+  , Boolean (BoolOf target)+  , AllTargetShow target+  , CommonTargetEqOrd target+  )++-- This is illegal:+-- type AllTargetShow target = forall y. KnownSTK y => Show (target y)++type AllTargetShow :: Target -> Constraint+class (forall y. KnownSTK y => Show (target y))+      => AllTargetShow target where+instance+      (forall y. KnownSTK y => Show (target y))+      => AllTargetShow target where++type CommonTargetEqOrd :: Target -> Constraint+class ( forall r. GoodScalar r => EqH target (TKScalar r)+      , forall r. GoodScalar r => OrdH target (TKScalar r)+      , forall r n. GoodScalar r => EqH target (TKR n r)+      , forall r n. GoodScalar r => OrdH target (TKR n r)+      , forall r sh. GoodScalar r => EqH target (TKS sh r)+      , forall r sh. GoodScalar r => OrdH target (TKS sh r)+      , forall r sh. GoodScalar r => EqH target (TKX sh r)+      , forall r sh. GoodScalar r => OrdH target (TKX sh r) )+      => CommonTargetEqOrd target where+instance+      ( forall r. GoodScalar r => EqH target (TKScalar r)+      , forall r. GoodScalar r => OrdH target (TKScalar r)+      , forall r n. GoodScalar r => EqH target (TKR n r)+      , forall r n. GoodScalar r => OrdH target (TKR n r)+      , forall r sh. GoodScalar r => EqH target (TKS sh r)+      , forall r sh. GoodScalar r => OrdH target (TKS sh r)+      , forall r sh. GoodScalar r => EqH target (TKX sh r)+      , forall r sh. GoodScalar r => OrdH target (TKX sh r) )+      => CommonTargetEqOrd target where
+ src/HordeAd/Core/OpsADVal.hs view
@@ -0,0 +1,643 @@+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Tensor class instances for dual numbers. All definitions+-- are generic over whether the dual numbers are built from concrete arrays+-- of floats or from AST terms or anything else (e.g., nested 'ADVal').+module HordeAd.Core.OpsADVal+  ( crevOnADInputs, crevOnParams, cfwdOnParams+  ) where++import Prelude hiding (foldl')++import Data.Maybe (fromMaybe)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import GHC.TypeLits (sameNat)++import Data.Array.Nested (Replicate, type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.CarriersADVal+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Delta+import HordeAd.Core.DeltaEval+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind++-- * Non-symbolic (or at least non-sharing) reverse and forward derivative computation++-- The user-written function f can do anything, so the inputs+-- argument has to be duplicable.+crevOnADInputs+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => Maybe (target (ADTensorKind z))+  -> (ADVal target x -> ADVal target z)+  -> FullShapeTK x -> ADVal target x+  -> (target (ADTensorKind x), target z)+-- Break the inline chain to prevent false positives in inspection testing.+-- {-# INLINE crevOnADInputs #-}+crevOnADInputs mdt f xftk inputs =+  let -- Evaluate completely after terms constructed, to free memory+      -- before evaluation allocates new memory and new FFI is started.+      !(D v delta) = f inputs in+  let zftk = ftkDelta delta+      dt = fromMaybe (treplTarget 1 $ adFTK zftk) mdt+      !gradient = gradientFromDelta xftk zftk dt delta+  in (gradient, v)++crevOnParams+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => Maybe (target (ADTensorKind z))+  -> (ADVal target x -> ADVal target z)+  -> FullShapeTK x -> target x+  -> (target (ADTensorKind x), target z)+{-# INLINE crevOnParams #-}+crevOnParams edt f xftk parameters =+  let deltaInputs = generateDeltaInputs xftk+      inputs = dDnotShared parameters deltaInputs+  in crevOnADInputs edt f xftk inputs++cfwdOnADInputs+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => FullShapeTK x -> ADVal target x+  -> (ADVal target x -> ADVal target z)+  -> target (ADTensorKind x)+  -> (target (ADTensorKind z), target z)+-- Break the inline chain to prevent false positives in inspection testing.+-- {-# INLINE cfwdOnADInputs #-}+cfwdOnADInputs xftk inputs f ds =+  let !(D v delta) = f inputs in+  let !derivative = derivativeFromDelta @x delta (adFTK xftk) ds+  in (derivative, v)++cfwdOnParams+  :: forall x z target. (ADReadyNoLet target, ShareTensor target)+  => FullShapeTK x -> target x+  -> (ADVal target x -> ADVal target z)+  -> target (ADTensorKind x)+  -> (target (ADTensorKind z), target z)+{-# INLINE cfwdOnParams #-}+cfwdOnParams xftk parameters f ds =+  let deltaInputs = generateDeltaInputs xftk+      inputs = dDnotShared parameters deltaInputs+  in cfwdOnADInputs xftk inputs f ds+++-- * Instances++fromPrimalFTK :: FullShapeTK z -> f z -> ADVal f z+fromPrimalFTK ftk a = dDnotShared a (DeltaZero ftk)++instance ( ADReadyNoLet target, ShareTensor target+         , ShareTensor (PrimalOf target) )+         => LetTensor (ADVal target) where+  ttlet (D u u') f =+    let !var2 = tshare u+    in f (dDnotShared var2 u')  -- u' was already shared+  ttletPrimal u f =+    let !var2 = tshare u+    in f var2+  toShare = id+  tunshare = id+  tD _stk = dD++instance (ADReadyNoLet target, ShareTensor target)+         => ShareTensor (ADVal target) where+  tshare (D u u') = dDnotShared (tshare u) u'  -- u' was already shared+  tunpair (D u u') = let (u1, u2) = tunpair u+                         (d1, d2) = unDeltaPair u'+                     in (dDnotShared u1 d1, dDnotShared u2 d2)++-- Note that this instance doesn't do vectorization. To enable it,+-- use the Ast instance and only then interpret in ADVal.+-- In any case, only the Ast instantiation of this instance+-- is used in the best pipeline, in particular, to satisfy the constraints+-- needed for the interpretation of Ast in ADVal.+-- The ADVal Concrete instantiation is used in other pipelines and tests.+instance ( ADReadyNoLet target, ShareTensor target+         , ShareTensor (PrimalOf target) )+         => BaseTensor (ADVal target) where+  -- Ranked ops+  rshape (D u _) = rshape u+  trsum (D u u') = withSNat (rwidth u) $ \snat ->+    dD (trsum u) (DeltaSum snat knownSTK u')+  trsum0 (D u u') = dD (trsum0 u) (DeltaSum0R u')+  trdot0 (D ue u') (D ve v') =+    -- The bangs below are neccessary for GHC 9.2.7 test results to match 9.4.+    let !u = tshare ue in+    let !v = tshare ve+    in dD (trdot0 u v) (dAdd (DeltaDot0R v u') (DeltaDot0R u v'))+  -- These two are manually vectorized to avoid delta blowup when run+  -- via primitive pipelines.+  trmatvecmul m v = trsum (rtr (trreplicate (rwidth m) v * m))+  trmatmul2 m1 m2 = case rshape m2 of+    _ :$: width2 :$: ZSR ->+      trsum (trtranspose [2,1,0] (trreplicate width2 m1)+             * trtranspose [1,0] (trreplicate (rwidth m1) m2))+  trreplicate k (D u u') = withSNat k $ \snat ->+    dD (trreplicate k u) (DeltaReplicate snat knownSTK u')+  trindex (D u u') i =+    let !ix = tshare . tprimalPart <$> i+    in dD (trindex u ix) (DeltaIndexR SNat u' ix)+  trscatter sh (D u u') f =+    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)+    in dD (trscatter sh u g) (DeltaScatterR SNat SNat SNat sh u' g)+  trgather sh (D u u') f =+    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)+    in dD (trgather sh u g) (DeltaGatherR SNat SNat SNat sh u' g)+      -- Note how f is not interpreted as a function on dual numbers+      -- but just on integers and so no cotangents for results of application+      -- of f have to be computed and stored in contangent maps later on.+      -- Note also how g is duplicated and this leads to loss of sharing+      -- of indexes in AST instances.+  trconcrete a =+    let v = trconcrete a+    in fromPrimalFTK (FTKR (Nested.rshape a) FTKScalar) v+  trfloor (D u _) =+    let v = trfloor u+    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v+  trfromIntegral (D u _) =+    let v = trfromIntegral u+    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v+  trcast (D u u') = dD (trcast u) (DeltaCastR u')+  trminIndex (D u _) =+    let v = trminIndex u+    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v+  trmaxIndex (D u _) =+    let v = trmaxIndex u+    in fromPrimalFTK (FTKR (rshape v) FTKScalar) v+  triota n = fromPrimalFTK (FTKR (n :$: ZSR) FTKScalar) $ triota n+  trappend (D u u') (D v v') = dD (trappend u v) (DeltaAppendR u' v')+  trslice i n (D u u') = dD (trslice i n u) (DeltaSliceR i n u')+  trreverse (D u u') = dD (trreverse u) (DeltaReverseR u')+  trtranspose perm (D u u') = dD (trtranspose perm u) (DeltaTransposeR perm u')+  trreshape sh (D u u') = dD (trreshape sh u) (DeltaReshapeR sh u')+  trbuild1 @n @x k f =+    let l = [0 .. fromIntegral k - 1]+    in if null l+       then case sameNat (Proxy @n) (Proxy @0) of+         Just Refl | Dict <- eltDictRep (knownSTK @x) ->+           let arr = Nested.remptyArray+           in tconcrete (tftkG knownSTK arr) (Concrete arr)+         Nothing -> error "rbuild1: shape ambiguity"+       else trfromVector $ V.fromList $ map (f . fromInteger) l+              -- hope this fuses++  -- Shaped ops+  sshape (D u _) = sshape u+  tssum (D u u') = dD (tssum u) (DeltaSum SNat knownSTK u')+  tssum0 (D u u') = dD (tssum0 u) (DeltaSum0S u')+  tsdot0 (D ue u') (D ve v') =+    -- The bangs below are neccessary for GHC 9.2.7 test results to match 9.4.+    let !u = tshare ue in+    let !v = tshare ve+    in dD (tsdot0 u v) (dAdd (DeltaDot0S v u') (DeltaDot0S u v'))+  -- These two are manually vectorized to avoid delta blowup when run+  -- via primitive pipelines.+  tsmatvecmul m v = tssum (str (tsreplicate SNat knownShS v * m))+  tsmatmul2 m1 m2 =+    tssum (tstranspose (Permutation.makePerm @'[2, 1, 0])+                       (tsreplicate SNat knownShS m1)+           * tstranspose (Permutation.makePerm @'[1, 0])+                         (tsreplicate SNat knownShS m2))+  tsindex (D u u') i =+    let !ix = tshare . tprimalPart <$> i+    in dD (tsindex u ix) (DeltaIndexS knownShS u' ix)+  tsscatter @shm @shn @shp (D u u') f =+    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)+    in dD (tsscatter @_ @shm @shn @shp u g)+          (DeltaScatterS @shm @shn @shp knownShS knownShS knownShS u' g)+  tsgather @shm @shn @shp (D u u') f =+    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)+    in dD (tsgather @_ @shm @shn @shp u g)+          (DeltaGatherS @shm @shn @shp knownShS knownShS knownShS u' g)+  tsconcrete a =+    let v = tsconcrete a+    in fromPrimalFTK (FTKS (Nested.sshape a) FTKScalar) v+  tsfloor (D u _) =+    let v = tsfloor u+    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v+  tsfromIntegral (D u _) =+    let v = tsfromIntegral u+    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v+  tscast (D u u') = dD (tscast u) (DeltaCastS u')+  tsminIndex (D u _) =+    let v = tsminIndex u+    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v+  tsmaxIndex (D u _) =+    let v = tsmaxIndex u+    in fromPrimalFTK (FTKS (sshape v) FTKScalar) v+  tsiota = fromPrimalFTK (FTKS (SNat :$$ ZSS) FTKScalar) tsiota+  tsappend (D u u') (D v v') = dD (tsappend u v) (DeltaAppendS u' v')+  tsslice i n k (D u u') = dD (tsslice i n k u) (DeltaSliceS i n k u')+  tsreverse (D u u') = dD (tsreverse u) (DeltaReverseS u')+  tsbuild1 @k @sh @r f | Dict <- eltDictRep (knownSTK @r) =+    let l = [0 .. valueOf @k - 1]+    in if null l+       then let arr = Nested.semptyArray @(RepConcrete r) (knownShS @sh)+            in gcastWith (unsafeCoerceRefl :: k :~: 0) $+               tconcrete (tftkG knownSTK arr) (Concrete arr)+       else tsfromVector $ V.fromList $ map (f . fromInteger) l+              -- hope this fuses++  -- Mixed ops+  xshape (D u _) = xshape u+  txsum (D u u') = dD (txsum u) (DeltaSum SNat knownSTK u')+  txsum0 (D u u') = dD (txsum0 u) (DeltaSum0X u')+  txdot0 (D ue u') (D ve v') =+    -- The bangs below are neccessary for GHC 9.2.7 test results to match 9.4.+    let !u = tshare ue in+    let !v = tshare ve+    in dD (txdot0 u v) (dAdd (DeltaDot0X v u') (DeltaDot0X u v'))+  -- These two are manually vectorized to avoid delta blowup when run+  -- via primitive pipelines.+  txmatvecmul mm mn m v =+    withKnownShX (ssxFromShX $ mn :$% ZSX) $+    withKnownShX (ssxFromShX $ mm :$% mn :$% ZSX) $+    withSNat (fromSMayNat' mm) $ \(SNat @m) ->+    withSNat (fromSMayNat' mn) $ \(SNat @n) ->+      xmcast (ssxFromShX (mm :$% ZSX))+      $ txsum (xtr (txreplicate (SNat @m) knownShX+                      (xmcast (ssxFromShX (Nested.SKnown (SNat @n)+                                             :$% ZSX)) v)+                    * xmcast (ssxFromShX (Nested.SKnown (SNat @m)+                                            :$% Nested.SKnown (SNat @n)+                                            :$% ZSX)) m))+  txmatmul2 m1 m2 =+    txsum (txtranspose (Permutation.makePerm @'[2, 1, 0])+                       (txreplicate SNat knownShX m1)+           * txtranspose (Permutation.makePerm @'[1, 0])+                         (txreplicate SNat knownShX m2))+  txreplicate snat sh (D u u') =+    dD (txreplicate snat sh u) (DeltaReplicate snat (STKX sh knownSTK) u')+  txindex (D u u') i =+    let !ix = tshare . tprimalPart <$> i+    in dD (txindex u ix) (DeltaIndexX knownShX u' ix)+  txscatter @shm @shn @shp sh (D u u') f =+    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)+    in dD (txscatter @_ @shm @shn @shp sh u g)+          (DeltaScatterX @shm @shn @shp knownShX knownShX knownShX sh u' g)+  txgather @shm @shn @shp sh (D u u') f =+    let g x = tprimalPart <$> f (tfromPrimal STKScalar <$> x)+    in dD (txgather @_ @shm @shn @shp sh u g)+          (DeltaGatherX @shm @shn @shp knownShX knownShX knownShX sh u' g)+  txconcrete a =+    let v = txconcrete a+    in fromPrimalFTK (FTKX (Nested.mshape a) FTKScalar) v+  txfloor (D u _) =+    let v = txfloor u+    in fromPrimalFTK (FTKX (xshape v) FTKScalar) v+  txfromIntegral (D u _) =+    let v = txfromIntegral u+    in fromPrimalFTK (FTKX (xshape v) FTKScalar) v+  txcast (D u u') = dD (txcast u) (DeltaCastX u')+  txminIndex (D u _) =+    let v = txminIndex u+    in fromPrimalFTK (FTKX (xshape v) FTKScalar) v+  txmaxIndex (D u _) =+    let v = txmaxIndex u+    in fromPrimalFTK (FTKX (xshape v) FTKScalar) v+  txiota = fromPrimalFTK (FTKX (Nested.SKnown SNat :$% ZSX) FTKScalar) txiota+  txappend (D u u') (D v v') = dD (txappend u v) (DeltaAppendX u' v')+  txslice i n k (D u u') = dD (txslice i n k u) (DeltaSliceX i n k u')+  txreverse (D u u') = dD (txreverse u) (DeltaReverseX u')+  txtranspose perm (D u u') =+    dD (txtranspose perm u) (DeltaTransposeX @_ @_ @_ @target perm u')+  txreshape sh (D u u') = dD (txreshape sh u) (DeltaReshapeX sh u')+  txbuild1 @k @sh @r f =+    let l = [0 .. valueOf @k - 1]+    in if null l+       then case testEquality (knownShX @sh) ZKX of+         Just Refl | Dict <- eltDictRep (knownSTK @r) ->+           let arr = Nested.memptyArray @(RepConcrete r) ZSX+           in gcastWith (unsafeCoerceRefl :: k :~: 0) $+              tconcrete (tftkG knownSTK arr) (Concrete arr)+         Nothing -> error "xbuild1: shape ambiguity"+       else txfromVector $ V.fromList $ map (f . fromInteger) l+              -- hope this fuses++  -- Scalar ops+  tkconcrete a =+    let v = tkconcrete a+    in fromPrimalFTK FTKScalar v+  tkfloor (D u _) =+    let v = tkfloor u+    in fromPrimalFTK FTKScalar v+  tkfromIntegral (D u _) =+    let v = tkfromIntegral u+    in fromPrimalFTK FTKScalar v+  tkcast (D u u') = dD (tkcast u) (DeltaCastK u')++  -- General operations that don't require LetTensor nor ShareTensor+  tftk _stk (D _ u') = ftkDelta u'+  tconcrete ftk t | Dict <- lemKnownSTK (ftkToSTK ftk) =+    fromPrimalFTK ftk $ tconcrete ftk t+  tpair (D u u') (D v v') = dDnotShared (tpair u v) (DeltaPair u' v')+  tproject1 (D u u') = dDnotShared (tproject1 u) (fst $ unDeltaPairUnshared u')+  tproject2 (D u u') = dDnotShared (tproject2 u) (snd $ unDeltaPairUnshared u')+  tsreplicate snat sh (D u u') =+    dD (tsreplicate snat sh u) (DeltaReplicate snat (STKS sh knownSTK) u')+  tstranspose perm (D u u') =+    dD (tstranspose perm u) (DeltaTransposeS @_ @_ @_ @target perm u')+  tsreshape sh (D u u') = dD (tsreshape sh u) (DeltaReshapeS sh u')+  tmapAccumRDer @accy @by @ey _ !k accftk bftk eftk f df rf acc0D esD+   | Dict <- lemKnownSTKOfBuild k (ftkToSTK accftk)+   , Dict <- lemKnownSTKOfBuild k (ftkToSTK eftk) =+    let !(D acc0 acc0') = acc0D in+    let !(D esNotShared es') = esD in+    let !es = tshare esNotShared+        codomainShs = FTKProduct accftk bftk+        g :: forall f. ADReady f+          => f (TKProduct accy ey)+          -> f (TKProduct accy (TKProduct accy by))+        g !acc_e =+          ttlet acc_e $ \ !acc_e1 ->+          ttlet (unHFun f acc_e) $ \ !accRes_bRes ->+            tpair (tproject1 accRes_bRes)+                  (tpair (tproject1 acc_e1) (tproject2 accRes_bRes))+        dg :: forall f. ADReady f+           => f (TKProduct (ADTensorKind (TKProduct accy ey))+                           (TKProduct accy ey))+           -> f (ADTensorKind (TKProduct accy (TKProduct accy by)))+        dg !dacc_de_acc_e =+          ttlet dacc_de_acc_e $ \ !dacc_de_acc_e1 ->+            let (!dacc_de, !_acc_e) =+                  (tproject1 dacc_de_acc_e1, tproject2 dacc_de_acc_e1)+                !dacc1 = tproject1 dacc_de+            in ttlet (unHFun df dacc_de_acc_e) $ \ !accRes_bRes ->+                 tpair (tproject1 accRes_bRes)+                       (tpair dacc1 (tproject2 accRes_bRes))+        rg :: forall f. ADReady f+           => f (TKProduct (ADTensorKind (TKProduct accy+                                         (TKProduct accy by)))+                           (TKProduct accy ey))+           -> f (ADTensorKind (TKProduct accy ey))+        rg !args =+          ttlet args $ \ args1 ->+            let (!dx_db, !acc_e) = (tproject1 args1, tproject2 args1)+            in ttlet dx_db $ \ !dx_db1 ->+              let (!dx, !db) = (tproject1 dx_db1, tproject2 dx_db1)+              in ttlet db $ \ !db1 ->+                let dx_dbRes = tpair dx (tproject2 db1)+                in ttlet (unHFun rf (tpair dx_dbRes acc_e))+                   $ \ !daccRes_deRes ->+                  let added = taddTarget (adSTK $ ftkToSTK accftk)+                                         (tproject1 daccRes_deRes)+                                         (tproject1 db1)+                  in tpair added (tproject2 daccRes_deRes)+        p = tmapAccumRDer (Proxy @target)+                          k accftk codomainShs eftk+                          (tlambda @target (FTKProduct accftk eftk)+                           $ HFun g)+                          (tlambda @target+                             (FTKProduct (adFTK (FTKProduct accftk eftk))+                                         (FTKProduct accftk eftk))+                           $ HFun dg)+                          (tlambda @target+                             (FTKProduct (adFTK (FTKProduct accftk codomainShs))+                                         (FTKProduct accftk eftk))+                           $ HFun rg)+                          acc0 es+        (accFin, qbs) = tunpair p+        (q, bs) = tunpair qbs+        dual = DeltaMapAccumR k bftk eftk q es df rf acc0' es'+    in dD (tpair accFin bs) dual+  tmapAccumLDer @accy @by @ey _ !k accftk bftk eftk f df rf acc0D esD+   | Dict <- lemKnownSTKOfBuild k (ftkToSTK accftk)+   , Dict <- lemKnownSTKOfBuild k (ftkToSTK eftk) =+    let !(D acc0 acc0') = acc0D in+    let !(D esNotShared es') = esD in+    let !es = tshare esNotShared+        codomainShs = FTKProduct accftk bftk+        g :: forall f. ADReady f+          => f (TKProduct accy ey)+          -> f (TKProduct accy (TKProduct accy by))+        g !acc_e =+          ttlet acc_e $ \ !acc_e1 ->+          ttlet (unHFun f acc_e) $ \ !accRes_bRes ->+            tpair (tproject1 accRes_bRes)+                  (tpair (tproject1 acc_e1) (tproject2 accRes_bRes))+        dg :: forall f. ADReady f+           => f (TKProduct (ADTensorKind (TKProduct accy ey))+                           (TKProduct accy ey))+           -> f (ADTensorKind (TKProduct accy (TKProduct accy by)))+        dg !dacc_de_acc_e =+          ttlet dacc_de_acc_e $ \ !dacc_de_acc_e1 ->+            let (!dacc_de, !_acc_e) =+                  (tproject1 dacc_de_acc_e1, tproject2 dacc_de_acc_e1)+                !dacc1 = tproject1 dacc_de+            in ttlet (unHFun df dacc_de_acc_e) $ \ !accRes_bRes ->+                 tpair (tproject1 accRes_bRes)+                       (tpair dacc1 (tproject2 accRes_bRes))+        rg :: forall f. ADReady f+           => f (TKProduct (ADTensorKind (TKProduct accy+                                         (TKProduct accy by)))+                           (TKProduct accy ey))+           -> f (ADTensorKind (TKProduct accy ey))+        rg !args =+          ttlet args $ \ args1 ->+            let (!dx_db, !acc_e) = (tproject1 args1, tproject2 args1)+            in ttlet dx_db $ \ !dx_db1 ->+              let (!dx, !db) = (tproject1 dx_db1, tproject2 dx_db1)+              in ttlet db $ \ !db1 ->+                let dx_dbRes = tpair dx (tproject2 db1)+                in ttlet (unHFun rf (tpair dx_dbRes acc_e))+                   $ \ !daccRes_deRes ->+                  let added = taddTarget (adSTK $ ftkToSTK accftk)+                                         (tproject1 daccRes_deRes)+                                         (tproject1 db1)+                  in tpair added (tproject2 daccRes_deRes)+        p = tmapAccumLDer (Proxy @target)+                          k accftk codomainShs eftk+                          (tlambda @target (FTKProduct accftk eftk)+                           $ HFun g)+                          (tlambda @target+                             (FTKProduct (adFTK (FTKProduct accftk eftk))+                                         (FTKProduct accftk eftk))+                           $ HFun dg)+                          (tlambda @target+                             (FTKProduct (adFTK (FTKProduct accftk codomainShs))+                                         (FTKProduct accftk eftk))+                           $ HFun rg)+                          acc0 es+        (accFin, qbs) = tunpair p+        (q, bs) = tunpair qbs+        dual = DeltaMapAccumL k bftk eftk q es df rf acc0' es'+    in dD (tpair accFin bs) dual+  tApply (HFun f) = f+  tlambda _ = id+  -- Bangs are for the proper order of sharing stamps.+  tcond !stk !b !u !v =+    let uv = tfromVector (SNat @2) stk (V.fromList [u, v])+    in tindexBuild (SNat @2) stk uv (tcond knownSTK b 0 1)+  tprimalPart (D u _) = u+  tdualPart _stk (D _ u') = u'+  tfromPrimal stk t = fromPrimalFTK (tftk stk t) t+  tfromDual t = dDnotShared (tdefTarget (ftkDelta t)) t+  tScale _stk = dScale+  tgrad @x xftk h =+    let rf :: forall f. ADReady f+           => f x+           -> f (ADTensorKind x)+        -- This computes the derivative of g again for each new a.+        rf !a = ttlet a $ \ !aShared ->+          tunshare $ fst $ crevOnParams+                             Nothing+                             (unHFun h @(ADVal (ShareOf f)))+                             xftk+                             (toShare aShared)+    in HFun rf+  tvjp @x @z xftk h =+    let rf :: forall f. ADReady f+           => f (TKProduct (ADTensorKind z) x)+           -> f (ADTensorKind x)+        -- This computes the derivative of g again for each new db and a.+        rf !db_a = ttlet db_a $ \ !db_aShared ->+          tunshare $ fst $ crevOnParams+                             (Just $ toShare $ tproject1 db_aShared)+                             (unHFun h @(ADVal (ShareOf f)))+                             xftk+                             (toShare $ tproject2 db_aShared)+    in HFun rf+  tjvp @x @z xftk h =+    let df :: forall f. ADReady f+           => f (TKProduct (ADTensorKind x) x)+           -> f (ADTensorKind z)+        -- This computes the derivative of g again for each new da and a.+        df !da_a = ttlet da_a $ \ !da_aShared ->+          tunshare $ fst $ cfwdOnParams+                             xftk+                             (toShare $ tproject2 da_aShared)+                             (unHFun h @(ADVal (ShareOf f)))+                             (toShare $ tproject1 da_aShared)+    in HFun df++  tfromVector snat stk lu =+    dD (tfromVector snat stk $ V.map (\(D u _) -> u) lu)+       (DeltaFromVector snat stk $ V.map (\(D _ u') -> u') lu)++  treplTarget r ftk = dDnotShared (treplTarget r ftk) (DeltaZero ftk)+  tdefTarget ftk = dDnotShared (tdefTarget ftk) (DeltaZero ftk)+  taddTarget = addTarget+  tmultTarget = multTarget+  tsum0Target = sum0Target+  tdot0Target = dot0Target++instance ( ADReadyNoLet target, ShareTensor target+         , ShareTensor (PrimalOf target) )+         => ConvertTensor (ADVal target) where+  tconvert c astk (D u u') =+    dDnotShared (tconvert c astk u)+                (DeltaConvert c u')++  rfromX a@(D _ u') = case ftkDelta u' of+    FTKX sh' _ ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        withKnownShS sh $+        rfromS $ sfromX @_ @sh a+  xfromR a@(D _ u') = case ftkDelta u' of+    FTKR shr _ ->+      withShsFromShR shr $ \(sh :: ShS sh) ->+        withKnownShS sh $+        xfromS @_ @sh $ sfromR a++  sfromR (D u u') = dDnotShared (sfromR u) (dSFromR knownShS u')+  sfromX (D u u') = dDnotShared (sfromX u) (dSFromX knownShS u')+  xfromS (D u u') = dDnotShared (xfromS u) (dXFromS knownShX u')++  rzip @_ @_ @n (D u u')+   | Refl <- lemRankReplicate (Proxy @n) = case ftkDelta u' of+    ftk@(FTKProduct (FTKR _sh y) (FTKR _ z)) ->+      let c = ConvCmp+                (ConvXR (ftkToSTK (FTKProduct y z)))+                (ConvCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvRX ConvRX))+      in dD (tconvert c (ftkToSTK ftk) u)+            (DeltaConvert c u')+  runzip @_ @_ @n (D u u')+   | Refl <- lemRankReplicate (Proxy @n) = case ftkDelta u' of+    ftk@(FTKR _sh (FTKProduct y z)) ->+      let c = ConvCmp+                (ConvT2 (ConvXR (ftkToSTK y)) (ConvXR (ftkToSTK z)))+                (ConvCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvRX)+      in dD (tconvert c (ftkToSTK ftk) u)+            (DeltaConvert c u')+  szip (D u u') = case ftkDelta u' of+    ftk@(FTKProduct (FTKS _sh y) (FTKS _ z)) ->+      let c = ConvCmp+                ConvXS+                (ConvCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvSX ConvSX))+      in dD (tconvert c (ftkToSTK ftk) u)+            (DeltaConvert c u')+  sunzip (D u u') = case ftkDelta u' of+    ftk@(FTKS _sh (FTKProduct y z)) ->+      let c = ConvCmp+                (ConvT2 ConvXS ConvXS)+                (ConvCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvSX)+      in dD (tconvert c (ftkToSTK ftk) u)+            (DeltaConvert c u')+  xzip (D u u') = case ftkDelta u' of+    ftk@(FTKProduct (FTKX _sh y) (FTKX _ z)) ->+      let c = ConvZip (ftkToSTK y) (ftkToSTK z)+      in dD (tconvert c (ftkToSTK ftk) u)+            (DeltaConvert c u')+  xunzip (D u u') = case ftkDelta u' of+    ftk@(FTKX _sh (FTKProduct y z)) ->+      let c = ConvUnzip (ftkToSTK y) (ftkToSTK z)+      in dD (tconvert c (ftkToSTK ftk) u)+            (DeltaConvert c u')++  xnestR @sh1 @m @x sh1 (D u u')+    | Refl <- lemRankReplicate (Proxy @m) =+      let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)+                            (TKX2 sh1 (TKR2 m x))+          c = ConvCmp+                (ConvXX (ConvXR (knownSTK @x)))+                (ConvNest @_ @_ @(Replicate m Nothing)+                          (STKX sh1 (knownSTK @x)))+      in dD (tconvert c (ftkToSTK $ ftkDelta u') u)+            (DeltaConvert c u')+  xnestS @_ @_ @x sh1 (D u u') =+    let c = ConvCmp (ConvXX ConvXS)+                    (ConvNest (STKX sh1 (knownSTK @x)))+    in dD (tconvert c (ftkToSTK $ ftkDelta u') u)+          (DeltaConvert c u')+  xnest @_ @_ @x sh1 (D u u') =+    let c = ConvNest (STKX sh1 (knownSTK @x))+    in dD (tconvert c (ftkToSTK $ ftkDelta u') u)+          (DeltaConvert c u')+  xunNestR (D u u') =+    let c = ConvCmp ConvUnnest+                    (ConvXX ConvRX)+    in dD (tconvert c (ftkToSTK $ ftkDelta u') u)+          (DeltaConvert c u')+  xunNestS (D u u') =+    let c = ConvCmp ConvUnnest+                   (ConvXX ConvSX)+    in dD (tconvert c (ftkToSTK $ ftkDelta u') u)+          (DeltaConvert c u')+  xunNest (D u u') =+    let c = ConvUnnest+    in dD (tconvert c (ftkToSTK $ ftkDelta u') u)+          (DeltaConvert c u')++  tpairConv = tpair+  tunpairConv = tunpair
+ src/HordeAd/Core/OpsAst.hs view
@@ -0,0 +1,1632 @@+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Tensor class instances for AST terms. Most of these instances+-- vectorize any term with the build constructor in the root.+-- The AST term instances can be used as building blocks for ADVal(AST)+-- instances defined in "HordeAd.Core.OpsADVal" but may also be used standalone.+module HordeAd.Core.OpsAst+  ( IncomingCotangentHandling(..)+  , forwardPassByInterpretation+  , revArtifactFromForwardPass, revProduceArtifact+  , fwdArtifactFromForwardPass, fwdProduceArtifact+  ) where++import Prelude++import Data.Maybe (fromMaybe)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import GHC.Exts (inline)+import GHC.TypeLits (OrderingI (..), cmpNat, type (+), type (-), type (<=?))+import System.IO.Unsafe (unsafePerformIO)+import Unsafe.Coerce (unsafeCoerce)++import Data.Array.Nested (Replicate, type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert+  ( ixrFromIxS+  , ixsFromIxR+  , ixsFromIxR'+  , ixsFromIxX'+  , ixxFromIxS+  , withShsFromShR+  , withShsFromShX+  )+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (Init, snatPlus, unsafeCoerceRefl)++import HordeAd.AstEngine+import HordeAd.Core.Ast+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstInline+import HordeAd.Core.AstInterpret+import HordeAd.Core.AstSimplify+import HordeAd.Core.AstTools+import HordeAd.Core.AstVectorize+import HordeAd.Core.CarriersADVal+import HordeAd.Core.CarriersAst+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Delta+import HordeAd.Core.DeltaEval+import HordeAd.Core.Ops+import HordeAd.Core.OpsConcrete ()+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind++-- * Symbolic reverse and forward derivative computation++data IncomingCotangentHandling = UseIncomingCotangent | IgnoreIncomingCotangent++-- Here a choice is made that derivatives are PrimalSpan terms.+-- This makes them easier to simplify and expresses via type that they+-- don't introduce tangents nor cotangents, but are purely primal functions.+-- They can still be liften to dual number functions via interpretations,+-- as is done, e.g., in tgrad below.+forwardPassByInterpretation+  :: forall x z.+     (AstTensor AstMethodLet FullSpan x+      -> AstTensor AstMethodLet FullSpan z)+  -> AstEnv (ADVal (AstRaw PrimalSpan))+  -> AstTensor AstMethodShare PrimalSpan x+  -> AstVarName FullSpan x+  -> AstTensor AstMethodLet FullSpan x+  -> ADVal (AstRaw PrimalSpan) z+{-# INLINE forwardPassByInterpretation #-}+forwardPassByInterpretation g envInit astVarPrimal var astVar0 =+  let deltaInputs = generateDeltaInputs $ varNameToFTK var+      varInputs = dDnotShared (AstRaw astVarPrimal) deltaInputs+      ast = g astVar0+      env = extendEnv var varInputs envInit+  in interpretAstFull env ast++revArtifactFromForwardPass+  :: forall x z.+     IncomingCotangentHandling+  -> (AstTensor AstMethodShare PrimalSpan x+      -> AstVarName FullSpan x+      -> AstTensor AstMethodLet FullSpan x+      -> ADVal (AstRaw PrimalSpan) z)+  -> FullShapeTK x+  -> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)+-- Break the inline chain to prevent false positives in inspection testing+-- and protect the unsafePerformIO.+{-# NOINLINE revArtifactFromForwardPass #-}+revArtifactFromForwardPass cotangentHandling+                           forwardPass xftk = unsafePerformIO $ do+  -- IO and bangs and the compound function to fix the numbering of variables+  -- for pretty-printing and prevent sharing the impure values+  -- in tests that reset the impure counters.+  (!varPrimal, astVarPrimal, var, astVar0) <- funToAstRevIO xftk+  -- Evaluate completely after terms constructed, to free memory+  -- before gradientFromDelta allocates new memory and new FFI is started.+  let !(D primalBody delta) = forwardPass astVarPrimal var astVar0+  let zftk = ftkAst $ unAstRaw primalBody+      (!varDt, astDt) = funToAst (adFTK zftk) Nothing id+  let oneAtF = treplTarget 1 $ adFTK zftk+      !dt = case cotangentHandling of+        UseIncomingCotangent -> AstRaw astDt+        IgnoreIncomingCotangent -> oneAtF+  let !gradient = gradientFromDelta xftk zftk dt delta+      !unGradient = unshareAstTensor $ unAstRaw gradient+      !unPrimal = unshareAstTensor $ unAstRaw primalBody+  return (AstArtifactRev varDt varPrimal unGradient unPrimal, delta)++revProduceArtifact+  :: forall x z.+     IncomingCotangentHandling+  -> (AstTensor AstMethodLet FullSpan x+      -> AstTensor AstMethodLet FullSpan z)+  -> AstEnv (ADVal (AstRaw PrimalSpan))+  -> FullShapeTK x+  -> AstArtifactRev x z+revProduceArtifact cotangentHandling g envInit xftk =+  fst $ inline revArtifactFromForwardPass+          cotangentHandling (forwardPassByInterpretation g envInit) xftk++fwdArtifactFromForwardPass+  :: forall x z.+     (AstTensor AstMethodShare PrimalSpan x+      -> AstVarName FullSpan x+      -> AstTensor AstMethodLet FullSpan x+      -> ADVal (AstRaw PrimalSpan) z)+  -> FullShapeTK x+  -> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)+-- Break the inline chain to prevent false positives in inspection testing+-- and protect the unsafePerformIO.+{-# NOINLINE fwdArtifactFromForwardPass #-}+fwdArtifactFromForwardPass forwardPass xftk = unsafePerformIO $ do+  (!varPrimalD, astVarD, varPrimal, astVarPrimal, var, astVar0)+    <- funToAstFwdIO xftk+  let !(D primalBody delta) = forwardPass astVarPrimal var astVar0+  let !derivative = derivativeFromDelta @x delta (adFTK xftk) (AstRaw astVarD)+      !unDerivative = unshareAstTensor $ unAstRaw derivative+      !unPrimal = unshareAstTensor $ unAstRaw primalBody+  return (AstArtifactFwd varPrimalD varPrimal unDerivative unPrimal, delta)++fwdProduceArtifact+  :: forall x z.+     (AstTensor AstMethodLet FullSpan x+      -> AstTensor AstMethodLet FullSpan z)+  -> AstEnv (ADVal (AstRaw PrimalSpan))+  -> FullShapeTK x+  -> AstArtifactFwd x z+{-# INLINE fwdProduceArtifact #-}+fwdProduceArtifact f envInit xftk =+  fst $ inline fwdArtifactFromForwardPass+          (forwardPassByInterpretation f envInit) xftk+++-- * AstTensor instances++-- | This is a vectorizing combinator that also simplifies+-- the terms touched during vectorization, but not any others.+-- Due to how the Ast tensor instances are defined, vectorization+-- works bottom-up, which removes the need to backtrack in the vectorization+-- pass or repeat until a fixed point is reached.+-- This combinator also introduces new variable names.+astBuild1Vectorize+  :: AstSpan s+  => SNat k -> SingletonTK y+  -> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)+  -> AstTensor AstMethodLet s (BuildTensorKind k y)+astBuild1Vectorize k@(SNat @k) stk f =+  build1Vectorize k stk $ funToAstI (Just (0, valueOf @k - 1)) f++instance AstSpan s => LetTensor (AstTensor AstMethodLet s) where+  ttlet = astLetFun+  ttletPrimal = astLetFun+  toShare t = AstRaw $ AstToShare t+  -- For convenience and simplicity we define this for all spans,+  -- but it can only ever be used for PrimalSpan.+  tunshare =+    case sameAstSpan @s @PrimalSpan of+      Just Refl -> unshareAstTensor . unAstRaw+      _ -> error "tunshare: used not at PrimalSpan"++-- | The checks and error messages in these functions result in complete+-- shape-checking of the ranked and mixed user code (shaped is already+-- fully checked by the Haskell type system).+instance AstSpan s => BaseTensor (AstTensor AstMethodLet s) where+  -- Ranked ops+  rshape t = case ftkAst t of+    FTKR sh _ -> sh+  trsum v = withSNat (rwidth v) $ \snat -> astSum snat knownSTK v+  trreplicate k = withSNat k $ \snat -> astReplicate snat knownSTK+  trindex @m @n a ix = case ftkAst a of+    FTKR @_ @x shmshn x ->+      withShsFromShR shmshn $ \(sh :: ShS sh) ->+        withKnownShS sh $+        gcastWith (unsafeCoerceRefl :: Rank (Take m sh) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Drop m sh) :~: n) $+        gcastWith (unsafeCoerceRefl :: Take m sh ++ Drop m sh :~: sh) $+        withKnownShS (shsTake @m sh) $+        astFromS' @(TKS2 (Drop m sh) x) (FTKR (shrDrop @m shmshn) x)+        $ astIndexS @(Take m sh) @(Drop m sh)+                    (shsDrop @m sh) (astSFromR' @sh sh a)+                    (ixsFromIxR' knownShS ix)+  trscatter @m @_ @p shpshn0 t f = case ftkAst t of+    FTKR @_ @x shmshn0 x ->+      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->+        withKnownShS shmshn $+        withKnownShS shpshn $+        withKnownShS (shsTake @m shmshn) $+        withKnownShS (shsDrop @m shmshn) $+        withKnownShS (shsTake @p shpshn) $+        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $+        gcastWith (unsafeCoerceRefl+                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $+        case testEquality (shsDrop @p shpshn) (shsDrop @m shmshn) of+          Just Refl ->+            astFromS' @(TKS2 shpshn x) (FTKR shpshn0 x)+            $ astScatterS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)+                          knownShS (astSFromR' shmshn t)+            $ funToAstIxS knownShS (ixsFromIxR' knownShS . f . ixrFromIxS)+                -- this introduces new variable names+          _ -> error $ "rscatter: shapes don't match: "+                       ++ show (shsDrop @p shpshn, shsDrop @m shmshn)+  trgather @m @_ @p shmshn0 t f = case ftkAst t of+    FTKR shpshn0 x ->+      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->+        withKnownShS shmshn $+        withKnownShS shpshn $+        withKnownShS (shsTake @m shmshn) $+        withKnownShS (shsDrop @m shmshn) $+        withKnownShS (shsTake @p shpshn) $+        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $+        gcastWith (unsafeCoerceRefl+                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $+        case testEquality (shsDrop @p shpshn) (shsDrop @m shmshn) of+          Just Refl ->+            astFromS' (FTKR shmshn0 x)+            $ astGatherS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)+                         knownShS (astSFromR' shpshn t)+            $ funToAstIxS knownShS (ixsFromIxR' knownShS . f . ixrFromIxS)+                -- this introduces new variable names+          _ -> error $ "rgather: shapes don't match: "+                       ++ show (shsDrop @p shpshn, shsDrop @m shmshn)+  trconcrete a = tconcrete (FTKR (Nested.rshape a) FTKScalar) (Concrete a)+  trfloor @_ @r2 a = case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        astFromS' @(TKS sh r2) (FTKR sh' FTKScalar)+        . fromPrimal . astFloorS . primalPart . astSFromR' @sh sh $ a+  trfromIntegral @_ @r2 a = case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        astFromS' @(TKS sh r2) (FTKR sh' FTKScalar)+        . fromPrimal . astFromIntegralS . primalPart . astSFromR' @sh sh $ a+  trcast @_ @r2 a = case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        astFromS' @(TKS sh r2) (FTKR sh' FTKScalar)+        . astCastS . astSFromR' sh $ a+  trminIndex @_ @_ @r2 a = case ftkAst a of+    FTKR sh' _ ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @_ @rest _ _ ->+          -- unfortunately, this is not enough:+          -- gcastWith (unsafeCoerceRefl :: Rank sh :~: 1 + Rank (Init sh)) $+          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $+          astFromS' @(TKS (Init sh) r2) (FTKR (shrInit sh') FTKScalar)+          . fromPrimal . AstMinIndexS . primalPart . astSFromR' @sh sh $ a+        ZSS -> error "rminIndex: impossible empty shape"+  trmaxIndex @_ @_ @r2 a = case ftkAst a of+    FTKR sh' _ ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @_ @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $+          astFromS' @(TKS (Init sh) r2) (FTKR (shrInit sh') FTKScalar)+          . fromPrimal . AstMaxIndexS . primalPart . astSFromR' @sh sh $ a+        ZSS -> error "rmaxIndex: impossible empty shape"+  triota @r n =+    withSNat n $ \(SNat @n) ->+      astFromS' (FTKR (n :$: ZSR) FTKScalar)+      $ fromPrimal $ AstIotaS @n @r SNat+  trappend u v = case ftkAst u of+    FTKR shu' x -> case ftkAst v of+      FTKR shv' _ ->+        withShsFromShR shu' $ \shu -> case shu of+          _ :$$ restu ->+            withShsFromShR shv' $ \shv -> case shv of+              _ :$$ restv ->+                case testEquality restu restv of+                  Just Refl ->+                    astFromS' (FTKR shu' x)+                    $ astAppendS (astSFromR' shu u) (astSFromR' shv v)+                  _ -> error $ "rappend: shapes don't match: "+                               ++ show (restu, restv)+              ZSS -> error "rappend: impossible shape"+          ZSS -> error "rappend: impossible shape"+  trslice i n a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh -> case sh of+        msnat@(SNat @m) :$$ _ ->+          withSNat i $ \isnat@(SNat @i) -> withSNat n $ \nsnat@(SNat @n) ->+            case cmpNat (snatPlus isnat nsnat) msnat of+              GTI -> error $ "rslice: argument tensor too narrow: "+                             ++ show (i, n, sNatValue msnat)+              EQI ->+                astFromS' (FTKR sh' x)+                . astSliceS isnat nsnat (SNat @(m - (i + n)))+                . astSFromR' sh $ a+              LTI ->+                astFromS' (FTKR sh' x)+                . astSliceS isnat nsnat (SNat @(m - (i + n)))+                . astSFromR' sh $ a+        ZSS -> error "xslice: impossible shape"+  trreverse a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh -> case sh of+        _ :$$ _ ->+          astFromS' (FTKR sh' x)+          . astReverseS . astSFromR' sh $ a+        ZSS -> error "xreverse: impossible shape"+  trtranspose @n @r permr a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \(sh :: ShS sh)  ->+        Permutation.permFromList permr $ \(perm :: Permutation.Perm perm) ->+          let result :: AstTensor AstMethodLet s (TKR2 n r)+              result =+                -- A noble lie, verified down below.+                gcastWith (unsafeCoerceRefl+                           :: (Rank perm <=? Rank sh) :~: True) $+                fromMaybe (error "rtranspose: impossible non-permutation")+                $ Permutation.permCheckPermutation perm+                $ astFromS' (FTKR (shrPermutePrefix permr sh') x)+                  . astTransposeS perm . astSFromR' sh $ a+          in case (Permutation.permRank perm, shsRank sh) of+            (psnat@SNat, shsnat@SNat) ->+              case cmpNat psnat shsnat of+                GTI -> error $ "rtranspose: rank mismatch: "+                               ++ show (sNatValue psnat, sNatValue shsnat)+                EQI -> result+                LTI -> result+  trreshape sh2' a = case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh ->+      withShsFromShR sh2' $ \sh2 ->+        case testEquality (shsProduct sh) (shsProduct sh2) of+          Just Refl -> astFromS' (FTKR sh2' x)+                       . astReshapeS sh2 . astSFromR' sh $ a+          _ -> error $ "rreshape: tensor size mismatch: "+                       ++ show ( sNatValue (shsProduct sh)+                               , sNatValue (shsProduct sh2) )+  trbuild1 @n @x k f = withSNat k $ \snat ->+    astBuild1Vectorize snat (STKR (SNat @n) (knownSTK @x)) f++  -- Shaped ops+  sshape t = case ftkAst t of+    FTKS sh _ -> sh+  tssum = astSum SNat knownSTK+  tsindex = astIndexS knownShS+  tsscatter @shm @shn @shp t f =+    astScatterS @shm @shn @shp knownShS t+    $ funToAstIxS knownShS f  -- this introduces new variable names+  tsgather @shm @shn @shp t f =+    astGatherS @shm @shn @shp knownShS t+    $ funToAstIxS knownShS f  -- this introduces new variable names+  tsconcrete = fromPrimal . AstConcreteS+  tsfloor = fromPrimal . astFloorS . primalPart+  tsfromIntegral = fromPrimal . astFromIntegralS . primalPart+  tscast = astCastS+  tsminIndex = fromPrimal . AstMinIndexS . primalPart+  tsmaxIndex = fromPrimal . AstMaxIndexS . primalPart+  tsiota = fromPrimal $ AstIotaS SNat+  tsappend = astAppendS+  tsslice = astSliceS+  tsreverse = astReverseS+  tsbuild1 @k @sh @x =+    astBuild1Vectorize (SNat @k) (STKS (knownShS @sh) (knownSTK @x))++  -- Mixed ops+  xshape t = case ftkAst t of+    FTKX sh _ -> sh+  txsum = astSum SNat knownSTK+  txreplicate snat sh = astReplicate snat (STKX sh knownSTK)+  txindex @sh1 @sh2 a ix = case ftkAst a of+    FTKX @sh1sh2 @x sh1sh2 x | SNat <- ssxRank (knownShX @sh1) ->+      withShsFromShX sh1sh2 $ \(sh :: ShS sh) ->+        withKnownShX (ssxFromShX sh1sh2) $+        withKnownShS sh $+        gcastWith (unsafeCoerceRefl :: Rank (Drop (Rank sh1) sh) :~: Rank sh2) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank sh1) sh ++ Drop (Rank sh1) sh :~: sh) $+        gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) sh1sh2 :~: sh2) $+        withKnownShS (shsTake @(Rank sh1) sh) $+        astFromS' @(TKS2 (Drop (Rank sh1) sh) x)+                  (FTKX (shxDrop @(Rank sh1) sh1sh2) x)+        $ astIndexS @(Take (Rank sh1) sh) @(Drop (Rank sh1) sh)+                    (shsDrop @(Rank sh1) sh) (astSFromX' @sh @sh1sh2 sh a)+                    (ixsFromIxX' knownShS ix)+  txscatter @shm @_ @shp shpshn0 t f = case ftkAst t of+    FTKX shmshn0 x | SNat <- ssxRank (knownShX @shm)+                   , SNat <- ssxRank (knownShX @shp) ->+      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->+        withKnownShS shmshn $+        withKnownShS shpshn $+        withKnownShS (shsTake @(Rank shm) shmshn) $+        withKnownShS (shsDrop @(Rank shm) shmshn) $+        withKnownShS (shsTake @(Rank shp) shpshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn+                      :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn+                      :~: shpshn) $+        case testEquality (shsDrop @(Rank shp) shpshn)+                          (shsDrop @(Rank shm) shmshn) of+          Just Refl ->+            astFromS' (FTKX shpshn0 x)+            $ astScatterS @(Take (Rank shm) shmshn)+                          @(Drop (Rank shm) shmshn)+                          @(Take (Rank shp) shpshn)+                          knownShS (astSFromX' shmshn t)+            $ funToAstIxS knownShS (ixsFromIxX' knownShS+                                    . f . ixxCast knownShX . ixxFromIxS)+                -- this introduces new variable names+          _ -> error $ "xscatter: shapes don't match: "+                       ++ show ( shsDrop @(Rank shp) shpshn+                               , shsDrop @(Rank shm) shmshn )+  txgather @shm @_ @shp shmshn0 t f = case ftkAst t of+    FTKX shpshn0 x | SNat <- ssxRank (knownShX @shm)+                   , SNat <- ssxRank (knownShX @shp) ->+      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->+        withKnownShS shmshn $+        withKnownShS shpshn $+        withKnownShS (shsTake @(Rank shm) shmshn) $+        withKnownShS (shsDrop @(Rank shm) shmshn) $+        withKnownShS (shsTake @(Rank shp) shpshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn+                      :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn+                      :~: shpshn) $+        case testEquality (shsDrop @(Rank shp) shpshn)+                          (shsDrop @(Rank shm) shmshn) of+          Just Refl ->+            astFromS' (FTKX shmshn0 x)+            $ astGatherS @(Take (Rank shm) shmshn)+                         @(Drop (Rank shm) shmshn)+                         @(Take (Rank shp) shpshn)+                         knownShS (astSFromX' shpshn t)+            $ funToAstIxS knownShS (ixsFromIxX' knownShS+                                    . f . ixxCast knownShX . ixxFromIxS)+                -- this introduces new variable names+          _ -> error $ "xgather: shapes don't match: "+                       ++ show ( shsDrop @(Rank shp) shpshn+                               , shsDrop @(Rank shm) shmshn )+  txconcrete a = tconcrete (FTKX (Nested.mshape a) FTKScalar) (Concrete a)+  txfloor @_ @r2 @sh' a = case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        astFromS' @(TKS sh r2) (FTKX sh' FTKScalar)+        . fromPrimal . astFloorS . primalPart . astSFromX' @sh @sh' sh $ a+  txfromIntegral @_ @r2 @sh' a = case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        astFromS' @(TKS sh r2) (FTKX sh' FTKScalar)+        . fromPrimal . astFromIntegralS+        . primalPart . astSFromX' @sh @sh' sh $ a+  txcast @_ @r2 a = case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        astFromS' @(TKS sh r2) (FTKX sh' FTKScalar)+        . astCastS . astSFromX' sh $ a+  txminIndex @_ @_ @_ @r2 a = case ftkAst a of+    FTKX @sh' sh' _ ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @n @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $+          astFromS' @(TKS (Init sh) r2)+                    (FTKX (shxInit sh') FTKScalar)+          . fromPrimal . AstMinIndexS @n @rest+          . primalPart . astSFromX' @sh @sh' sh $ a+  txmaxIndex @_ @_ @_ @r2 a = case ftkAst a of+    FTKX @sh' sh' _ ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @n @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $+          astFromS' @(TKS (Init sh) r2)+                   (FTKX (shxInit sh') FTKScalar)+          . fromPrimal . AstMaxIndexS @n @rest+          . primalPart . astSFromX' @sh @sh' sh $ a+  txiota @n @r = astFromS' (FTKX (SKnown (SNat @n) :$% ZSX) FTKScalar)+                 $ fromPrimal $ AstIotaS @n @r SNat+  txappend u v = case ftkAst u of+    FTKX (Nested.SKnown m@SNat :$% shu') x -> case ftkAst v of+      FTKX (Nested.SKnown n@SNat :$% shv') _ ->+        withShsFromShX shu' $ \(shu :: ShS shu) ->+          withShsFromShX shv' $ \(shv :: ShS shv) ->+            case shxEqual shu' shv' of+              Just Refl ->+                gcastWith (unsafeCoerceRefl :: shu :~: shv) $+                astFromS' (FTKX (Nested.SKnown (snatPlus m n) :$% shu') x)+                $ astAppendS (astSFromX' (m :$$ shu) u)+                             (astSFromX' (n :$$ shv) v)+              _ -> error $ "xappend: shapes don't match: "+                           ++ show (shu', shv')+  txslice i n@SNat k a = case ftkAst a of+    FTKX sh'@(_ :$% sh2') x ->+      withShsFromShX sh' $ \sh@(msnat :$$ _) ->+        case testEquality (snatPlus i (snatPlus n k)) msnat of+          Just Refl ->+            astFromS' (FTKX (SKnown n :$% sh2') x)+            . astSliceS i n k . astSFromX' sh $ a+          _ -> error $ "xslice: argument tensor too narrow: "+                       ++ show ( sNatValue i, sNatValue n, sNatValue k+                               , sNatValue msnat )+  txreverse a = case ftkAst a of+    FTKX sh' x ->+      withShsFromShX sh' $ \(sh@(_ :$$ _) :: ShS sh) ->+        astFromS' (FTKX sh' x)+        . astReverseS . astSFromX' @sh sh $ a+  txtranspose perm a = case ftkAst a of+    FTKX sh' x ->+      let sh2' = shxPermutePrefix perm sh'+      in withShsFromShX sh' $ \sh ->+           astFromS' (FTKX sh2' x)+           . astTransposeS perm . astSFromX' sh $ a+  txreshape sh2' a = case ftkAst a of+    FTKX sh' x ->+      withShsFromShX sh' $ \sh ->+      withShsFromShX sh2' $ \sh2 ->+        case testEquality (shsProduct sh) (shsProduct sh2) of+          Just Refl ->+            astFromS' (FTKX sh2' x)+            . astReshapeS sh2 . astSFromX' sh $ a+          _ -> error $ "xreshape: tensor size mismatch: "+                       ++ show ( sNatValue (shsProduct sh)+                               , sNatValue (shsProduct sh2) )+  txbuild1 @k @sh @x =+    astBuild1Vectorize (SNat @k) (STKX (knownShX @sh) (knownSTK @x))++  -- Scalar ops+  tkconcrete = fromPrimal . AstConcreteK+  tkfloor = fromPrimal . astFloorK . primalPart+  tkfromIntegral = fromPrimal . astFromIntegralK . primalPart+  tkcast = astCastK++  -- General operations that don't require LetTensor nor ShareTensor+  tftk _stk = ftkAst+  tconcrete ftk a = fromPrimal $ astConcrete ftk a+  tpair = astPair+  tproject1 = astProject1+  tproject2 = astProject2+  tsreplicate snat sh = astReplicate snat (STKS sh knownSTK)+  tstranspose = astTransposeS+  tsreshape = astReshapeS+  tmapAccumRDer _ !k _ !bftk !eftk f df rf acc0 es =+    astMapAccumRDer k bftk eftk f df rf acc0 es+  tmapAccumLDer _ !k _ !bftk !eftk f df rf acc0 es =+    astMapAccumLDer k bftk eftk f df rf acc0 es+  tApply = astApply+  tlambda ftk f =+    let (var, ast) = funToAst ftk Nothing $ unHFun f+    in AstLambda var ast+  tcond _ !b !u !v = astCond b u v+  tprimalPart = primalPart+  tdualPart _ = dualPart+  tfromPrimal _ = fromPrimal+  tfromDual = fromDual+  tgrad xftk f =+    -- We don't have an AST constructor to hold it, so we compute outright.+    --+    -- This computes the (AST of) derivative of f once and interprets it again+    -- for each new tensor of arguments, which is better than computing it anew.+    let -- No bangs here, because this goes under lambda and should not be+        -- evaluated too early (which at some point was even incorrect+        -- and triggered error "tunshare: used not at PrimalSpan"; maybe this+        -- is related to terms getting spans converted when interpreted)+        AstArtifactRev{..} =+          revProduceArtifact+            IgnoreIncomingCotangent (simplifyInline . unHFun f) emptyEnv xftk+        -- A new variable is created to give it the right span as opposed+        -- to the fixed PrimalSpan that artVarDomainRev has.+        (varP, ast) = funToAst xftk Nothing $ \ !astP ->+          let env = extendEnv artVarDomainRev astP emptyEnv+          in interpretAst env $ simplifyInline artDerivativeRev+    in AstLambda varP ast+  tvjp ftkx f =+    -- This computes the (AST of) derivative of f once and interprets it again+    -- for each new tensor of arguments, which is better than computing it anew.+    let AstArtifactRev{..} =+          revProduceArtifact+            UseIncomingCotangent (simplifyInline . unHFun f) emptyEnv ftkx+        ftkz = varNameToFTK artVarDtRev+        ftk2 = FTKProduct ftkz ftkx+        (varP, ast) = funToAst ftk2 Nothing $ \ !astP ->+          let env = extendEnv artVarDtRev (astProject1 astP)+                    $ extendEnv artVarDomainRev (astProject2 astP) emptyEnv+          in interpretAst env $ simplifyInline artDerivativeRev+    in AstLambda varP ast+  tjvp ftkx f =+    -- This computes the (AST of) derivative of f once and interprets it again+    -- for each new tensor of arguments, which is better than computing it anew.+    let AstArtifactFwd{..} =+          fwdProduceArtifact (simplifyInline . unHFun f) emptyEnv ftkx+        ftk2 = FTKProduct (adFTK ftkx) ftkx+        (varP, ast) = funToAst ftk2 Nothing $ \ !astP ->+          let env = extendEnv artVarDsFwd (astProject1 astP)+                    $ extendEnv artVarDomainFwd (astProject2 astP) emptyEnv+          in interpretAst env $ simplifyInline artDerivativeFwd+    in AstLambda varP ast++  tfromVector = astFromVector+  tsum snat@SNat stk u = case stk of+    STKScalar -> kfromS $ tssum u+    STKR SNat x | Dict <- lemKnownSTK x -> trsum u+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (tsum snat stk1 (tproject1 u3))+              (tsum snat stk2 (tproject2 u3))+  treplicate snat@SNat stk u = case stk of+    STKScalar -> tsreplicate snat ZSS $ sfromK u+    STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (sNatValue snat) u+    STKS sh x | Dict <- lemKnownSTK x -> tsreplicate snat sh u+    STKX sh x | Dict <- lemKnownSTK x -> txreplicate snat sh u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (treplicate snat stk1 (tproject1 u3))+              (treplicate snat stk2 (tproject2 u3))+  tindexBuild snat@SNat stk u i = case stk of+    STKScalar -> kfromS $ tsindex u (i :.$ ZIS)+    STKR SNat x | Dict <- lemKnownSTK x -> trindex u (i :.: ZIR)+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsindex u (i :.$ ZIS)+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txindex u (i :.% ZIX)+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (tindexBuild snat stk1 (tproject1 u3) i)+              (tindexBuild snat stk2 (tproject2 u3) i)++  treplTarget = replTarget+  tdefTarget = defTarget+  taddTarget = addTarget+  tmultTarget = multTarget+  tsum0Target = sum0Target+  tdot0Target = dot0Target+++-- * AstRaw instances++instance AstSpan s => ShareTensor (AstRaw s) where+  -- For convenience and simplicity we define this for all spans,+  -- but it can only ever be used for PrimalSpan.+  tshare t =+    let u = unAstRaw t+    in if astIsSmall True u+       then t+       else AstRaw $ fun1ToAst (ftkAst u) $ \ !var -> AstShare var u+  tunpair (AstRaw (AstPair t1 t2)) = (AstRaw t1, AstRaw t2)+  tunpair t = let tShared = tshare t+              in (tproject1 tShared, tproject2 tShared)++instance AstSpan s => BaseTensor (AstRaw s) where+  -- Ranked ops+  rshape t = case ftkAst $ unAstRaw t of+    FTKR sh _ -> sh+  trsum v = withSNat (rwidth v) $ \snat ->+              AstRaw . AstSum snat knownSTK . unAstRaw $ v+  trreplicate k = withSNat k $ \snat ->+    AstRaw . AstReplicate snat knownSTK . unAstRaw+  trindex @m @n (AstRaw a) ix = AstRaw $ case ftkAst a of+    FTKR @_ @x shmshn x ->+      withShsFromShR shmshn $ \(sh :: ShS sh) ->+        withKnownShS sh $+        gcastWith (unsafeCoerceRefl :: Rank (Take m sh) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Drop m sh) :~: n) $+        gcastWith (unsafeCoerceRefl :: Take m sh ++ Drop m sh :~: sh) $+        withKnownShS (shsTake @m sh) $+        cAstFromS @(TKS2 (Drop m sh) x) (FTKR (shrDrop @m shmshn) x)+        $ AstIndexS @(Take m sh) @(Drop m sh)+                    (shsDrop @m sh) (cAstSFromR @sh sh a)+                    (ixsFromIxR knownShS (unAstRaw <$> ix))+  trscatter @m @_ @p shpshn0 (AstRaw t) f = AstRaw $ case ftkAst t of+    FTKR @_ @x shmshn0 x ->+      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->+        withKnownShS shmshn $+        withKnownShS shpshn $+        withKnownShS (shsTake @m shmshn) $+        withKnownShS (shsDrop @m shmshn) $+        withKnownShS (shsTake @p shpshn) $+        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $+        gcastWith (unsafeCoerceRefl+                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $+        case testEquality (shsDrop @p shpshn) (shsDrop @m shmshn) of+          Just Refl ->+            cAstFromS @(TKS2 shpshn x) (FTKR shpshn0 x)+            $ AstScatterS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)+                          knownShS (cAstSFromR shmshn t)+            $ funToAstIxS knownShS (fmap unAstRaw . ixsFromIxR knownShS+                                    . f . ixrFromIxS+                                    . fmap AstRaw)+                -- this introduces new variable names+          _ -> error $ "rscatter: shapes don't match: "+                       ++ show (shsDrop @p shpshn, shsDrop @m shmshn)+  trgather @m @_ @p shmshn0 (AstRaw t) f = AstRaw $ case ftkAst t of+    FTKR shpshn0 x ->+      withShsFromShR shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShR shpshn0 $ \(shpshn :: ShS shpshn) ->+        withKnownShS shmshn $+        withKnownShS shpshn $+        withKnownShS (shsTake @m shmshn) $+        withKnownShS (shsDrop @m shmshn) $+        withKnownShS (shsTake @p shpshn) $+        gcastWith (unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) $+        gcastWith (unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) $+        gcastWith (unsafeCoerceRefl+                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) $+        case testEquality (shsDrop @p shpshn) (shsDrop @m shmshn) of+          Just Refl ->+            cAstFromS (FTKR shmshn0 x)+            $ AstGatherS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)+                         knownShS (cAstSFromR shpshn t)+            $ funToAstIxS knownShS (fmap unAstRaw . ixsFromIxR knownShS+                                    . f . ixrFromIxS+                           . fmap AstRaw)+                -- this introduces new variable names+          _ -> error $ "rgather: shapes don't match: "+                       ++ show (shsDrop @p shpshn, shsDrop @m shmshn)+  trconcrete a = tconcrete (FTKR (Nested.rshape a) FTKScalar) (Concrete a)+  trfloor @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        cAstFromS @(TKS sh r2) (FTKR sh' FTKScalar)+        . fromPrimal . AstFloorS . primalPart . cAstSFromR @sh sh $ a+  trfromIntegral @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        cAstFromS @(TKS sh r2) (FTKR sh' FTKScalar)+        . fromPrimal . AstFromIntegralS . primalPart . cAstSFromR @sh sh $ a+  trcast @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' FTKScalar ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        cAstFromS @(TKS sh r2) (FTKR sh' FTKScalar)+        . AstCastS . cAstSFromR sh $ a+  trminIndex @_ @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' _ ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @_ @rest _ _ ->+          -- unfortunately, this is not enough:+          -- gcastWith (unsafeCoerceRefl :: Rank sh :~: 1 + Rank (Init sh)) $+          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $+          cAstFromS @(TKS (Init sh) r2) (FTKR (shrInit sh') FTKScalar)+          . fromPrimal . AstMinIndexS . primalPart . cAstSFromR @sh sh $ a+        ZSS -> error "rminIndex: impossible shape"+  trmaxIndex @_ @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' _ ->+      withShsFromShR sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @_ @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) $+          cAstFromS @(TKS (Init sh) r2) (FTKR (shrInit sh') FTKScalar)+          . fromPrimal . AstMaxIndexS . primalPart . cAstSFromR @sh sh $ a+        ZSS -> error "rmaxIndex: impossible shape"+  triota @r n =+    AstRaw+    $ withSNat n $ \(SNat @n) ->+        cAstFromS (FTKR (n :$: ZSR) FTKScalar)+        $ fromPrimal $ AstIotaS @n @r SNat+  trappend (AstRaw u) (AstRaw v) = AstRaw $ case ftkAst u of+    ftk@(FTKR shu' _) -> case ftkAst v of+      FTKR shv' _ ->+        withShsFromShR shu' $ \shu -> case shu of+          _ :$$ restu ->+            withShsFromShR shv' $ \shv -> case shv of+              _ :$$ restv ->+                case testEquality restu restv of+                  Just Refl ->+                    cAstFromS ftk+                    $ AstAppendS (cAstSFromR shu u) (cAstSFromR shv v)+                  _ -> error $ "rappend: shapes don't match: "+                               ++ show (restu, restv)+              ZSS -> error "rappend: impossible shape"+          ZSS -> error "rappend: impossible shape"+  trslice i n (AstRaw a) = AstRaw $ case ftkAst a of+    ftk@(FTKR sh' _) ->+      withShsFromShR sh' $ \sh -> case sh of+        msnat@(SNat @m) :$$ _ ->+          withSNat i $ \isnat@(SNat @i) -> withSNat n $ \nsnat@(SNat @n) ->+            case cmpNat (snatPlus isnat nsnat) msnat of+              GTI -> error $ "rslice: argument tensor too narrow: "+                             ++ show (i, n, sNatValue msnat)+              EQI ->+                cAstFromS ftk+                . AstSliceS isnat nsnat (SNat @(m - (i + n)))+                . cAstSFromR sh $ a+              LTI ->+                cAstFromS ftk+                . AstSliceS isnat nsnat (SNat @(m - (i + n)))+                . cAstSFromR sh $ a+        ZSS -> error "xslice: impossible shape"+  trreverse (AstRaw a) = AstRaw $ case ftkAst a of+    ftk@(FTKR sh' _) ->+      withShsFromShR sh' $ \sh -> case sh of+        _ :$$ _ ->+          cAstFromS ftk+          . AstReverseS . cAstSFromR sh $ a+        ZSS -> error "xreverse: impossible shape"+  trtranspose @n @r permr (AstRaw a) = AstRaw $ case ftkAst a of+    ftk@(FTKR sh' _) ->+      withShsFromShR sh' $ \(sh :: ShS sh) ->+        Permutation.permFromList permr $ \(perm :: Permutation.Perm perm) ->+          let result :: AstTensor AstMethodShare s (TKR2 n r)+              result =+                -- A noble lie, verified down below.+                gcastWith (unsafeCoerceRefl+                           :: (Rank perm <=? Rank sh) :~: True) $+                fromMaybe (error "rtranspose: impossible non-permutation")+                $ Permutation.permCheckPermutation perm+                $ cAstFromS ftk+                  . AstTransposeS perm . cAstSFromR sh $ a+          in case (Permutation.permRank perm, shsRank sh) of+            (psnat@SNat, shsnat@SNat) ->+              case cmpNat psnat shsnat of+                GTI -> error $ "rtranspose: rank mismatch: "+                               ++ show (sNatValue psnat, sNatValue shsnat)+                EQI -> result+                LTI -> result+  trreshape sh2' (AstRaw a) = AstRaw $ case ftkAst a of+    FTKR sh' x ->+      withShsFromShR sh' $ \sh ->+      withShsFromShR sh2' $ \sh2 ->+        case testEquality (shsProduct sh) (shsProduct sh2) of+          Just Refl -> cAstFromS (FTKR sh2' x)+                       . AstReshapeS sh2 . cAstSFromR sh $ a+          _ -> error $ "rreshape: tensor size mismatch: "+                       ++ show ( sNatValue (shsProduct sh)+                               , sNatValue (shsProduct sh2) )+  trbuild1 k f = withSNat k $ \snat ->+    AstRaw $ AstBuild1 snat knownSTK+    $ funToAstI (Just (0, fromIntegral k - 1))+        -- this introduces new variable names+    $ unAstRaw . f . AstRaw++  -- Shaped ops+  sshape t = case ftkAst $ unAstRaw t of+    FTKS sh _ -> sh+  tssum = AstRaw . AstSum SNat knownSTK . unAstRaw+  tsindex v ix = AstRaw $ AstIndexS knownShS (unAstRaw v) (unAstRaw <$> ix)+  tsscatter @shm @shn @shp t f =+    AstRaw $ AstScatterS @shm @shn @shp knownShS (unAstRaw t)+           $ funToAstIxS knownShS (fmap unAstRaw . f . fmap AstRaw)+               -- this introduces new variable names+  tsgather @shm @shn @shp t f =+    AstRaw $ AstGatherS @shm @shn @shp knownShS (unAstRaw t)+           $ funToAstIxS knownShS (fmap unAstRaw . f . fmap AstRaw)+               -- this introduces new variable names+  tsconcrete = AstRaw . fromPrimal . AstConcreteS+  tsfloor = AstRaw . fromPrimal . AstFloorS . primalPart . unAstRaw+  tsfromIntegral =+    AstRaw . fromPrimal . AstFromIntegralS . primalPart . unAstRaw+  tscast = AstRaw . AstCastS . unAstRaw+  tsminIndex = AstRaw . fromPrimal . AstMinIndexS . primalPart . unAstRaw+  tsmaxIndex = AstRaw . fromPrimal . AstMaxIndexS . primalPart . unAstRaw+  tsiota = AstRaw . fromPrimal $ AstIotaS SNat+  tsappend u v = AstRaw $ AstAppendS (unAstRaw u) (unAstRaw v)+  tsslice i n k = AstRaw . AstSliceS i n k . unAstRaw+  tsreverse = AstRaw . AstReverseS . unAstRaw+  tsbuild1 @k f = AstRaw $ AstBuild1 (SNat @k) knownSTK+                  $ funToAstI (Just (0, valueOf @k - 1))+                      -- this introduces new variable names+                  $ unAstRaw . f . AstRaw++  -- Mixed ops+  xshape t = case ftkAst $ unAstRaw t of+    FTKX sh _ -> sh+  txsum = AstRaw . AstSum SNat knownSTK . unAstRaw+  txreplicate snat sh = AstRaw . AstReplicate snat (STKX sh knownSTK) . unAstRaw+  txindex @sh1 @sh2 (AstRaw a) ix = case ftkAst a of+    FTKX @sh1sh2 @x sh1sh2 x | SNat <- ssxRank (knownShX @sh1) ->+      withShsFromShX sh1sh2 $ \(sh :: ShS sh) ->+        withKnownShX (ssxFromShX sh1sh2) $+        withKnownShS sh $+        gcastWith (unsafeCoerceRefl :: Rank (Drop (Rank sh1) sh) :~: Rank sh2) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank sh1) sh ++ Drop (Rank sh1) sh :~: sh) $+        gcastWith (unsafeCoerceRefl :: Drop (Rank sh1) sh1sh2 :~: sh2) $+        withKnownShS (shsTake @(Rank sh1) sh) $+        AstRaw $ cAstFromS @(TKS2 (Drop (Rank sh1) sh) x)+                           (FTKX (shxDrop @(Rank sh1) sh1sh2) x)+        $ AstIndexS @(Take (Rank sh1) sh) @(Drop (Rank sh1) sh)+                    (shsDrop @(Rank sh1) sh) (cAstSFromX @sh @sh1sh2 sh a)+                    (ixsFromIxX' knownShS (unAstRaw <$> ix))+  txscatter @shm @_ @shp shpshn0 (AstRaw t) f = AstRaw $ case ftkAst t of+    FTKX shmshn0 x | SNat <- ssxRank (knownShX @shm)+                   , SNat <- ssxRank (knownShX @shp) ->+      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->+        withKnownShS shmshn $+        withKnownShS shpshn $+        withKnownShS (shsTake @(Rank shm) shmshn) $+        withKnownShS (shsDrop @(Rank shm) shmshn) $+        withKnownShS (shsTake @(Rank shp) shpshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn+                      :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn+                      :~: shpshn) $+        case testEquality (shsDrop @(Rank shp) shpshn)+                          (shsDrop @(Rank shm) shmshn) of+          Just Refl ->+            cAstFromS (FTKX shpshn0 x)+            $ AstScatterS @(Take (Rank shm) shmshn)+                          @(Drop (Rank shm) shmshn)+                          @(Take (Rank shp) shpshn)+                          knownShS (cAstSFromX shmshn t)+            $ funToAstIxS knownShS (fmap unAstRaw+                                    . ixsFromIxX' knownShS+                                    . f . ixxCast knownShX . ixxFromIxS+                                    . fmap AstRaw)+                -- this introduces new variable names+          _ -> error $ "xscatter: shapes don't match: "+                       ++ show ( shsDrop @(Rank shp) shpshn+                               , shsDrop @(Rank shm) shmshn )+  txgather @shm @_ @shp shmshn0 (AstRaw t) f = AstRaw $ case ftkAst t of+    FTKX shpshn0 x | SNat <- ssxRank (knownShX @shm)+                   , SNat <- ssxRank (knownShX @shp) ->+      withShsFromShX shmshn0 $ \(shmshn :: ShS shmshn) ->+      withShsFromShX shpshn0 $ \(shpshn :: ShS shpshn) ->+        withKnownShS shmshn $+        withKnownShS shpshn $+        withKnownShS (shsTake @(Rank shm) shmshn) $+        withKnownShS (shsDrop @(Rank shm) shmshn) $+        withKnownShS (shsTake @(Rank shp) shpshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn+                      :~: shmshn) $+        gcastWith (unsafeCoerceRefl+                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn+                      :~: shpshn) $+        case testEquality (shsDrop @(Rank shp) shpshn)+                          (shsDrop @(Rank shm) shmshn) of+          Just Refl ->+            cAstFromS (FTKX shmshn0 x)+            $ AstGatherS @(Take (Rank shm) shmshn)+                         @(Drop (Rank shm) shmshn)+                         @(Take (Rank shp) shpshn)+                         knownShS (cAstSFromX shpshn t)+            $ funToAstIxS knownShS (fmap unAstRaw+                                    . ixsFromIxX' knownShS+                                    . f . ixxCast knownShX . ixxFromIxS+                                    . fmap AstRaw)+                -- this introduces new variable names+          _ -> error $ "xgather: shapes don't match: "+                       ++ show ( shsDrop @(Rank shp) shpshn+                               , shsDrop @(Rank shm) shmshn )+  txconcrete a = tconcrete (FTKX (Nested.mshape a) FTKScalar) (Concrete a)+  txfloor @_ @r2 @sh' (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        cAstFromS @(TKS sh r2) (FTKX sh' FTKScalar)+        . fromPrimal . AstFloorS . primalPart . cAstSFromX @sh @sh' sh $ a+  txfromIntegral @_ @r2 @sh' (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        cAstFromS @(TKS sh r2) (FTKX sh' FTKScalar)+        . fromPrimal . AstFromIntegralS+        . primalPart . cAstSFromX @sh @sh' sh $ a+  txcast @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' FTKScalar ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        cAstFromS @(TKS sh r2) (FTKX sh' FTKScalar)+        . AstCastS . cAstSFromX sh $ a+  txminIndex @_ @_ @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX @sh' sh' _ ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @n @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $+          cAstFromS @(TKS (Init sh) r2)+                    (FTKX (shxInit sh') FTKScalar)+          . fromPrimal . AstMinIndexS @n @rest+          . primalPart . cAstSFromX @sh @sh' sh $ a+  txmaxIndex @_ @_ @_ @r2 (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX @sh' sh' _ ->+      withShsFromShX sh' $ \(sh :: ShS sh) -> case sh of+        (:$$) @n @rest _ _ ->+          gcastWith (unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) $+          cAstFromS @(TKS (Init sh) r2)+                   (FTKX (shxInit sh') FTKScalar)+          . fromPrimal . AstMaxIndexS @n @rest+          . primalPart . cAstSFromX @sh @sh' sh $ a+  txiota @n @r = AstRaw $ cAstFromS (FTKX (SKnown (SNat @n) :$% ZSX) FTKScalar)+                 $ fromPrimal $ AstIotaS @n @r SNat+  txappend (AstRaw u) (AstRaw v) = AstRaw $ case ftkAst u of+    FTKX (Nested.SKnown m@SNat :$% shu') x -> case ftkAst v of+      FTKX (Nested.SKnown n@SNat :$% shv') _ ->+        withShsFromShX shu' $ \(shu :: ShS shu) ->+          withShsFromShX shv' $ \(shv :: ShS shv) ->+            case shxEqual shu' shv' of+              Just Refl ->+                gcastWith (unsafeCoerceRefl :: shu :~: shv) $+                cAstFromS (FTKX (SKnown (snatPlus m n) :$% shu') x)+                $ AstAppendS (cAstSFromX (m :$$ shu) u)+                             (cAstSFromX (n :$$ shv) v)+              _ -> error $ "xappend: shapes don't match: "+                           ++ show (shu', shv')+  txslice i n@SNat k (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh'@(_ :$% sh2') x ->+      withShsFromShX sh' $ \sh@(msnat :$$ _) ->+        case testEquality (snatPlus i (snatPlus n k)) msnat of+          Just Refl ->+            cAstFromS (FTKX (SKnown n :$% sh2') x)+            . AstSliceS i n k . cAstSFromX sh $ a+          _ -> error $ "xslice: argument tensor too narrow: "+                       ++ show ( sNatValue i, sNatValue n, sNatValue k+                               , sNatValue msnat )+  txreverse (AstRaw a) = AstRaw $ case ftkAst a of+    ftk@(FTKX sh' _) ->+      withShsFromShX sh' $ \(sh@(_ :$$ _) :: ShS sh) ->+        cAstFromS ftk+        . AstReverseS . cAstSFromX @sh sh $ a+  txtranspose perm (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' x ->+      let sh2' = shxPermutePrefix perm sh'+      in withShsFromShX sh' $ \sh ->+           cAstFromS (FTKX sh2' x)+           . AstTransposeS perm+           . cAstSFromX sh $ a+  txreshape sh2' (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX sh' x ->+      withShsFromShX sh' $ \sh ->+      withShsFromShX sh2' $ \sh2 ->+        case testEquality (shsProduct sh) (shsProduct sh2) of+          Just Refl ->+            cAstFromS (FTKX sh2' x)+            . AstReshapeS sh2 . cAstSFromX sh $ a+          _ -> error $ "xreshape: tensor size mismatch: "+                       ++ show ( sNatValue (shsProduct sh)+                               , sNatValue (shsProduct sh2) )+  txbuild1 @k f = AstRaw $ AstBuild1 (SNat @k) knownSTK+                  $ funToAstI (Just (0, valueOf @k - 1))+                      -- this introduces new variable names+                  $ unAstRaw . f . AstRaw++  -- Scalar ops+  tkconcrete = AstRaw . fromPrimal . AstConcreteK+  tkfloor = AstRaw . fromPrimal . AstFloorK . primalPart . unAstRaw+  tkfromIntegral = AstRaw . fromPrimal . AstFromIntegralK+                   . primalPart . unAstRaw+  tkcast = AstRaw . AstCastK . unAstRaw++  -- General operations that don't require LetTensor nor ShareTensor+  tftk _stk = ftkAst . unAstRaw+  tconcrete ftk a = AstRaw $ fromPrimal $ unAstRaw $ astConcreteRaw ftk a+  tpair t1 t2 = AstRaw $ AstPair (unAstRaw t1) (unAstRaw t2)+  tproject1 t = AstRaw $ AstProject1 $ unAstRaw t+  tproject2 t = AstRaw $ AstProject2 $ unAstRaw t+  tsreplicate snat sh = AstRaw . AstReplicate snat (STKS sh knownSTK) . unAstRaw+  tstranspose perm = AstRaw . AstTransposeS perm . unAstRaw+  tsreshape sh = AstRaw . AstReshapeS sh . unAstRaw+  tmapAccumRDer _ !k _ !bftk !eftk f df rf acc0 es =+    AstRaw $ AstMapAccumRDer k bftk eftk f df rf (unAstRaw acc0) (unAstRaw es)+  tmapAccumLDer _ !k _ !bftk !eftk f df rf acc0 es =+    AstRaw $ AstMapAccumLDer k bftk eftk f df rf (unAstRaw acc0) (unAstRaw es)+  tApply t ll = AstRaw $ AstApply t (unAstRaw ll)+  tlambda = tlambda @(AstTensor AstMethodLet s)+  tcond _ !b !u !v = AstRaw $ AstCond b (unAstRaw u) (unAstRaw v)+  tprimalPart t = AstRaw $ primalPart $ unAstRaw t+  tdualPart _ t = dualPart $ unAstRaw t+  tfromPrimal _ t = AstRaw $ fromPrimal $ unAstRaw t+  tfromDual t = AstRaw $ fromDual t+  -- These three methods are called at this type in delta evaluation via+  -- tmapAccumR and tmapAccumL, so they have to work. We could refrain from+  -- simplifying the resulting terms, but it's not clear that's more consistent.+  tgrad = tgrad @(AstTensor AstMethodLet s)+  tvjp = tvjp @(AstTensor AstMethodLet s)+  tjvp = tjvp @(AstTensor AstMethodLet s)++  tfromVector k stk =+    AstRaw . AstFromVector k stk . V.map unAstRaw++  treplTarget = replTarget+  tdefTarget = defTarget+  taddTarget = addTarget+  tmultTarget = multTarget+  tsum0Target = sum0Target+  tdot0Target = dot0Target++instance AstSpan s => ConvertTensor (AstRaw s) where+  tconvert c _astk = AstRaw . cAstConvert c . unAstRaw++  rfromX a = case ftkAst $ unAstRaw a of+    FTKX sh' _ ->+      withShsFromShX sh' $ \(sh :: ShS sh) ->+        withKnownShS sh $+        rfromS $ sfromX @_ @sh a+  xfromR a = case ftkAst $ unAstRaw a of+    FTKR shr _ ->+      withShsFromShR shr $ \(sh :: ShS sh) ->+        withKnownShS sh $+        xfromS @_ @sh $ sfromR a++  sfromR = AstRaw . cAstSFromR knownShS . unAstRaw+  sfromX = AstRaw . cAstSFromX knownShS . unAstRaw+  xfromS @_ @sh' = AstRaw . cAstXFromS (knownShX @sh') . unAstRaw++  rzip @_ @_ @n (AstRaw a)+   | Refl <- lemRankReplicate (Proxy @n) = AstRaw $ case ftkAst a of+    FTKProduct (FTKR _sh y) (FTKR _ z) ->+      let c = ConvCmp+                (ConvXR (ftkToSTK (FTKProduct y z)))+                (ConvCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvRX ConvRX))+      in cAstConvert c a+  runzip @_ @_ @n (AstRaw a)+   | Refl <- lemRankReplicate (Proxy @n) = AstRaw $ case ftkAst a of+    FTKR _sh (FTKProduct y z) ->+      let c = ConvCmp+                (ConvT2 (ConvXR (ftkToSTK y)) (ConvXR (ftkToSTK z)))+                (ConvCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvRX)+      in cAstConvert c a+  szip (AstRaw a) = AstRaw $ case ftkAst a of+    FTKProduct (FTKS _sh y) (FTKS _ z) ->+      let c = ConvCmp+                ConvXS+                (ConvCmp+                   (ConvZip (ftkToSTK y) (ftkToSTK z))+                   (ConvT2 ConvSX ConvSX))+      in cAstConvert c a+  sunzip (AstRaw a) = AstRaw $ case ftkAst a of+    FTKS _sh (FTKProduct y z) ->+      let c = ConvCmp+                (ConvT2 ConvXS ConvXS)+                (ConvCmp+                   (ConvUnzip (ftkToSTK y) (ftkToSTK z))+                   ConvSX)+      in cAstConvert c a+  xzip (AstRaw a) = AstRaw $ case ftkAst a of+    FTKProduct (FTKX _sh y) (FTKX _ z) ->+      let c = ConvZip (ftkToSTK y) (ftkToSTK z)+      in cAstConvert c a+  xunzip (AstRaw a) = AstRaw $ case ftkAst a of+    FTKX _sh (FTKProduct y z) ->+      let c = ConvUnzip (ftkToSTK y) (ftkToSTK z)+      in cAstConvert c a++  xnestR @sh1 @m @x sh1 (AstRaw a)+    | Refl <- lemRankReplicate (Proxy @m) = AstRaw $+      let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)+                            (TKX2 sh1 (TKR2 m x))+          c = ConvCmp+                (ConvXX (ConvXR (knownSTK @x)))+                (ConvNest @_ @_ @(Replicate m Nothing)+                          (STKX sh1 (knownSTK @x)))+      in cAstConvert c a+  xnestS @_ @_ @x sh1 (AstRaw a) = AstRaw $+    let c = ConvCmp (ConvXX ConvXS)+                    (ConvNest (STKX sh1 (knownSTK @x)))+    in cAstConvert c a+  xnest @_ @_ @x sh1 (AstRaw a) = AstRaw $+    let c = ConvNest (STKX sh1 (knownSTK @x))+    in cAstConvert c a+  xunNestR (AstRaw a) = AstRaw $+    let c = ConvCmp ConvUnnest+                    (ConvXX ConvRX)+    in cAstConvert c a+  xunNestS (AstRaw a) = AstRaw $+    let c = ConvCmp ConvUnnest+                    (ConvXX ConvSX)+    in cAstConvert c a+  xunNest (AstRaw a) = AstRaw $+    let c = ConvUnnest+    in cAstConvert c a++  tpairConv = tpair+  tunpairConv = tunpair++-- All but the last case are shortcuts for common forms.+astConcreteRaw :: FullShapeTK y -> Concrete y+               -> AstRaw PrimalSpan y+astConcreteRaw ftk v = case ftk of+  FTKScalar -> AstRaw $ AstConcreteK $ unConcrete v+  FTKR sh' FTKScalar -> AstRaw $+    withShsFromShR sh' $ \(sh :: ShS sh) ->+      withKnownShS sh $+      cAstFromS ftk $ AstConcreteS (unConcrete $ sfromR @_ @sh v)+  FTKS _ FTKScalar -> AstRaw $ AstConcreteS $ unConcrete v+  FTKX sh' FTKScalar -> AstRaw $+    withShsFromShX sh' $ \(sh :: ShS sh) ->+      withKnownShS sh $+      cAstFromS ftk $ AstConcreteS (unConcrete $ sfromX @_ @sh v)+  FTKProduct ftk1 ftk2 -> AstRaw $+    AstPair (unAstRaw $ astConcreteRaw ftk1 (tproject1 v))+            (unAstRaw $ astConcreteRaw ftk2 (tproject2 v))+  _ -> concreteTarget (tkconcrete . unConcrete) (tsconcrete . unConcrete)+                      (\ftk2 a -> AstRaw $ cAstFromS ftk2 $ unAstRaw a)+                      (ftkToSTK ftk) v+++-- * AstNoVectorize instances++instance AstSpan s => LetTensor (AstNoVectorize s) where+  ttlet u f = AstNoVectorize+              $ ttlet (unAstNoVectorize u)+                      (unAstNoVectorize . f . AstNoVectorize)+  ttletPrimal u f = AstNoVectorize+                    $ ttletPrimal (unAstNoVectorize u)+                                  (unAstNoVectorize . f . AstNoVectorize)+  toShare t = toShare $ unAstNoVectorize t++instance AstSpan s => BaseTensor (AstNoVectorize s) where+  -- Ranked ops+  rshape = rshape . unAstNoVectorize+  trsum = AstNoVectorize . trsum . unAstNoVectorize+  trreplicate k = AstNoVectorize . trreplicate k . unAstNoVectorize+  trindex v ix =+    AstNoVectorize $ trindex (unAstNoVectorize v) (unAstNoVectorize <$> ix)+  trscatter sh t f =+    AstNoVectorize $ trscatter sh (unAstNoVectorize t)+                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize+  trgather sh t f =+    AstNoVectorize $ trgather sh (unAstNoVectorize t)+                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize+  trconcrete = AstNoVectorize . trconcrete+  trfloor = AstNoVectorize . trfloor . unAstNoVectorize+  trfromIntegral = AstNoVectorize . trfromIntegral . unAstNoVectorize+  trcast = AstNoVectorize . trcast . unAstNoVectorize+  trminIndex = AstNoVectorize . trminIndex . unAstNoVectorize+  trmaxIndex = AstNoVectorize . trmaxIndex . unAstNoVectorize+  triota = AstNoVectorize . triota+  trappend u v =+    AstNoVectorize $ trappend (unAstNoVectorize u) (unAstNoVectorize v)+  trslice i n = AstNoVectorize . trslice i n . unAstNoVectorize+  trreverse = AstNoVectorize . trreverse . unAstNoVectorize+  trtranspose perm = AstNoVectorize . trtranspose perm . unAstNoVectorize+  trreshape sh = AstNoVectorize . trreshape sh . unAstNoVectorize+  trbuild1 k f = withSNat k $ \snat ->+    AstNoVectorize $ AstBuild1 snat knownSTK+    $ funToAstI (Just (0, fromIntegral k - 1))+        -- this introduces new variable names+    $ unAstNoVectorize . f . AstNoVectorize++  -- Shaped ops+  sshape = sshape . unAstNoVectorize+  tssum = AstNoVectorize . tssum . unAstNoVectorize+  tsindex v ix =+    AstNoVectorize $ tsindex (unAstNoVectorize v) (unAstNoVectorize <$> ix)+  tsscatter @_ @shm @shn @shp t f =+    AstNoVectorize $ tsscatter @_ @_ @shm @shn @shp (unAstNoVectorize t)+                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize+  tsgather @_ @shm @shn @shp t f =+    AstNoVectorize $ tsgather @_ @_ @shm @shn @shp (unAstNoVectorize t)+                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize+  tsconcrete = AstNoVectorize . tsconcrete+  tsfloor = AstNoVectorize . tsfloor . unAstNoVectorize+  tsfromIntegral = AstNoVectorize . tsfromIntegral . unAstNoVectorize+  tscast = AstNoVectorize . tscast . unAstNoVectorize+  tsminIndex = AstNoVectorize . tsminIndex . unAstNoVectorize+  tsmaxIndex = AstNoVectorize . tsmaxIndex . unAstNoVectorize+  tsiota = AstNoVectorize tsiota+  tsappend u v =+    AstNoVectorize $ tsappend (unAstNoVectorize u) (unAstNoVectorize v)+  tsslice i n k = AstNoVectorize . tsslice i n k . unAstNoVectorize+  tsreverse = AstNoVectorize . tsreverse . unAstNoVectorize+  tsbuild1 @k f = AstNoVectorize $ AstBuild1 (SNat @k) knownSTK+                  $ funToAstI (Just (0, valueOf @k - 1))+                      -- this introduces new variable names+                  $ unAstNoVectorize . f . AstNoVectorize++  -- Mixed ops+  xshape = xshape . unAstNoVectorize+  txsum = AstNoVectorize . txsum . unAstNoVectorize+  txreplicate snat sh = AstNoVectorize . txreplicate snat sh . unAstNoVectorize+  txindex v ix =+    AstNoVectorize $ txindex (unAstNoVectorize v) (unAstNoVectorize <$> ix)+  txscatter @_ @shm @shn @shp sh t f =+    AstNoVectorize $ txscatter @_ @_ @shm @shn @shp sh (unAstNoVectorize t)+                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize+  txgather @_ @shm @shn @shp sh t f =+    AstNoVectorize $ txgather @_ @_ @shm @shn @shp sh (unAstNoVectorize t)+                   $ fmap unAstNoVectorize . f . fmap AstNoVectorize+  txconcrete = AstNoVectorize . txconcrete+  txfloor = AstNoVectorize . txfloor . unAstNoVectorize+  txfromIntegral = AstNoVectorize . txfromIntegral . unAstNoVectorize+  txcast = AstNoVectorize . txcast . unAstNoVectorize+  txminIndex = AstNoVectorize . txminIndex . unAstNoVectorize+  txmaxIndex = AstNoVectorize . txmaxIndex . unAstNoVectorize+  txiota @n = AstNoVectorize $ txiota @_ @n+  txappend u v =+    AstNoVectorize $ txappend (unAstNoVectorize u) (unAstNoVectorize v)+  txslice i n k = AstNoVectorize . txslice i n k . unAstNoVectorize+  txreverse = AstNoVectorize . txreverse . unAstNoVectorize+  txtranspose perm = AstNoVectorize . txtranspose perm . unAstNoVectorize+  txreshape sh = AstNoVectorize . txreshape sh . unAstNoVectorize+  txbuild1 @k f = AstNoVectorize $ AstBuild1 (SNat @k) knownSTK+                  $ funToAstI (Just (0, valueOf @k - 1))+                      -- this introduces new variable names+                  $ unAstNoVectorize . f . AstNoVectorize++  -- Scalar ops+  tkconcrete = AstNoVectorize . tkconcrete+  tkfloor = AstNoVectorize . tkfloor . unAstNoVectorize+  tkfromIntegral = AstNoVectorize . tkfromIntegral . unAstNoVectorize+  tkcast = AstNoVectorize . tkcast . unAstNoVectorize++  -- General operations that don't require LetTensor nor ShareTensor+  tftk stk = tftk stk . unAstNoVectorize+  tconcrete ftk a = AstNoVectorize $ tconcrete ftk a+  tpair t1 t2 =+    AstNoVectorize $ tpair (unAstNoVectorize t1) (unAstNoVectorize t2)+  tproject1 t = AstNoVectorize $ tproject1 $ unAstNoVectorize t+  tproject2 t = AstNoVectorize $ tproject2 $ unAstNoVectorize t+  tsreplicate snat sh = AstNoVectorize . tsreplicate snat sh. unAstNoVectorize+  tstranspose perm =+    AstNoVectorize . tstranspose perm . unAstNoVectorize+  tsreshape sh = AstNoVectorize . tsreshape sh . unAstNoVectorize+  tmapAccumRDer _ !k !accftk !bftk !eftk f df rf acc0 es =+    AstNoVectorize $ tmapAccumRDer Proxy k accftk bftk eftk f df rf+                       (unAstNoVectorize acc0) (unAstNoVectorize es)+  tmapAccumLDer _ !k !accftk !bftk !eftk f df rf acc0 es =+    AstNoVectorize $ tmapAccumLDer Proxy k accftk bftk eftk f df rf+                       (unAstNoVectorize acc0) (unAstNoVectorize es)+  tApply t ll = AstNoVectorize $ tApply t (unAstNoVectorize ll)+  tlambda = tlambda @(AstTensor AstMethodLet s)+  tcond !stk !b !u !v =+    AstNoVectorize $ tcond stk b (unAstNoVectorize u) (unAstNoVectorize v)+  tprimalPart t = AstNoVectorize $ tprimalPart $ unAstNoVectorize t+  tdualPart stk t = tdualPart stk $ unAstNoVectorize t+  tfromPrimal stk t = AstNoVectorize $ tfromPrimal stk $ unAstNoVectorize t+  tfromDual t = AstNoVectorize $ tfromDual t+  tgrad = tgrad @(AstTensor AstMethodLet s)+  tvjp = tvjp @(AstTensor AstMethodLet s)+  tjvp = tjvp @(AstTensor AstMethodLet s)++  tfromVector k stk =+    AstNoVectorize . tfromVector k stk . V.map unAstNoVectorize+  tsum k stk =+    AstNoVectorize . tsum k stk . unAstNoVectorize+  treplicate k stk =+    AstNoVectorize . treplicate k stk . unAstNoVectorize+  tindexBuild k stk u i =+    AstNoVectorize $ tindexBuild k stk (unAstNoVectorize u) (unAstNoVectorize i)++  treplTarget r ftk = AstNoVectorize $ treplTarget r ftk+  tdefTarget = AstNoVectorize . tdefTarget+  taddTarget stk a b = AstNoVectorize $ taddTarget stk (unAstNoVectorize a)+                                                       (unAstNoVectorize b)+  tmultTarget stk a b = AstNoVectorize $ tmultTarget stk (unAstNoVectorize a)+                                                         (unAstNoVectorize b)+  tsum0Target stk a = AstNoVectorize $ tsum0Target stk (unAstNoVectorize a)+  tdot0Target stk a b = AstNoVectorize $ tdot0Target stk (unAstNoVectorize a)+                                                         (unAstNoVectorize b)++instance AstSpan s => ConvertTensor (AstNoVectorize s) where+  tconvert c _astk = AstNoVectorize . astConvert c . unAstNoVectorize++  rfromX = AstNoVectorize . rfromX . unAstNoVectorize+  xfromR = AstNoVectorize . xfromR . unAstNoVectorize++  sfromR = AstNoVectorize . sfromR . unAstNoVectorize+  sfromX = AstNoVectorize . sfromX . unAstNoVectorize+  xfromS = AstNoVectorize . xfromS . unAstNoVectorize++  rzip = AstNoVectorize . rzip . unAstNoVectorize+  runzip = AstNoVectorize . runzip . unAstNoVectorize+  szip = AstNoVectorize . szip . unAstNoVectorize+  sunzip = AstNoVectorize . sunzip . unAstNoVectorize+  xzip = AstNoVectorize . xzip . unAstNoVectorize+  xunzip = AstNoVectorize . xunzip . unAstNoVectorize++  xnestR sh = AstNoVectorize . xnestR sh . unAstNoVectorize+  xnestS sh = AstNoVectorize . xnestS sh . unAstNoVectorize+  xnest sh = AstNoVectorize . xnest sh . unAstNoVectorize+  xunNestR = AstNoVectorize . xunNestR . unAstNoVectorize+  xunNestS = AstNoVectorize . xunNestS . unAstNoVectorize+  xunNest = AstNoVectorize . xunNest . unAstNoVectorize++  tpairConv = tpair+  tunpairConv a = let (b, c) = tunpairConv $ unAstNoVectorize a+                  in (AstNoVectorize b, AstNoVectorize c)+++-- * AstNoSimplify instances++instance AstSpan s => LetTensor (AstNoSimplify s) where+  ttlet u f = AstNoSimplify+              $ astLetFunNoSimplify (unAstNoSimplify u)+                                    (unAstNoSimplify . f . AstNoSimplify)+  ttletPrimal u f = AstNoSimplify+                    $ astLetFunNoSimplify (unAstNoSimplify u)+                                          (unAstNoSimplify . f . AstNoSimplify)+  toShare t = AstRaw $ AstToShare $ unAstNoSimplify t++wAstNoSimplify :: AstRaw s y -> AstNoSimplify s y+wAstNoSimplify =+  AstNoSimplify+  . (unsafeCoerce :: AstTensor AstMethodShare s y+                  -> AstTensor AstMethodLet s y)+  . unAstRaw++wunAstNoSimplify :: AstNoSimplify s y -> AstRaw s y+wunAstNoSimplify =+  AstRaw+  . (unsafeCoerce :: AstTensor AstMethodLet s y+                  -> AstTensor AstMethodShare s y)+  . unAstNoSimplify++instance AstSpan s => BaseTensor (AstNoSimplify s) where+  -- The implementation of these methods differs from the AstRaw instance:+  trbuild1 @n @x k f = withSNat k $ \snat ->+    AstNoSimplify+    $ astBuild1Vectorize snat (STKR (SNat @n) (knownSTK @x))+                         (unAstNoSimplify . f . AstNoSimplify)+  tsbuild1 @k @sh @x f =+    AstNoSimplify+    $ astBuild1Vectorize (SNat @k) (STKS (knownShS @sh) (knownSTK @x))+                         (unAstNoSimplify . f . AstNoSimplify)+  txbuild1 @k @sh @x f =+    AstNoSimplify+    $ astBuild1Vectorize (SNat @k) (STKX (knownShX @sh) (knownSTK @x))+                         (unAstNoSimplify . f . AstNoSimplify)+  -- These three have tricky types, so we repaat the AstRaw definitions:+  tcond _ !b !u !v =+    AstNoSimplify $ AstCond b (unAstNoSimplify u) (unAstNoSimplify v)+  tdualPart _ t = dualPart $ unAstNoSimplify t+  tfromDual t = AstNoSimplify $ fromDual t++  -- All the following implementations piggy-back on AstRaw implementations.+  -- Ranked ops+  rshape = rshape . wunAstNoSimplify+  trsum = wAstNoSimplify . trsum . wunAstNoSimplify+  trreplicate k = wAstNoSimplify . trreplicate k . wunAstNoSimplify+  trindex v ix =+    wAstNoSimplify $ trindex (wunAstNoSimplify v) (wunAstNoSimplify <$> ix)+  trscatter sh t f =+    wAstNoSimplify $ trscatter sh (wunAstNoSimplify t)+                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify+  trgather sh t f =+    wAstNoSimplify $ trgather sh (wunAstNoSimplify t)+                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify+  trconcrete = wAstNoSimplify . trconcrete+  trfloor = wAstNoSimplify . trfloor . wunAstNoSimplify+  trfromIntegral = wAstNoSimplify . trfromIntegral . wunAstNoSimplify+  trcast = wAstNoSimplify . trcast . wunAstNoSimplify+  trminIndex = wAstNoSimplify . trminIndex . wunAstNoSimplify+  trmaxIndex = wAstNoSimplify . trmaxIndex . wunAstNoSimplify+  triota = wAstNoSimplify . triota+  trappend u v =+    wAstNoSimplify $ trappend (wunAstNoSimplify u) (wunAstNoSimplify v)+  trslice i n = wAstNoSimplify . trslice i n . wunAstNoSimplify+  trreverse = wAstNoSimplify . trreverse . wunAstNoSimplify+  trtranspose perm = wAstNoSimplify . trtranspose perm . wunAstNoSimplify+  trreshape sh = wAstNoSimplify . trreshape sh . wunAstNoSimplify++  -- Shaped ops+  sshape = sshape . wunAstNoSimplify+  tssum = wAstNoSimplify . tssum . wunAstNoSimplify+  tsindex v ix =+    wAstNoSimplify $ tsindex (wunAstNoSimplify v) (wunAstNoSimplify <$> ix)+  tsscatter @_ @shm @shn @shp t f =+    wAstNoSimplify $ tsscatter @_ @_ @shm @shn @shp (wunAstNoSimplify t)+                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify+  tsgather @_ @shm @shn @shp t f =+    wAstNoSimplify $ tsgather @_ @_ @shm @shn @shp (wunAstNoSimplify t)+                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify+  tsconcrete = wAstNoSimplify . tsconcrete+  tsfloor = wAstNoSimplify . tsfloor . wunAstNoSimplify+  tsfromIntegral = wAstNoSimplify . tsfromIntegral . wunAstNoSimplify+  tscast = wAstNoSimplify . tscast . wunAstNoSimplify+  tsminIndex = wAstNoSimplify . tsminIndex . wunAstNoSimplify+  tsmaxIndex = wAstNoSimplify . tsmaxIndex . wunAstNoSimplify+  tsiota = wAstNoSimplify tsiota+  tsappend u v =+    wAstNoSimplify $ tsappend (wunAstNoSimplify u) (wunAstNoSimplify v)+  tsslice i n k = wAstNoSimplify . tsslice i n k . wunAstNoSimplify+  tsreverse = wAstNoSimplify . tsreverse . wunAstNoSimplify++  -- Mixed ops+  xshape = xshape . wunAstNoSimplify+  txsum = wAstNoSimplify . txsum . wunAstNoSimplify+  txreplicate snat sh = wAstNoSimplify . txreplicate snat sh . wunAstNoSimplify+  txindex v ix =+    wAstNoSimplify $ txindex (wunAstNoSimplify v) (wunAstNoSimplify <$> ix)+  txscatter @_ @shm @shn @shp sh t f =+    wAstNoSimplify $ txscatter @_ @_ @shm @shn @shp sh (wunAstNoSimplify t)+                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify+  txgather @_ @shm @shn @shp sh t f =+    wAstNoSimplify $ txgather @_ @_ @shm @shn @shp sh (wunAstNoSimplify t)+                   $ fmap wunAstNoSimplify . f . fmap wAstNoSimplify+  txconcrete = wAstNoSimplify . txconcrete+  txfloor = wAstNoSimplify . txfloor . wunAstNoSimplify+  txfromIntegral = wAstNoSimplify . txfromIntegral . wunAstNoSimplify+  txcast = wAstNoSimplify . txcast . wunAstNoSimplify+  txminIndex = wAstNoSimplify . txminIndex . wunAstNoSimplify+  txmaxIndex = wAstNoSimplify . txmaxIndex . wunAstNoSimplify+  txiota @n = wAstNoSimplify $ txiota @_ @n+  txappend u v =+    wAstNoSimplify $ txappend (wunAstNoSimplify u) (wunAstNoSimplify v)+  txslice i n k = wAstNoSimplify . txslice i n k . wunAstNoSimplify+  txreverse = wAstNoSimplify . txreverse . wunAstNoSimplify+  txtranspose perm = wAstNoSimplify . txtranspose perm . wunAstNoSimplify+  txreshape sh = wAstNoSimplify . txreshape sh . wunAstNoSimplify++  -- Scalar ops+  tkconcrete = wAstNoSimplify . tkconcrete+  tkfloor = wAstNoSimplify . tkfloor . wunAstNoSimplify+  tkfromIntegral = wAstNoSimplify . tkfromIntegral . wunAstNoSimplify+  tkcast = wAstNoSimplify . tkcast . wunAstNoSimplify++  -- General operations that don't require LetTensor nor ShareTensor+  tftk stk = tftk stk . wunAstNoSimplify+  tconcrete ftk a = wAstNoSimplify $ tconcrete ftk a+  tpair t1 t2 =+    wAstNoSimplify $ tpair (wunAstNoSimplify t1) (wunAstNoSimplify t2)+  tproject1 t = wAstNoSimplify $ tproject1 $ wunAstNoSimplify t+  tproject2 t = wAstNoSimplify $ tproject2 $ wunAstNoSimplify t+  tsreplicate snat sh = wAstNoSimplify . tsreplicate snat sh . wunAstNoSimplify+  tstranspose perm =+    wAstNoSimplify . tstranspose perm . wunAstNoSimplify+  tsreshape sh = wAstNoSimplify . tsreshape sh . wunAstNoSimplify+  tmapAccumRDer _ !k !accftk !bftk !eftk f df rf acc0 es =+    wAstNoSimplify $ tmapAccumRDer Proxy k accftk bftk eftk f df rf+                       (wunAstNoSimplify acc0) (wunAstNoSimplify es)+  tmapAccumLDer _ !k !accftk !bftk !eftk f df rf acc0 es =+    wAstNoSimplify $ tmapAccumLDer Proxy k accftk bftk eftk f df rf+                       (wunAstNoSimplify acc0) (wunAstNoSimplify es)+  tApply t ll = wAstNoSimplify $ tApply t (wunAstNoSimplify ll)+  tlambda = tlambda @(AstRaw s)+  tprimalPart t = wAstNoSimplify $ tprimalPart $ wunAstNoSimplify t+  tfromPrimal stk t = wAstNoSimplify $ tfromPrimal stk $ wunAstNoSimplify t+  tgrad = tgrad @(AstRaw s)+  tvjp = tvjp @(AstRaw s)+  tjvp = tjvp @(AstRaw s)++  tfromVector k stk =+    wAstNoSimplify . tfromVector k stk . V.map wunAstNoSimplify+  tsum snat@SNat stk u = case stk of+    STKScalar -> kfromS $ tssum u+    STKR SNat x | Dict <- lemKnownSTK x -> trsum u+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (tsum snat stk1 (tproject1 u3))+              (tsum snat stk2 (tproject2 u3))+  treplicate snat@SNat stk u = case stk of+    STKScalar -> tsreplicate snat ZSS $ sfromK u+    STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (sNatValue snat) u+    STKS sh x | Dict <- lemKnownSTK x -> tsreplicate snat sh u+    STKX sh x | Dict <- lemKnownSTK x -> txreplicate snat sh u+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (treplicate snat stk1 (tproject1 u3))+              (treplicate snat stk2 (tproject2 u3))+  tindexBuild snat@SNat stk u i = case stk of+    STKScalar -> kfromS $ tsindex u (i :.$ ZIS)+    STKR SNat x | Dict <- lemKnownSTK x -> trindex u (i :.: ZIR)+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsindex u (i :.$ ZIS)+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txindex u (i :.% ZIX)+    STKProduct stk1 stk2 ->+      ttlet u $ \ !u3 ->+        tpair (tindexBuild snat stk1 (tproject1 u3) i)+              (tindexBuild snat stk2 (tproject2 u3) i)++  treplTarget r ftk = wAstNoSimplify $ treplTarget r ftk+  tdefTarget = wAstNoSimplify . tdefTarget+  taddTarget stk a b = wAstNoSimplify $ taddTarget stk (wunAstNoSimplify a)+                                                       (wunAstNoSimplify b)+  tmultTarget stk a b = wAstNoSimplify $ tmultTarget stk (wunAstNoSimplify a)+                                                         (wunAstNoSimplify b)+  tsum0Target stk a = wAstNoSimplify $ tsum0Target stk (wunAstNoSimplify a)+  tdot0Target stk a b = wAstNoSimplify $ tdot0Target stk (wunAstNoSimplify a)+                                                         (wunAstNoSimplify b)++instance AstSpan s => ConvertTensor (AstNoSimplify s) where+  tconvert c astk = wAstNoSimplify . tconvert c astk . wunAstNoSimplify++  rfromX = wAstNoSimplify . rfromX . wunAstNoSimplify+  xfromR = wAstNoSimplify . xfromR . wunAstNoSimplify++  sfromR = wAstNoSimplify . sfromR . wunAstNoSimplify+  sfromX = wAstNoSimplify . sfromX . wunAstNoSimplify+  xfromS = wAstNoSimplify . xfromS . wunAstNoSimplify++  rzip = wAstNoSimplify . rzip . wunAstNoSimplify+  runzip = wAstNoSimplify . runzip . wunAstNoSimplify+  szip = wAstNoSimplify . szip . wunAstNoSimplify+  sunzip = wAstNoSimplify . sunzip . wunAstNoSimplify+  xzip = wAstNoSimplify . xzip . wunAstNoSimplify+  xunzip = wAstNoSimplify . xunzip . wunAstNoSimplify++  xnestR sh = wAstNoSimplify . xnestR sh . wunAstNoSimplify+  xnestS sh = wAstNoSimplify . xnestS sh . wunAstNoSimplify+  xnest sh = wAstNoSimplify . xnest sh . wunAstNoSimplify+  xunNestR = wAstNoSimplify . xunNestR . wunAstNoSimplify+  xunNestS = wAstNoSimplify . xunNestS . wunAstNoSimplify+  xunNest = wAstNoSimplify . xunNest . wunAstNoSimplify++  tpairConv = tpair+  tunpairConv (AstNoSimplify (AstPair t1 t2)) =  -- a tiny bit of simplification+    (AstNoSimplify t1, AstNoSimplify t2)+  tunpairConv t = (tproject1 t, tproject2 t)
+ src/HordeAd/Core/OpsConcrete.hs view
@@ -0,0 +1,1315 @@+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | Tensor class instances for concrete arrays backed+-- by 'Data.Vector.Storable.Vector'+-- and defined in @ox-arrays@ on the basis of @orthotope@.+module HordeAd.Core.OpsConcrete+  () where++import Prelude hiding (foldl')++import Control.Arrow (second)+import Control.Exception.Assert.Sugar+import Data.Coerce (Coercible, coerce)+import Data.Foldable qualified as Foldable+import Data.Function ((&))+import Data.Int (Int64)+import Data.List (foldl', mapAccumL, mapAccumR)+import Data.List.Index (imap)+import Data.List.NonEmpty qualified as NonEmpty+import Data.Map.Strict qualified as M+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import Data.Vector.Generic qualified as V+import Data.Vector.Storable qualified as VS+import Data.Vector.Strict qualified as Data.Vector+import GHC.Exts (IsList (..))+import GHC.TypeLits (KnownNat, sameNat, type (+))++import Data.Array.Nested (MapJust, Replicate, type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed qualified as Mixed+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked qualified as Ranked+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped qualified as Shaped+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (Init, unsafeCoerceRefl)+import Data.Array.Strided.Orthotope (liftVEltwise1)++import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.OpsADVal+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind++-- * Tensor classes instance++instance LetTensor Concrete where+  ttlet = (&)+  ttletPrimal = (&)+  toShare = id+  tunshare = id+  tD _stk t DummyDualTarget{} = t+  tfold k _ stk f x0 as = foldl' f x0 (tunravelToListShare k stk as)+  tscan k@(SNat @k) nstk stk f x0 as =+    tfromVector (SNat @(1 + k)) nstk+    $ V.scanl' f x0 (V.fromList $ tunravelToListShare k stk as)++instance ShareTensor Concrete where+  tshare = id+  tunpair (Concrete (t1, t2)) = (Concrete t1, Concrete t2)++instance BaseTensor Concrete where+  -- Ranked ops+  rshape @_ @r | Dict <- eltDictRep (knownSTK @r) = Nested.rshape . unConcrete+  trfromVector @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.rfromListOuter . NonEmpty.fromList . V.toList+    . fmapUnConcrete+  trfromVector0N @_ @r sh | Dict <- eltDictRep (knownSTK @r) =+    Concrete . tfromVector0NR sh . fmapUnConcrete+  trunravelToList @_ @r | Dict <- eltDictRep (knownSTK @r) =+    fmapConcrete . Nested.rtoListOuter . unConcrete+  trsum t = case tftk knownSTK t of+    FTKR _ FTKScalar ->  -- optimized+      Concrete . Nested.rsumOuter1 . unConcrete $ t+    FTKR _ x ->+      let l = trunravelToList t+          sh = shrTail $ rshape t+      in foldr (taddTarget knownSTK) (tdefTarget (FTKR sh x)) l+        -- Concrete has a ShareTensor instance, so taddTarget arguments+        -- don't need to be duplicable+  trsum0 @_ @r t = case knownSTK @r of+    STKScalar ->  -- optimized+      Concrete . Nested.rscalar . Nested.rsumAllPrim . unConcrete $ t+    _ -> trsum . rflatten $ t+  {-# INLINE trdot0 #-}+  trdot0 u v =+    Concrete $ Nested.rscalar $ Nested.rdot (unConcrete u) (unConcrete v)+  trdot1In u v = Concrete $ Nested.rdot1Inner (unConcrete u) (unConcrete v)+  {-# INLINE trmatvecmul #-}+  trmatvecmul m v = trdot1In m (trreplicate (rwidth m) v)+  trmatmul2 m1 m2 = case rshape m2 of+    _ :$: width2 :$: ZSR ->+      trdot1In (trtranspose [1, 0] (trreplicate width2 m1))+               (trtranspose [0, 2, 1] (trreplicate (rwidth m1) m2))+  trreplicate @_ @r k | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.rreplicate (k :$: ZSR) . unConcrete+  trreplicate0N @_ @r sh | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.rreplicate sh . unConcrete+  trindex = tindexZR+  trindex0 = tindex0R+  trscatter = tscatterZR+  trscatter1 = tscatterZ1R+  trgather = tgatherZR+  trgather1 = tgatherZ1R+  trconcrete = Concrete+  trfloor = Concrete . liftVR (V.map floor) . unConcrete+  trfromIntegral = Concrete . liftVR (V.map fromIntegral) . unConcrete+  {-# INLINE trcast #-}+  trcast = Concrete . liftVR (V.map realToFrac) . unConcrete+  trminIndex = Concrete . tminIndexR . unConcrete+  trmaxIndex = Concrete . tmaxIndexR . unConcrete+  triota n = Concrete $ Nested.rfromList1 $ NonEmpty.map fromInteger+             $ NonEmpty.fromList [0 .. fromIntegral n - 1]+  trappend @_ @r u v | Dict <- eltDictRep (knownSTK @r) =+    Concrete $ Nested.rappend (unConcrete u) (unConcrete v)+  trslice @_ @r i n | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.rslice i n . unConcrete+  trreverse @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.rrev1 . unConcrete+  trtranspose @_ @r perm | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.rtranspose perm . unConcrete+  trreshape @_ @_ @r sh | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.rreshape sh . unConcrete+  trbuild1 @_ @r k f | Dict <- eltDictRep (knownSTK @r) =+    Concrete $ tbuild1R k (unConcrete . f . Concrete)+  trmap0N @_ @r @r1 f t = case (knownSTK @r1, knownSTK @r) of+    (STKScalar, STKScalar) ->+      Concrete $ tmap0NR (unConcrete . f . Concrete) (unConcrete t)+    _ ->  -- this is the default implementation from the class+      rbuild (rshape t) (f . trindex t)+  trzipWith0N @_ @r1 @r2 @r f t u =+    case (knownSTK @r1, knownSTK @r2, knownSTK @r) of+      (STKScalar, STKScalar, STKScalar) ->+        Concrete+        $ tzipWith0NR (\v w -> unConcrete $ f (Concrete v) (Concrete w))+                      (unConcrete t) (unConcrete u)+      _ ->  -- this is the default implementation from the class+        rbuild (rshape u) (\ix -> f (trindex t ix) (trindex u ix))++  -- Shaped ops+  sshape @_ @r | Dict <- eltDictRep (knownSTK @r) = Nested.sshape . unConcrete+  tsfromVector @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.sfromListOuter SNat . NonEmpty.fromList . V.toList+    . fmapUnConcrete+  tsfromVector0N @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . tfromVector0NS . fmapUnConcrete+  tsunravelToList @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =+    fmapConcrete . Nested.stoListOuter . unConcrete+  tssum t = case tftk knownSTK t of+    FTKS _ FTKScalar ->  -- optimized+      Concrete . Nested.ssumOuter1 . unConcrete $ t+    FTKS _ x ->+      let l = tsunravelToList t+          sh = shsTail $ sshape t+      in foldr (taddTarget knownSTK) (tdefTarget (FTKS sh x)) l+  tssum0 @sh @r t | SNat <- shsProduct (knownShS @sh) = case knownSTK @r of+    STKScalar ->  -- optimized+      Concrete . Nested.sscalar . Nested.ssumAllPrim . unConcrete $ t+    _ -> tssum . sflatten $ t+  {-# INLINE tsdot0 #-}  -- this doesn't want to specialize+  tsdot0 u v  =+    Concrete $ Nested.sscalar $ Nested.sdot (unConcrete u) (unConcrete v)+  tsdot1In @_ (SNat @n) u v =+    Concrete $ Nested.sdot1Inner (Proxy @n) (unConcrete u) (unConcrete v)+  {-# INLINE tsmatvecmul #-}  -- this doesn't want to specialize+  tsmatvecmul m v = tsdot1In SNat m (tsreplicate SNat knownShS v)+  tsmatmul2 m1 m2 =+    tsdot1In SNat+             (tstranspose (Permutation.makePerm @'[1, 0])+                          (tsreplicate SNat knownShS m1))+             (tstranspose (Permutation.makePerm @'[0, 2, 1])+                          (tsreplicate SNat knownShS m2))+  tsindex = tindexZS+  tsindex0 = tindex0S+  -- Performance depends a lot on the number and size of tensors.+  -- If tensors are not tiny, memory taken by underlying vectors matters most+  -- and this implementation is probbaly optimal in this respect+  -- (the only new vectors are created by V.concat, but this is done on demand).+  -- TODO: optimize updateNS and make it consume and forget arguments+  -- one by one to make the above true+  --+  -- Note how ix being in bounds is checked. The semantics of the operation+  -- permits index out of bounds and then no tensors is added at such an index.+  tsscatter @shm @shn @shp t f =+    let shpshn = knownShS @shp `shsAppend` knownShS @shn+    in withKnownShS (knownShS @shm `shsAppend` knownShS @shn) $+       case tftk knownSTK t of+         FTKS _ x@FTKScalar ->  -- optimized+           gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn)+                                          :~: shp) $+           gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn)+                                          :~: shn) $+           let zero = tdefTarget (FTKS shpshn x)+               shm = knownShS @shm+               s = shsSize shm+               g ix =+                 let ix2 = f $ fmapConcrete ix+                 in if ixInBounds (fmapUnConcrete $ toList ix2)+                                  (shsToList shpshn)+                    then M.insertWith (V.zipWith (+)) ix2+                           (Nested.stoVector+                            $ tindexNS @_ @shm @shn (unConcrete t) ix)+                    else id+               ivs = foldr g M.empty [ fromLinearIdxS fromIntegral shm+                                       $ fromIntegral i+                                     | i <- [0 .. s - 1] ]+           in withKnownShS shpshn $+              updateNS @(Rank shp) zero+              $ map (second $ Concrete . Nested.sfromVector (knownShS @shn))+              $ M.assocs ivs+         FTKS _ x | Dict <- eltDictRep (ftkToSTK x) ->+           gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn)+                                       :~: shp) $+           gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn)+                                       :~: shn) $+           let zero = tdefTarget (FTKS shpshn x)+               shm = knownShS @shm+               s = shsSize shm+               g ix =+                 let ix2 = f $ fmapConcrete ix+                 in if ixInBounds (fmapUnConcrete $ toList ix2)+                                  (shsToList shpshn)+                    then M.insertWith (taddTarget knownSTK) ix2+                           (Concrete+                            $ tindexNS @_ @shm @shn (unConcrete t) ix)+                    else id+               ivs = foldr g M.empty [ fromLinearIdxS fromIntegral shm+                                       $ fromIntegral i+                                     | i <- [0 .. s - 1] ]+           in withKnownShS shpshn $+              updateNS @(Rank shp) zero+              $ M.assocs ivs+  tsscatter1 = tscatterZ1S+  -- The semantics of the operation permits index out of bounds+  -- and the result of such indexing is def, which is 0.+  -- TODO: are bounds checked in the optimized case?+  -- The same question also elsewhere.+  tsgather @shm @shn @_ @r t f =+    gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $+    gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $+    case knownSTK @r of+      STKScalar ->  -- optimized+        let shm = knownShS @shm+            s = shsSize shm+            l = [ stoVector+                  $ tsindex @_ @_ @shn+                      t (f (fmapConcrete+                            $ fromLinearIdxS fromIntegral shm i))+                | i <- [0 .. fromIntegral s - 1] ]+        in Concrete+           $ Nested.sfromVector (knownShS @shm `shsAppend` knownShS @shn)+           $ V.concat l+      _ ->+        withKnownShS (knownShS @shm `shsAppend` knownShS @shn) $+        sbuild @(Rank shm) (\ix -> t `tsindex` f ix)+  tsgather1 = tgatherZ1S+  tsconcrete = Concrete+  tsfloor = Concrete . liftVS (V.map floor) . unConcrete+  tsfromIntegral = Concrete . tfromIntegralS . unConcrete+  {-# INLINE tscast #-}  -- this doesn't want to specialize+  tscast = Concrete . liftVS (V.map realToFrac) . unConcrete+  tsminIndex = Concrete . tminIndexS . unConcrete+  tsmaxIndex = Concrete . tmaxIndexS . unConcrete+  tsiota @n = case NonEmpty.nonEmpty [0 .. valueOf @n - 1] of+    Nothing -> case sameNat (Proxy @n) (Proxy @0) of+      Just Refl -> Concrete $ Nested.semptyArray ZSS+      Nothing -> error "siota: wrong rank"+    Just l -> Concrete $ Nested.sfromList1 SNat $ NonEmpty.map fromInteger l+  tsappend @_ @_ @_ @r u v | Dict <- eltDictRep (knownSTK @r) =+    Concrete $ Nested.sappend (unConcrete u) (unConcrete v)+  tsslice @_ @_ @_ @_ @r i n _ | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.sslice i n . unConcrete+  tsreverse @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.srev1 . unConcrete+  tsbuild1 @_ @_ @r f | Dict <- eltDictRep (knownSTK @r) =+    Concrete $ tbuild1S (unConcrete . f . Concrete)+  tsmap0N @sh @r @r1 f v = case (knownSTK @r1, knownSTK @r) of+    (STKScalar, STKScalar) ->+      Concrete $ tmap0NS (unConcrete . f . Concrete) (unConcrete v)+    _ | Refl <- lemAppNil @sh ->+      -- this is the default implementation from the class+      gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])+      $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)+      $ sbuild @(Rank sh) (f . tsindex v)+  tszipWith0N @sh @r1 @r2 @r f t u =+    case (knownSTK @r1, knownSTK @r2, knownSTK @r) of+      (STKScalar, STKScalar, STKScalar) ->+        Concrete+        $ tzipWith0NS (\v w -> unConcrete $ f (Concrete v) (Concrete w))+                      (unConcrete t) (unConcrete u)+      _ | Refl <- lemAppNil @sh ->+        -- this is the default implementation from the class+        gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])+        $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)+        $ sbuild @(Rank sh) (\ix -> f (tsindex t ix) (tsindex u ix))++  -- Mixed ops+  xshape @_ @r | Dict <- eltDictRep (knownSTK @r) = Nested.mshape . unConcrete+  txfromVector @n @sh @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.mcast (Nested.SKnown (SNat @n) :!% knownShX @sh)+    . Nested.mfromListOuter . NonEmpty.fromList . V.toList+    . fmapUnConcrete+  txfromVector0N @_ @r sh | Dict <- eltDictRep (knownSTK @r) =+    Concrete . tfromVector0NX sh . fmapUnConcrete+  txunravelToList @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =+    fmapConcrete . Nested.mtoListOuter . unConcrete+  txsum t = case tftk knownSTK t of+    FTKX _ FTKScalar ->  -- optimized+      Concrete . Nested.msumOuter1 . unConcrete $ t+    FTKX _ x ->+      let l = txunravelToList t+          sh = shxTail $ xshape t+      in foldr (taddTarget knownSTK) (tdefTarget (FTKX sh x)) l+  txsum0 @_ @r t =+   case knownSTK @r of+    STKScalar ->  -- optimized+      Concrete . Nested.mscalar . Nested.msumAllPrim . unConcrete $ t+    _ -> withSNat (shxSize $ xshape t) $ \snat ->+      txsum (xmcast (Nested.SKnown snat :!% ZKX) $ xflatten t)+  {-# INLINE txdot0 #-}+  txdot0 u v =+    Concrete $ Nested.mscalar $ Nested.mdot (unConcrete u) (unConcrete v)+  txdot1In @_ (SNat @n) u v =+    Concrete $ Nested.mdot1Inner (Proxy @(Just n)) (unConcrete u) (unConcrete v)+  txmatvecmul mm mn m v =+    withKnownShX (ssxFromShX $ mn :$% ZSX) $+    withKnownShX (ssxFromShX $ mm :$% mn :$% ZSX) $+    withSNat (fromSMayNat' mm) $ \(SNat @m) ->+    withSNat (fromSMayNat' mn) $ \(SNat @n) ->+      xmcast (ssxFromShX (mm :$% ZSX))+      $ txsum (xtr (txreplicate (SNat @m) knownShX+                      (xmcast (ssxFromShX (Nested.SKnown (SNat @n)+                                             :$% ZSX)) v)+                    * xmcast (ssxFromShX (Nested.SKnown (SNat @m)+                                            :$% Nested.SKnown (SNat @n)+                                            :$% ZSX)) m))+  {-# INLINE txmatvecmul #-}+  txmatmul2 m1 m2 =+    txdot1In SNat+             (txtranspose (Permutation.makePerm @'[1, 0])+                          (txreplicate SNat knownShX m1))+             (txtranspose (Permutation.makePerm @'[0, 2, 1])+                          (txreplicate SNat knownShX m2))+  txreplicate @_ @_ @r snat _sh | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.mreplicate (Nested.SKnown snat :$% ZSX) . unConcrete+  txreplicate0N @sh @r sh | Refl <- lemAppNil @sh+                          , Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.mreplicate sh . unConcrete+  txindex = tindexZX+  txindex0 = tindex0X+  txscatter @shm @shn @shp sh t f =+    withKnownShX (knownShX @shm `ssxAppend` knownShX @shn) $+    gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn) :~: shp) $+    gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn) :~: shn) $+    case tftk knownSTK t of+      FTKX _ x@FTKScalar ->  -- optimized+        let zero = tdefTarget (FTKX sh x)+            shm = shxTakeSSX (Proxy @shn) (knownShX @shm) (xshape t)+            shDropP = shxDropSSX (knownShX @shm) (xshape t)+            s = shxSize shm+            g ix =+              let ix2 = f $ fmapConcrete ix+              in if ixInBounds (fmapUnConcrete $ toList ix2) (shxToList sh)+                 then M.insertWith (V.zipWith (+)) ix2+                        (Nested.mtoVector+                         $ tindexNX @_ @shm @shn (unConcrete t) ix)+                 else id+            ivs = foldr g M.empty [ fromLinearIdxX fromIntegral shm+                                    $ fromIntegral i+                                  | i <- [0 .. s - 1] ]+        in updateNX @(Rank shp) zero+           $ map (second $ Concrete . Nested.mfromVector shDropP)+           $ M.assocs ivs+      FTKX _ x | Dict <- eltDictRep (ftkToSTK x) ->+        let zero = tdefTarget (FTKX sh x)+            shm = shxTakeSSX (Proxy @shn) (knownShX @shm) (xshape t)+            s = shxSize shm+            g ix =+              let ix2 = f $ fmapConcrete ix+              in if ixInBounds (fmapUnConcrete $ toList ix2) (shxToList sh)+                 then M.insertWith (taddTarget knownSTK) ix2+                        (Concrete+                         $ tindexNX @_ @shm @shn (unConcrete t) ix)+                 else id+            ivs = foldr g M.empty [ fromLinearIdxX fromIntegral shm+                                    $ fromIntegral i+                                  | i <- [0 .. s - 1] ]+        in updateNX @(Rank shp) zero+           $ M.assocs ivs+  txscatter1 = tscatterZ1X+  txgather @shm @shn @_ @r sh t f =+    gcastWith (unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) $+    gcastWith (unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) $+    case knownSTK @r of+      STKScalar ->  -- optimized+        let shm = shxTakeSSX (Proxy @shn) (knownShX @shm) sh+            s = shxSize shm+            l = [ xtoVector+                  $ txindex @_ @_ @shn+                      t (f (fmapConcrete+                            $ fromLinearIdxX fromIntegral shm i))+                | i <- [0 .. fromIntegral s - 1] ]+        in Concrete $ Nested.mfromVector sh $ V.concat l+      _ ->+        withKnownShX (ssxFromShX sh) $+        xbuild @(Rank shm) sh (\ix -> t `txindex` f ix)+  txgather1 = tgatherZ1X+  txconcrete = Concrete+  txfloor = Concrete . liftVX (V.map floor) . unConcrete+  txfromIntegral = Concrete . liftVX (V.map fromIntegral) . unConcrete+  {-# INLINE txcast #-}+  txcast = Concrete . liftVX (V.map realToFrac) . unConcrete+  txminIndex = Concrete . tminIndexX . unConcrete+  txmaxIndex = Concrete . tmaxIndexX . unConcrete+  txiota @n = let n = valueOf @n+                  t = Nested.mfromList1 $ NonEmpty.map fromInteger+                      $ NonEmpty.fromList [0 .. n - 1]+              in Concrete $ Nested.mcast (Nested.SKnown (SNat @n) :!% ZKX) t+  txappend @_ @_ @_ @r u v | Dict <- eltDictRep (knownSTK @r) =+    Concrete $ Nested.mappend (unConcrete u) (unConcrete v)+  txslice @_ @_ @_ @_ @r i n _ | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.mslice i n . unConcrete+  txreverse @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.mrev1 . unConcrete+  txtranspose @_ @_ @r perm | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.mtranspose perm . unConcrete+  txreshape @_ @_ @r sh | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.mreshape sh . unConcrete+  txbuild1 @_ @_ @r f | Dict <- eltDictRep (knownSTK @r) =+    Concrete $ tbuild1X (unConcrete . f . Concrete)++  -- Scalar ops+  tkconcrete = Concrete+  tkfloor = Concrete . floor . unConcrete+  tkfromIntegral = Concrete . fromIntegral . unConcrete+  tkcast = Concrete . realToFrac . unConcrete++  -- General operations that don't require LetTensor nor ShareTensor+  tftk stk (Concrete t) = tftkG stk t+  tconcrete _ = id+  tpair !u !v = Concrete (unConcrete u, unConcrete v)+  tproject1 = Concrete . fst . unConcrete+  tproject2 = Concrete . snd . unConcrete+  tsreplicate @_ @_ @x snat@SNat _sh | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.sreplicate (snat :$$ ZSS) . unConcrete+  tsreplicate0N @sh @r sh | Refl <- lemAppNil @sh+                          , Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.sreplicate sh . unConcrete+  tstranspose @_ @_ @r perm | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.stranspose perm . unConcrete+  tsreshape @_ @_ @x sh | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.sreshape sh . unConcrete+  -- The eta-expansion below is needed for typing.+  tmapAccumRDer _ k _ bftk eftk f _df _rf acc0 es =+    oRtmapAccumR k bftk eftk (\ (Concrete a) (Concrete b) ->+                                Concrete $ f (a, b)) acc0 es+  tmapAccumLDer _ k _ bftk eftk f _df _rf acc0 es =+    oRtmapAccumL k bftk eftk (\ (Concrete a) (Concrete b) ->+                                Concrete $ f (a, b)) acc0 es+  tApply f x = Concrete $ f $ unConcrete x+  tlambda _ f x = unConcrete $ unHFun f $ Concrete x+  tcond _ b u v = if b then u else v+  tprimalPart = id+  tdualPart stk t = DummyDualTarget (tftk stk t)+  tfromPrimal _ t = t+  tfromDual (DummyDualTarget ftk) = tdefTarget ftk+  tScale _ _ t = t+  -- The code for tvjp and tjvp in this instance is similar as for the+  -- ADVal ranked instance, because the type family instance is the same.+  tgrad @x xftk h =+    let rf :: RepConcrete x -> RepConcrete (ADTensorKind x)+        rf !a = unConcrete $ fst $ crevOnParams Nothing (unHFun h)+                                                 xftk (Concrete a)+    in rf+  tvjp @x @z xftk h =+    let rf :: RepConcrete (TKProduct (ADTensorKind z) x)+           -> RepConcrete (ADTensorKind x)+        rf !db_a = unConcrete $ fst+                   $ crevOnParams (Just $ Concrete $ fst db_a) (unHFun h)+                                   xftk (Concrete $ snd db_a)+    in rf+  tjvp @x @z xftk h =+    let df :: RepConcrete (TKProduct (ADTensorKind x) x)+           -> RepConcrete (ADTensorKind z)+        df !da_a = unConcrete $ fst+                   $ cfwdOnParams xftk (Concrete $ snd da_a)+                                   (unHFun h) (Concrete $ fst da_a)+    in df++  tfromVector snat@SNat stk v = assert (V.length v == sNatValue snat)+                                $ case stk of+    STKScalar -> tsfromVector $ V.map sfromK v+    STKR SNat x | Dict <- lemKnownSTK x -> trfromVector v+    STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsfromVector v+    STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txfromVector v+    STKProduct stk1 stk2 ->+      let (v1, v2) = V.unzip $ V.map tunpair v+      in tpair (tfromVector snat stk1 v1) (tfromVector snat stk2 v2)++  treplTarget = replTarget+  tdefTarget = defTarget+  taddTarget = addTarget+  tmultTarget = multTarget+  tsum0Target = sum0Target+  tdot0Target = dot0Target++instance ConvertTensor Concrete where+  tconvert c astk a | Dict <- eltDictRep astk+                    , Dict <- eltDictRep (convertSTK c astk) =+    Concrete $ Nested.convert (interpretTKConversion c) (unConcrete a)++  kfromR = Concrete . Nested.runScalar . unConcrete+  kfromS = Concrete . Nested.sunScalar . unConcrete+  kfromX = Concrete . Nested.munScalar . unConcrete+  rfromK = Concrete . Nested.rscalar . unConcrete+  rfromS @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.stoRanked . unConcrete+  {-# SPECIALIZE rfromS :: KnownShS sh => Concrete (TKS sh Double) -> Concrete (TKR (Rank sh) Double) #-}+  {-# SPECIALIZE rfromS :: KnownShS sh => Concrete (TKS sh Float) -> Concrete (TKR (Rank sh) Float) #-}+  rfromX @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.mtoRanked . unConcrete+  sfromK = Concrete . Nested.sscalar . unConcrete+  sfromR @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . flip Nested.rcastToShaped knownShS . unConcrete+  sfromX @_ @_ @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.mcastToShaped knownShS . unConcrete+  xfromK = Concrete . Nested.mscalar . unConcrete+  xfromR @sh @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.rcastToMixed (knownShX @sh) . unConcrete+  xfromS @_ @sh' @r | Dict <- eltDictRep (knownSTK @r) =+    Concrete . Nested.scastToMixed (knownShX @sh') . unConcrete++  rzip @y @z (Concrete (a, b)) | Dict <- eltDictRep (knownSTK @y)+                               , Dict <- eltDictRep (knownSTK @z) =+    Concrete $ Nested.rzip a b+  runzip a = let (!a1, !a2) = Nested.runzip $ unConcrete a+             in Concrete (a1, a2)+  szip @y @z  (Concrete (a, b)) | Dict <- eltDictRep (knownSTK @y)+                                , Dict <- eltDictRep (knownSTK @z) =+    Concrete $ Nested.szip a b+  sunzip a = let (!a1, !a2) = Nested.sunzip $ unConcrete a+             in Concrete (a1, a2)+  xzip @y @z  (Concrete (a, b)) | Dict <- eltDictRep (knownSTK @y)+                                , Dict <- eltDictRep (knownSTK @z) =+    Concrete $ Nested.mzip a b+  xunzip a = let (!a1, !a2) = Nested.munzip $ unConcrete a+             in Concrete (a1, a2)++  xnestR @sh1 @m @x sh | Dict <- eltDictRep (knownSTK @x)+                       , Refl <- lemRankReplicate (SNat @m) =+    Concrete+    . Nested.convert+        @(Nested.Mixed sh1 (Nested.Mixed (Replicate m Nothing) (RepConcrete x)))+        (Nested.ConvXX Nested.ConvXR)+    . Nested.mnest sh+    . unConcrete+  xnestS @sh1 @sh2 @x sh | Dict <- eltDictRep (knownSTK @x) =+    Concrete+    . Nested.convert+        @(Nested.Mixed sh1 (Nested.Mixed (MapJust sh2) (RepConcrete x)))+        (Nested.ConvXX Nested.ConvXS)+    . Nested.mnest sh+    . unConcrete+  xnest @_ @_ @x sh | Dict <- eltDictRep (knownSTK @x) =+    Concrete . Nested.mnest sh . unConcrete+  xunNestR @sh1 @m @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete+    . Nested.munNest+    . Nested.convert+        @(Nested.Mixed sh1 (Nested.Ranked m (RepConcrete x)))+        (Nested.ConvXX Nested.ConvRX)+    . unConcrete+  xunNestS @sh1 @sh2 @x | Dict <- eltDictRep (knownSTK @x) =+    Concrete+    . Nested.munNest+    . Nested.convert+        @(Nested.Mixed sh1 (Nested.Shaped sh2 (RepConcrete x)))+        (Nested.ConvXX Nested.ConvSX)+    . unConcrete+  xunNest = Concrete . Nested.munNest . unConcrete++  tpairConv = tpair+  tunpairConv = tunpair++interpretTKConversion :: TKConversion a b+                    -> Nested.Conversion (RepConcrete a) (RepConcrete b)+interpretTKConversion c0 = case c0 of+  ConvId -> Nested.ConvId+  ConvCmp c1 c2 -> Nested.ConvCmp (interpretTKConversion c1)+                                  (interpretTKConversion c2)+  ConvRX -> Nested.ConvRX+  ConvSX -> Nested.ConvSX+  ConvXR stk | Dict <- eltDictRep stk -> Nested.ConvXR+  ConvXS -> Nested.ConvXS+  ConvXS' (FTKS sh' ftk) | Dict <- eltDictRep (ftkToSTK ftk) ->+    Nested.ConvXS' sh'+  ConvXX' (FTKX shx ftk) | Dict <- eltDictRep (ftkToSTK ftk) ->+    Nested.ConvXX' (ssxFromShX shx)+  ConvRR c -> Nested.ConvRR (interpretTKConversion c)+  ConvSS c -> Nested.ConvSS (interpretTKConversion c)+  ConvXX c -> Nested.ConvXX (interpretTKConversion c)+  ConvT2 c1 c2 ->+    Nested.ConvT2 (interpretTKConversion c1) (interpretTKConversion c2)+  Conv0X stk | Dict <- eltDictRep stk -> Nested.Conv0X+  ConvX0 -> Nested.ConvX0+  ConvNest (STKX sh x) | Dict <- eltDictRep x -> Nested.ConvNest sh+  ConvUnnest -> Nested.ConvUnnest+  ConvZip stk1 stk2 | Dict <- eltDictRep stk1+                    , Dict <- eltDictRep stk2 -> Nested.ConvZip+  ConvUnzip stk1 stk2 | Dict <- eltDictRep stk1+                      , Dict <- eltDictRep stk2 -> Nested.ConvUnzip+++-- * MapAccum internal definitions++ravel :: forall k y.+         SNat k -> SingletonTK y -> [Concrete y]+      -> Concrete (BuildTensorKind k y)+ravel k stk l = tfromVector k stk (V.fromList l)++unravel :: forall k y.+           SNat k -> SingletonTK y -> Concrete (BuildTensorKind k y)+        -> [Concrete y]+unravel = tunravelToListShare++oRtmapAccumR+  :: forall k accy by ey.+     SNat k+  -> FullShapeTK by+  -> FullShapeTK ey+  -> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))+  -> Concrete accy+  -> Concrete (BuildTensorKind k ey)+  -> Concrete (TKProduct accy (BuildTensorKind k by))+oRtmapAccumR k bftk eftk f acc0 es = case sNatValue k of+  0 -> tpair acc0 (treplicate k (ftkToSTK bftk) (tdefTarget bftk))+  _ ->+    let g a b = let res = f a b+                in (tproject1 res, tproject2 res)+        (xout, lout) = mapAccumR g acc0 (unravel k (ftkToSTK eftk) es)+    in tpair xout (ravel k (ftkToSTK bftk) lout)+      -- TODO: reimplement not with Haskell's mapAccumR to avoid the ravels++oRtmapAccumL+  :: forall k accy by ey.+     SNat k+  -> FullShapeTK by+  -> FullShapeTK ey+  -> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))+  -> Concrete accy+  -> Concrete (BuildTensorKind k ey)+  -> Concrete (TKProduct accy (BuildTensorKind k by))+oRtmapAccumL k bftk eftk f acc0 es = case sNatValue k of+  0 -> tpair acc0 (treplicate k (ftkToSTK bftk) (tdefTarget bftk))+  _ ->+    let g a b = let res = f a b+                in (tproject1 res, tproject2 res)+        (xout, lout) = mapAccumL g acc0 (unravel k (ftkToSTK eftk) es)+    in tpair xout (ravel k (ftkToSTK bftk) lout)+++-- * Ranked internal definitions++-- TODO: check what the following did in tsum0R and if worth emulating+-- (also in sum1Inner and extremum and maybe tdot0R):+-- LA.sumElements $ OI.toUnorderedVectorT sh t+{-+tdot0R t u = OR.toVector t LA.<.> OR.toVector u+  -- TODO: if offset 0 and same strides, use toUnorderedVectorT+  -- TODO: if either has length 1 values, it may or may not be faster to do+  -- tsum0R (t * u) -}++-- TODO: check what the following did in tsum0R and if worth emulating+-- (also in sum1Inner and extremum and maybe tdot0R):+-- LA.sumElements $ OI.toUnorderedVectorT sh t++-- TODO: try to weave a similar magic as in tindex0R+-- TODO: for the non-singleton case see+-- https://github.com/Mikolaj/horde-ad/pull/81#discussion_r1096532164+updateNR :: forall n m x. (KnownNat n, KnownNat m, KnownSTK x)+         => Concrete (TKR2 (n + m) x)+         -> [(IxROf Concrete n, Concrete (TKR2 m x))]+         -> Concrete (TKR2 (n + m) x)+updateNR arr upd = case knownSTK @x of+  STKScalar ->  -- optimized+    let values = rtoVector arr+        sh = rshape arr+        f !t (ix, u) =+          let v = rtoVector u+              i = fromIntegral $ unConcrete+                  $ toLinearIdxR @n @m fromIntegral sh ix+          in V.concat [V.take i t, v, V.drop (i + V.length v) t]+    in Concrete $ Nested.rfromVector sh (foldl' f values upd)+  _ ->+    let arrNested = rnest (SNat @n) arr+        shNested = rshape arrNested+        f i v = case lookup (fromLinearIdxR+                               @n (Concrete . fromIntegral)+                               shNested ((Concrete . fromIntegral) i)) upd of+          Just u -> rnest (SNat @0) u+          Nothing -> v+    in runNest $ trfromVector0N shNested $ V.fromList+       $ imap f $ trunravelToList $ rflatten arrNested++tminIndexR+  :: forall r r2 n.+     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)+  => Nested.Ranked (1 + n) r -> Nested.Ranked n r2+tminIndexR v | SNat <- Nested.rrank v =+  let f :: Nested.Ranked 1 r -> Nested.Ranked 0 r2+      f = Nested.rscalar . fromIntegral . ixrHead . Nested.rminIndexPrim+  in Nested.rrerank SNat ZSR f v++tmaxIndexR+  :: forall r r2 n.+     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)+  => Nested.Ranked (1 + n) r -> Nested.Ranked n r2+tmaxIndexR v | SNat <- Nested.rrank v =+  let f :: Nested.Ranked 1 r -> Nested.Ranked 0 r2+      f = Nested.rscalar . fromIntegral . ixrHead . Nested.rmaxIndexPrim+  in Nested.rrerank SNat ZSR f v++-- We could generalize by unwinding and only then doing the PrimElt things,+-- but we'd need a type family that says "replace this underlying scalars+-- by this one", which makes things too complicated.+--+-- We could also expose `liftVR` in the user API, but in addition+-- to the main function argument such as floor or cast, it'd need the function's+-- derivative, just as with mapAccums. Maybe it's better to generalize even more+-- and permit arbitrary extra ops if given their derivatives.+liftVR+  :: (Nested.PrimElt r1, Nested.PrimElt r2)+  => (VS.Vector r1 -> VS.Vector r2)+  -> Nested.Ranked n r1 -> Nested.Ranked n r2+liftVR f = Ranked.liftRanked1 (Mixed.mliftNumElt1 (`liftVEltwise1` f))+{-# SPECIALIZE liftVR :: (VS.Vector Double -> VS.Vector Double) -> Nested.Ranked n Double -> Nested.Ranked n Double #-}+{-# SPECIALIZE liftVR :: (VS.Vector Float -> VS.Vector Float) -> Nested.Ranked n Float -> Nested.Ranked n Float #-}+{-# SPECIALIZE liftVR :: (VS.Vector Double -> VS.Vector Float) -> Nested.Ranked n Double -> Nested.Ranked n Float #-}+{-# SPECIALIZE liftVR :: (VS.Vector Float -> VS.Vector Double) -> Nested.Ranked n Float -> Nested.Ranked n Double #-}++ixInBounds :: [Int64] -> [Int] -> Bool+ixInBounds ix sh =+  and $ zipWith (\i dim -> 0 <= i && i < fromIntegral dim) ix sh++tindexNR+  :: (Nested.Elt r, Show r, KnownNat m, KnownNat n)+  => Nested.Ranked (m + n) r -> IxR m Int64 -> Nested.Ranked n r+tindexNR v ix = let sh = Nested.rshape v+                    !_A = assert (ixInBounds (toList ix) (toList sh)+                                  `blame` (v, ix)) ()+                in Nested.rindexPartial v (fmap fromIntegral ix)+{- TODO: benchmark if this is faster enough for its complexity;+         probably not, becasue orthotope's index does no canonicalization either+tindexNR v@(RS.A (RG.A sh OI.T{strides, offset, values})) ix =+  let l = indexToList ix+      linear = offset + sum (zipWith (*) (map fromIntegral l) strides)+      plen = valueOf @m  -- length of prefix being indexed out of+      !_A = assert (ixInBounds l sh `blame` (ix, sh, v)) ()+  in+    RS.A (RG.A (drop plen sh) OI.T{ strides = drop plen strides+                                  , offset = linear+                                  , values })+-}++tindexZR+  :: forall r m n. (KnownSTK r, KnownNat m, KnownNat n)+  => Concrete (TKR2 (m + n) r) -> IxROf Concrete m -> Concrete (TKR2 n r)+tindexZR v ixConcrete | Dict <- showDictRep (knownSTK @r)+                  , Dict <- eltDictRep (knownSTK @r) =+  let ix = fmapUnConcrete ixConcrete+  in case tftk knownSTK v of+    FTKR sh x ->+     if ixInBounds (Foldable.toList ix) (Foldable.toList sh)+     then Concrete $ tindexNR (unConcrete v) ix+     else tdefTarget (FTKR (shrDrop @m sh) x)++tindex0R+  :: forall r m. (KnownSTK r, KnownNat m)+  => Concrete (TKR2 m r) -> IxROf Concrete m -> Concrete (TKR2 0 r)+tindex0R v ixConcrete | Dict <- eltDictRep (knownSTK @r) =+  let ix = fmapUnConcrete ixConcrete+  in case tftk knownSTK v of+    FTKR sh x ->+      if ixInBounds (toList ix) (toList sh)+      then let arr = Nested.rscalar+                     $ Nested.rindex (unConcrete v) (fmap fromIntegral ix)+           in Concrete arr+      else tdefTarget (FTKR ZSR x)+{- TODO: see above+tindex0R (RS.A (RG.A _ OI.T{..})) ix =+  values V.! (offset + sum (zipWith (*) (map fromIntegral $ indexToList ix)+                                        strides))+-}++-- Performance depends a lot on the number and size of tensors.+-- If tensors are not tiny, memory taken by underlying vectors matters most+-- and this implementation is probbaly optimal in this respect+-- (the only new vectors are created by V.concat, but this is done on demand).+-- TODO: optimize updateNR and make it consume and forget arguments+-- one by one to make the above true+--+-- Note how ix being in bounds is checked. The semantics of the operation+-- permits index out of bounds and then no tensors is added at such an index.+tscatterZR :: forall m p n r.+              (KnownNat p, KnownNat m, KnownNat n, KnownSTK r)+           => IShR (p + n) -> Concrete (TKR2 (m + n) r)+           -> (IxROf Concrete m -> IxROf Concrete p)+           -> Concrete (TKR2 (p + n) r)+tscatterZR sh t f+ | Dict <- eltDictRep (knownSTK @r) = case tftk knownSTK t of+  FTKR _ x@FTKScalar ->  -- optimized+    let zero = tdefTarget (FTKR sh x)+        (shm, shDropP) = shrSplitAt @m $ rshape t+        s = shrSize shm+        g ix =+          let ix2 = f $ fmapConcrete ix+          in if ixInBounds (fmapUnConcrete $ toList ix2) (toList sh)+             then M.insertWith (V.zipWith (+)) ix2+                               (Nested.rtoVector $ unConcrete t `tindexNR` ix)+             else id+        ivs = foldr g M.empty [ fromLinearIdxR fromIntegral shm i+                              | i <- [0 .. fromIntegral s - 1] ]+    in updateNR zero+       $ map (second $ Concrete . Nested.rfromVector shDropP)+       $ M.assocs ivs+  FTKR _ x | Dict <- showDictRep (ftkToSTK x) ->+    let zero = tdefTarget (FTKR sh x)+        (shm, _) = shrSplitAt @m $ rshape t+        s = shrSize shm+        g ix =+          let ix2 = f $ fmapConcrete ix+          in if ixInBounds (fmapUnConcrete $ toList ix2) (toList sh)+             then M.insertWith (taddTarget knownSTK) ix2+                               (Concrete $ unConcrete t `tindexNR` ix)+             else id+        ivs = foldr g M.empty [ fromLinearIdxR fromIntegral shm i+                              | i <- [0 .. fromIntegral s - 1] ]+    in updateNR zero+       $ M.assocs ivs++-- TODO: update in place in ST or with a vector builder, but that requires+-- building the underlying value vector with crafty index computations+-- and then freezing it and calling Nested.rfromVector+-- or optimize tscatterNR and instantiate it instead+tscatterZ1R :: (KnownSTK r, KnownNat p, KnownNat n)+            => IShR (p + n) -> Concrete (TKR2 (1 + n) r)+            -> (IntOf Concrete -> IxROf Concrete p)+            -> Concrete (TKR2 (p + n) r)+tscatterZ1R sh t f = case tftk knownSTK t of+  FTKR _ x ->+    let zero = tdefTarget (FTKR sh x)+        lt = trunravelToList t+        g i ti = let ix2 = f $ Concrete $ fromIntegral i+                 in if ixInBounds (fmapUnConcrete $ toList ix2) (toList sh)+                    then updateNR zero [(ix2, ti)]+                    else zero+        lu = imap g lt+    in foldr (taddTarget knownSTK) zero lu++tfromVector0NR+  :: Nested.KnownElt r+  => IShR n -> Data.Vector.Vector (Nested.Ranked 0 r) -> Nested.Ranked n r+tfromVector0NR sh l = case NonEmpty.nonEmpty $ V.toList l of+  Nothing -> Nested.rreshape sh Nested.remptyArray+  Just nl -> Nested.rfromListLinear sh $ NonEmpty.map Nested.runScalar nl++tbuild1R+  :: forall r n. (Nested.KnownElt r, KnownNat n)+  => Int -> (Int64 -> Nested.Ranked n r) -> Nested.Ranked (1 + n) r+tbuild1R k f = case NonEmpty.nonEmpty [0 .. fromIntegral k - 1] of+  Nothing -> case sameNat (Proxy @n) (Proxy @0) of+    Just Refl -> Nested.remptyArray+    Nothing -> error "rbuild1: shape ambiguity"+  Just l -> Nested.rfromListOuter $ NonEmpty.map f l  -- hope this fuses++tmap0NR+  :: (Nested.PrimElt r1, Nested.PrimElt r)+  => (Nested.Ranked 0 r1 -> Nested.Ranked 0 r) -> Nested.Ranked n r1+  -> Nested.Ranked n r+tmap0NR f = Ranked.liftRanked1+              (Mixed.mliftPrim (Nested.runScalar . f . Nested.rscalar ))+  -- too slow: tbuildNR (Nested.rshape v) (\ix -> f $ v `tindexNR` ix)++tzipWith0NR+  :: (Nested.PrimElt r, Nested.PrimElt r1, Nested.PrimElt r2)+  => (Nested.Ranked 0 r1 -> Nested.Ranked 0 r2 -> Nested.Ranked 0 r)+  -> Nested.Ranked n r1 -> Nested.Ranked n r2 -> Nested.Ranked n r+tzipWith0NR f =+  Ranked.liftRanked2+    (Mixed.mliftPrim2+       (\x y -> Nested.runScalar $ f (Nested.rscalar x) (Nested.rscalar y)))++-- The semantics of the operation permits index out of bounds+-- and the result of such indexing is def, which is 0.+tgatherZR :: forall m p n r.+             (KnownNat m, KnownNat p, KnownNat n, KnownSTK r)+          => IShR (m + n) -> Concrete (TKR2 (p + n) r)+          -> (IxROf Concrete m -> IxROf Concrete p)+          -> Concrete (TKR2 (m + n) r)+tgatherZR sh t f = case knownSTK @r of+  STKScalar ->  -- optimized+    let shm = shrTake @m sh+        s = shrSize shm+        l = [ rtoVector+              $ t `trindex` f (fmapConcrete $ fromLinearIdxR fromIntegral shm i)+            | i <- [0 .. fromIntegral s - 1] ]+    in Concrete $ Nested.rfromVector sh $ V.concat l+  _ -> rbuild sh (\ix -> t `trindex` f ix)++tgatherZ1R :: forall p n r.+              (KnownNat p, KnownNat n, KnownSTK r)+           => Int -> Concrete (TKR2 (p + n) r)+           -> (IntOf Concrete -> IxROf Concrete p)+           -> Concrete (TKR2 (1 + n) r)+tgatherZ1R k t f = case knownSTK @r of+  STKScalar ->  -- optimized+    trfromVector $ V.fromList $ map (\i -> t `trindex` f (Concrete i))+                                    [0 .. fromIntegral k - 1]+  _ -> trbuild1 k (\ix -> t `trindex` f ix)+++-- * Shaped internal definitions++-- TODO: try to weave a similar magic as in tindex0R+-- TODO: for the non-singleton case see+-- https://github.com/Mikolaj/horde-ad/pull/81#discussion_r1096532164+updateNS :: forall n sh r.+            ( KnownSTK r, KnownShS sh, KnownShS (Drop n sh)+            , KnownShS (Take n sh) )+         => Concrete (TKS2 sh r)+         -> [(IxSOf Concrete (Take n sh), Concrete (TKS2 (Drop n sh) r))]+         -> Concrete (TKS2 sh r)+updateNS arr upd = case knownSTK @r of+  STKScalar ->+    let values = stoVector arr+        sh = knownShS @sh+        f !t (ix, u) =+          let v = stoVector u+              i = gcastWith (unsafeCoerceRefl+                             :: sh :~: Take n sh ++ Drop n sh)+                  $ fromIntegral $ unConcrete+                  $ toLinearIdxS @(Take n sh) @(Drop n sh)+                                 fromIntegral sh ix+          in V.concat [V.take i t, v, V.drop (i + V.length v) t]+    in Concrete $ Nested.sfromVector knownShS (foldl' f values upd)+  _ -> case shsProduct (knownShS @(Take n sh)) of+    SNat ->+      gcastWith (unsafeCoerceRefl :: sh :~: Take n sh ++ Drop n sh) $+      let arrNested = snest (knownShS @(Take n sh)) arr+          shNested = sshape arrNested+          f i v = case lookup (fromLinearIdxS+                                 @(Take n sh) (Concrete . fromIntegral)+                                 shNested ((Concrete . fromIntegral) i)) upd of+            Just u -> snest (knownShS @'[]) u+            Nothing -> v+      in sunNest @_ @(Take n sh) $ tsfromVector0N $ V.fromList+         $ imap f $ tsunravelToList $ sflatten arrNested++tfromIntegralS :: (GoodScalar r1, Integral r1, GoodScalar r2)+               => Nested.Shaped sh r1 -> Nested.Shaped sh r2+tfromIntegralS = liftVS (V.map fromIntegral)++tminIndexS+  :: forall n sh r r2.+     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)+  => Nested.Shaped (n ': sh) r -> Nested.Shaped (Init (n ': sh)) r2+tminIndexS v | sh1@(_ :$$ sh) <- Nested.sshape v =+  let f :: Nested.Shaped '[m] r -> Nested.Shaped '[] r2+      f = Nested.sscalar . fromIntegral . ixsHead . Nested.sminIndexPrim+  in case sh of+    ZSS -> f @n v+    _ | SNat @m <- shsLast sh1 ->+      gcastWith (unsafeCoerceRefl+                 :: Init (n ': sh) ++ '[m] :~: n ': sh) $+      gcastWith (unsafeCoerceRefl+                 :: Init (n ': sh) :~: Init (n ': sh) ++ '[]) $+      Nested.srerank @'[m] @'[] @(Init (n ': sh))+                     (shsInit sh1) ZSS (f @m) v++tmaxIndexS+  :: forall n sh r r2.+     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)+  => Nested.Shaped (n ': sh) r -> Nested.Shaped (Init (n ': sh)) r2+tmaxIndexS v | sh1@(_ :$$ sh) <- Nested.sshape v =+  let f :: Nested.Shaped '[m] r -> Nested.Shaped '[] r2+      f = Nested.sscalar . fromIntegral . ixsHead . Nested.smaxIndexPrim+  in case sh of+    ZSS -> f @n v+    _ | SNat @m <- shsLast sh1 ->+      gcastWith (unsafeCoerceRefl+                 :: Init (n ': sh) ++ '[m] :~: n ': sh) $+      gcastWith (unsafeCoerceRefl+                 :: Init (n ': sh) :~: Init (n ': sh) ++ '[]) $+      Nested.srerank @'[m] @'[] @(Init (n ': sh))+                     (shsInit sh1) ZSS (f @m) v++liftVS+  :: (Nested.PrimElt r1, Nested.PrimElt r)+  => (VS.Vector r1 -> VS.Vector r)+  -> Nested.Shaped sh r1 -> Nested.Shaped sh r+liftVS f = Shaped.liftShaped1 (Mixed.mliftNumElt1 (`liftVEltwise1` f))+{-# SPECIALIZE liftVS :: (VS.Vector Double -> VS.Vector Double) -> Nested.Shaped sh Double -> Nested.Shaped sh Double #-}+{-# SPECIALIZE liftVS :: (VS.Vector Float -> VS.Vector Float) -> Nested.Shaped sh Float -> Nested.Shaped sh Float #-}+{-# SPECIALIZE liftVS :: (VS.Vector Double -> VS.Vector Float) -> Nested.Shaped sh Double -> Nested.Shaped sh Float #-}+{-# SPECIALIZE liftVS :: (VS.Vector Float -> VS.Vector Double) -> Nested.Shaped sh Float -> Nested.Shaped sh Double #-}++tindexNS+  :: Nested.Elt r+  => Nested.Shaped (sh1 ++ sh2) r -> IxS sh1 Int64 -> Nested.Shaped sh2 r+tindexNS v ix = Nested.sindexPartial v (fmap fromIntegral ix)+{- TODO+tindexNS (SS.A (SG.A OI.T{strides, offset, values})) ix =+  let l = ShapedList.indexToList ix+      linear = offset + sum (zipWith (*) (map fromIntegral l) strides)+      plen = length l  -- length of prefix being indexed out of+  in+    SS.A (SG.A OI.T{ strides = drop plen strides+                   , offset = linear+                   , values })+-}++-- Note that after vectorization, the index may not fit within+-- the type-level shape, which we catch in the @ixInBounds@+-- and return def, so it's fine. Similarly in gather and scatter.+tindexZS+  :: forall r sh1 sh2. (KnownSTK r, KnownShS sh1, KnownShS sh2)+  => Concrete (TKS2 (sh1 ++ sh2) r) -> IxSOf Concrete sh1+  -> Concrete (TKS2 sh2 r)+tindexZS v ixConcrete | Dict <- eltDictRep (knownSTK @r) =+  let ix = fmapUnConcrete ixConcrete+  in withKnownShS (knownShS @sh1 `shsAppend` knownShS @sh2) $+     case tftk knownSTK v of+       FTKS sh x ->+         if ixInBounds (Foldable.toList ix) (shsToList sh)+         then Concrete $ tindexNS (unConcrete v) ix+         else tdefTarget (FTKS knownShS x)++tindex0S+  :: forall r sh. (KnownSTK r, KnownShS sh)+  => Concrete (TKS2 sh r) -> IxSOf Concrete sh -> Concrete (TKS2 '[] r)+tindex0S v ixConcrete | Dict <- eltDictRep (knownSTK @r) =+  let ix = fmapUnConcrete ixConcrete+  in case tftk knownSTK v of+    FTKS sh x ->+      if ixInBounds (toList ix) (toList sh)+      then let arr = Nested.sscalar+                     $ Nested.sindex (unConcrete v) (fmap fromIntegral ix)+           in Concrete arr+      else tdefTarget (FTKS ZSS x)+{- TODO: benchmark if this is faster enough for its complexity;+         probably not, becasue orthotope's index does no canonicalization either+tindex0S (SS.A (SG.A OI.T{..})) ix =+  values V.! (offset + sum (zipWith (*) (map fromIntegral+                                         $ ShapedList.indexToList ix)+                                        strides))+    -- to avoid linearizing @values@, we do everything in unsized way+-}++-- TODO: update in place in ST or with a vector builder, but that requires+-- building the underlying value vector with crafty index computations+-- and then freezing it and calling OS.fromVector+-- or optimize tscatterNS and instantiate it instead+tscatterZ1S+  :: forall r n2 shn shp.+     (KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp)+  => Concrete (TKS2 (n2 ': shn) r)+  -> (IntOf Concrete -> IxSOf Concrete shp)+  -> Concrete (TKS2 (shp ++ shn) r)+tscatterZ1S t f = case tftk knownSTK t of+  FTKS _ x ->+    gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn) :~: shp) $+    gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn) :~: shn) $+    let shpshn = knownShS @shp `shsAppend` knownShS @shn+        zero = tdefTarget (FTKS shpshn x)+        lt = tsunravelToList t+        g i ti = let ix2 = f $ Concrete $ fromIntegral i+                 in if ixInBounds (fmapUnConcrete $ Foldable.toList ix2)+                                  (shsToList shpshn)+                    then withKnownShS shpshn $+                         updateNS @(Rank shp) zero [(ix2, ti)]+                    else zero+        lu = imap g lt+    in foldr (taddTarget (STKS shpshn (knownSTK @r))) zero lu++tfromVector0NS+  :: forall r sh. (Nested.KnownElt r, KnownShS sh)+  => Data.Vector.Vector (Nested.Shaped '[] r) -> Nested.Shaped sh r+tfromVector0NS l = case NonEmpty.nonEmpty $ V.toList l of+  Nothing -> case testEquality (shsProduct (knownShS @sh)) (SNat @0) of+    Just Refl -> Nested.sreshape (knownShS @sh)+                 $ Nested.semptyArray (knownShS @sh)+    Nothing -> error "tfromVector0N: empty list, but not shape"+  Just nl -> Nested.sfromListLinear knownShS $ NonEmpty.map Nested.sunScalar nl++tbuild1S+  :: forall k sh r. (KnownNat k, KnownShS sh, Nested.KnownElt r)+  => (Int64 -> Nested.Shaped sh r) -> Nested.Shaped (k ': sh) r+tbuild1S f = case NonEmpty.nonEmpty [0 .. valueOf @k - 1] of+  Nothing -> gcastWith (unsafeCoerceRefl :: k :~: 0) $+             Nested.semptyArray knownShS+  Just l -> Nested.sfromListOuter SNat $ NonEmpty.map f l  -- hope this fuses++tmap0NS+  :: forall r1 r sh. (Nested.PrimElt r1, Nested.PrimElt r)+  => (Nested.Shaped '[] r1 -> Nested.Shaped '[] r) -> Nested.Shaped sh r1+  -> Nested.Shaped sh r+tmap0NS f =+  Shaped.liftShaped1+    (Mixed.mliftPrim (Nested.sunScalar . f . Nested.sscalar))+      -- too slow: tbuildNS (tshapeS v) (\ix -> f $ v `tindexNS` ix)++tzipWith0NS+  :: forall r1 r2 r sh. (Nested.PrimElt r, Nested.PrimElt r1, Nested.PrimElt r2)+  => (Nested.Shaped '[] r1 -> Nested.Shaped '[] r2 -> Nested.Shaped '[] r)+  -> Nested.Shaped sh r1 -> Nested.Shaped sh r2 -> Nested.Shaped sh r+tzipWith0NS f =+  Shaped.liftShaped2+    (Mixed.mliftPrim2+       (\x y -> Nested.sunScalar $ f (Nested.sscalar x) (Nested.sscalar y)))++tgatherZ1S+  :: forall r n2 shn shp.+     (KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp)+  => Concrete (TKS2 (shp ++ shn) r)+  -> (IntOf Concrete -> IxSOf Concrete shp)+  -> Concrete (TKS2 (n2 ': shn) r)+tgatherZ1S t f =+  case knownSTK @r of+    STKScalar ->  -- optimized+      tsfromVector $ V.fromList $ map (\i -> t `tsindex` f (Concrete i))+                                      [0 .. valueOf @n2 - 1]+    _ -> tsbuild1 (\ix -> t `tsindex` f ix)+++-- * Mixed internal definitions++updateNX :: forall n sh r.+            (KnownSTK r, KnownShX (Drop n sh), KnownShX (Take n sh))+         => Concrete (TKX2 sh r)+         -> [(IxXOf Concrete (Take n sh), Concrete (TKX2 (Drop n sh) r))]+         -> Concrete (TKX2 sh r)+updateNX arr upd = case knownSTK @r of+  STKScalar ->+    let values = xtoVector arr+        sh = xshape arr+        f !t (ix, u) =+          let v = xtoVector u+              i = gcastWith (unsafeCoerceRefl+                             :: sh :~: Take n sh ++ Drop n sh)+                  $ fromIntegral $ unConcrete+                  $ toLinearIdxX @(Take n sh) @(Drop n sh)+                                 fromIntegral sh ix+          in V.concat [V.take i t, v, V.drop (i + V.length v) t]+    in Concrete $ Nested.mfromVector (xshape arr) (foldl' f values upd)+  _ | Dict <- eltDictRep (knownSTK @r) ->+      gcastWith (unsafeCoerceRefl :: sh :~: Take n sh ++ Drop n sh) $+      let arrNested = xnest (knownShX @(Take n sh)) arr+          shNested = xshape arrNested+          f i v = case lookup (fromLinearIdxX+                                 @(Take n sh) (Concrete . fromIntegral)+                                 shNested ((Concrete . fromIntegral) i)) upd of+            Just u -> xnest ZKX u+            Nothing -> v+      in withSNat (shxSize shNested) $ \snat ->+           xunNest @_ @(Take n sh) $ txfromVector0N shNested $ V.fromList+           $ imap f $ txunravelToList+           $ Concrete $ Nested.mcast (Nested.SKnown snat :!% ZKX)+           $ unConcrete $ xflatten arrNested++tminIndexX+  :: forall mn sh r r2.+     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)+  => Nested.Mixed (mn ': sh) r -> Nested.Mixed (Init (mn ': sh)) r2+tminIndexX v | sh1@(_ :$% sh) <- Nested.mshape v =+  let f :: Nested.Mixed '[mm] r -> Nested.Mixed '[] r2+      f = Nested.mscalar . fromIntegral . ixxHead+          . Nested.mminIndexPrim+  in case sh of+    ZSX -> f @mn v+    _ -> withSNat (fromSMayNat' (shxLast sh1)) $ \(_ :: SNat m) ->+      gcastWith (unsafeCoerceRefl+                 :: Init (mn ': sh) ++ '[Just m] :~: mn ': sh) $+      gcastWith (unsafeCoerceRefl+                 :: Init (mn ': sh) :~: Init (mn ': sh) ++ '[]) $+      Nested.mrerank @'[Just m] @'[] @(Init (mn ': sh))+                     (ssxFromShX $ shxInit sh1) ZSX (f @(Just m)) v++tmaxIndexX+  :: forall mn sh r r2.+     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)+  => Nested.Mixed (mn ': sh) r -> Nested.Mixed (Init (mn ': sh)) r2+tmaxIndexX v | sh1@(_ :$% sh) <- Nested.mshape v =+  let f :: Nested.Mixed '[mm] r -> Nested.Mixed '[] r2+      f = Nested.mscalar . fromIntegral . ixxHead+          . Nested.mmaxIndexPrim+  in case sh of+    ZSX -> f @mn v+    _ -> withSNat (fromSMayNat' (shxLast sh1)) $ \(_ :: SNat m) ->+      gcastWith (unsafeCoerceRefl+                 :: Init (mn ': sh) ++ '[Just m] :~: mn ': sh) $+      gcastWith (unsafeCoerceRefl+                 :: Init (mn ': sh) :~: Init (mn ': sh) ++ '[]) $+      Nested.mrerank @'[Just m] @'[] @(Init (mn ': sh))+                     (ssxFromShX $ shxInit sh1) ZSX (f @(Just m)) v++liftVX+  :: (Nested.PrimElt r1, Nested.PrimElt r)+  => (VS.Vector r1 -> VS.Vector r)+  -> Nested.Mixed sh r1 -> Nested.Mixed sh r+liftVX f = Mixed.mliftNumElt1 (`liftVEltwise1` f)+{-# SPECIALIZE liftVX :: (VS.Vector Double -> VS.Vector Double) -> Nested.Mixed sh Double -> Nested.Mixed sh Double #-}+{-# SPECIALIZE liftVX :: (VS.Vector Float -> VS.Vector Float) -> Nested.Mixed sh Float -> Nested.Mixed sh Float #-}+{-# SPECIALIZE liftVX :: (VS.Vector Double -> VS.Vector Float) -> Nested.Mixed sh Double -> Nested.Mixed sh Float #-}+{-# SPECIALIZE liftVX :: (VS.Vector Float -> VS.Vector Double) -> Nested.Mixed sh Float -> Nested.Mixed sh Double #-}++tindexNX+  :: Nested.Elt r+  => Nested.Mixed (sh1 ++ sh2) r -> IxX sh1 Int64 -> Nested.Mixed sh2 r+tindexNX v ix = Nested.mindexPartial v (fmap fromIntegral ix)++tindexZX+  :: forall r sh1 sh2. (KnownSTK r, KnownShX sh1, KnownShX sh2)+  => Concrete (TKX2 (sh1 ++ sh2) r) -> IxXOf Concrete sh1+  -> Concrete (TKX2 sh2 r)+tindexZX v ixConcrete | Dict <- eltDictRep (knownSTK @r) =+  let ix = fmapUnConcrete ixConcrete+  in withKnownShX (knownShX @sh1 `ssxAppend` knownShX @sh2) $+     case tftk knownSTK v of+       FTKX sh x ->+         if ixInBounds (Foldable.toList ix) (shxToList sh)+         then Concrete $ tindexNX (unConcrete v) ix+         else tdefTarget (FTKX (shxDropSSX (knownShX @sh1) sh) x)++tindex0X+  :: forall r sh. (KnownSTK r, KnownShX sh)+  => Concrete (TKX2 sh r) -> IxXOf Concrete sh -> Concrete (TKX2 '[] r)+tindex0X v ixConcrete | Dict <- eltDictRep (knownSTK @r) =+  let ix = fmapUnConcrete ixConcrete+  in case tftk knownSTK v of+    FTKX sh x ->+      if ixInBounds (toList ix) (toList sh)+      then let arr = Nested.mscalar+                     $ Nested.mindex (unConcrete v) (fmap fromIntegral ix)+           in Concrete arr+      else tdefTarget (FTKX ZSX x)++tscatterZ1X+  :: forall r n2 shn shp.+     (KnownSTK r, KnownNat n2, KnownShX shn, KnownShX shp)+  => IShX (shp ++ shn) -> Concrete (TKX2 (Just n2 ': shn) r)+  -> (IntOf Concrete -> IxXOf Concrete shp)+  -> Concrete (TKX2 (shp ++ shn) r)+tscatterZ1X sh t f =+  case tftk knownSTK t of+    FTKX _ x ->+      withKnownShX (ssxFromShX sh) $+      gcastWith (unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn) :~: shp) $+      gcastWith (unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn) :~: shn) $+      let zero = tdefTarget (FTKX sh x)+          lt = txunravelToList t+          g i ti = let ix2 = f $ Concrete $ fromIntegral i+                   in if ixInBounds (fmapUnConcrete $ Foldable.toList ix2)+                                    (shxToList sh)+                      then updateNX @(Rank shp) zero [(ix2, ti)]+                      else zero+          lu = imap g lt+      in foldr (taddTarget knownSTK) zero lu++tfromVector0NX+  :: forall r sh. Nested.KnownElt r+  => IShX sh -> Data.Vector.Vector (Nested.Mixed '[] r) -> Nested.Mixed sh r+tfromVector0NX sh l = case NonEmpty.nonEmpty $ V.toList l of+  Nothing -> if shxSize sh == 0+             then Nested.mreshape sh $ Nested.memptyArray sh+             else error "tfromVector0N: empty list, but not shape"+  Just nl -> Nested.mfromListLinear sh $ NonEmpty.map Nested.munScalar nl++tbuild1X+  :: forall k sh r. (KnownNat k, KnownShX sh, Nested.KnownElt r)+  => (Int64 -> Nested.Mixed sh r)+  -> Nested.Mixed (Just k ': sh) r+tbuild1X f = case NonEmpty.nonEmpty [0 .. valueOf @k - 1] of+  Nothing -> case testEquality (knownShX @sh) ZKX of+    Just Refl -> gcastWith (unsafeCoerceRefl :: k :~: 0) $+                 Nested.memptyArray ZSX+    Nothing -> error "xbuild1: shape ambiguity"+  Just l -> Nested.mcast (Nested.SKnown (SNat @k) :!% knownShX)+            $ Nested.mfromListOuter $ NonEmpty.map f l  -- hope this fuses++tgatherZ1X+  :: forall r n2 shn shp.+     (KnownSTK r, KnownShX shn, KnownShX shp)+  => SNat n2 -> Concrete (TKX2 (shp ++ shn) r)+  -> (IntOf Concrete -> IxXOf Concrete shp)+  -> Concrete (TKX2 (Just n2 ': shn) r)+tgatherZ1X SNat t f =+  case knownSTK @r of+    STKScalar ->  -- optimized+      txfromVector $ V.fromList $ map (\i -> t `txindex` f (Concrete i))+                                      [0 .. valueOf @n2 - 1]+    _ -> txbuild1 @_ @n2 (\ix -> t `txindex` f ix)++fmapConcrete :: Coercible (f (RepConcrete y)) (f (Concrete y))+               => f (RepConcrete y) -> f (Concrete y)+fmapConcrete = coerce++fmapUnConcrete :: Coercible (f (Concrete y)) (f (RepConcrete y))+               => f (Concrete y) -> f (RepConcrete y)+fmapUnConcrete = coerce
+ src/HordeAd/Core/TensorKind.hs view
@@ -0,0 +1,456 @@+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Two kinds of singletons for tensor kindss and constraints+-- and lemmas associated with the singletons.+module HordeAd.Core.TensorKind+  ( -- * Tensor kind singletons+    SingletonTK(..), KnownSTK(..)+  , TKConversion(..), convertSTK, convertFTK, buildTKConversion+  , withKnownSTK, lemKnownSTK, sameKnownSTK, sameSTK+  , stkUnit, buildSTK, razeSTK, adSTK+  , lemKnownSTKOfBuild, lemKnownSTKOfAD, lemBuildOfAD, lengthSTK, widthSTK+    -- * Full shape tensor kind quasi-singletons+  , FullShapeTK(..)+  , matchingFTK, ftkToSTK, ftkUnit, buildFTK, razeFTK, adFTK, differentiableFTK+  , DummyDualTarget(..)+  ) where++import Prelude hiding ((.))++import Control.Category+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import GHC.Exts (withDict)+import GHC.TypeLits (KnownNat, OrderingI (..), cmpNat, fromSNat, type (+))+import Type.Reflection (typeRep)++import Data.Array.Nested (MapJust, Replicate, type (++))+import Data.Array.Nested.Convert (shrFromShX, shsFromShX, shxFromShS, shxFromShR, shsFromSSX)+import Data.Array.Nested.Lemmas+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.Types++-- * Tensor kind singletons++-- | Tensor kind singleton type.+type role SingletonTK nominal+data SingletonTK y where+  STKScalar :: GoodScalar r+            => SingletonTK (TKScalar r)+  STKR :: SNat n -> SingletonTK x -> SingletonTK (TKR2 n x)+  STKS :: ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)+  STKX :: StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)+  STKProduct :: SingletonTK y -> SingletonTK z+             -> SingletonTK (TKProduct y z)++deriving instance Show (SingletonTK y)++-- | The constraints corresponding to 'SingletonTK'.+class KnownSTK (y :: TK) where+  knownSTK :: SingletonTK y++instance GoodScalar r => KnownSTK (TKScalar r) where+  knownSTK = STKScalar++instance (KnownSTK x, KnownNat n)+         => KnownSTK (TKR2 n x) where+  knownSTK = STKR SNat knownSTK++instance (KnownSTK x, KnownShS sh)+         => KnownSTK (TKS2 sh x) where+  knownSTK = STKS knownShS knownSTK++instance (KnownSTK x, KnownShX sh)+         => KnownSTK (TKX2 sh x) where+  knownSTK = STKX knownShX knownSTK++instance (KnownSTK y, KnownSTK z)+         => KnownSTK (TKProduct y z) where+  knownSTK = STKProduct (knownSTK @y) (knownSTK @z)++-- | Turning a singleton into a constraint via a continuation.+withKnownSTK :: forall y r. SingletonTK y -> (KnownSTK y => r) -> r+withKnownSTK = withDict @(KnownSTK y)++-- | Turning a singleton into a dictionary containing constraint.+lemKnownSTK :: SingletonTK y -> Dict KnownSTK y+lemKnownSTK = \case+  STKScalar -> Dict+  STKR SNat x | Dict <- lemKnownSTK x -> Dict+  STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh Dict+  STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh Dict+  STKProduct stk1 stk2 | Dict <- lemKnownSTK stk1+                       , Dict <- lemKnownSTK stk2 -> Dict++sameKnownSTK :: forall y1 y2. (KnownSTK y1, KnownSTK y2)+             => Maybe (y1 :~: y2)+sameKnownSTK = sameSTK (knownSTK @y1) (knownSTK @y2)++-- | A plausible implementation of `testEquality` on `SingletonTK`.+sameSTK :: SingletonTK y1 -> SingletonTK y2 -> Maybe (y1 :~: y2)+sameSTK stk1 stk2 = case (stk1, stk2) of+  (STKScalar @r1, STKScalar @r2)+    | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->+      Just Refl+  (STKR snat1 x1, STKR snat2 x2)+    | Just Refl <- sameSTK x1 x2, Just Refl <- testEquality snat1 snat2 ->+      Just Refl+  (STKS sh1 x1, STKS sh2 x2)+    | Just Refl <- sameSTK x1 x2, Just Refl <- testEquality sh1 sh2 ->+      Just Refl+  (STKX sh1 x1, STKX sh2 x2)+    | Just Refl <- sameSTK x1 x2, Just Refl <- testEquality sh1 sh2 ->+      Just Refl+  (STKProduct x1 y1, STKProduct x2 y2)+    | Just Refl <- sameSTK x1 x2, Just Refl <- sameSTK y1 y2 ->+      Just Refl+  _ -> Nothing++stkUnit :: SingletonTK TKUnit+stkUnit = STKScalar++buildSTK :: SNat k -> SingletonTK y -> SingletonTK (BuildTensorKind k y)+buildSTK snat@SNat = \case+  stk@STKScalar -> STKS (snat :$$ ZSS) stk+  STKR SNat x -> STKR SNat x+  STKS sh x -> STKS (snat :$$ sh) x+  STKX sh x -> STKX (SKnown snat :!% sh) x+  STKProduct stk1 stk2 -> STKProduct (buildSTK snat stk1) (buildSTK snat stk2)++razeSTK :: SingletonTK z -> SingletonTK (RazeTensorKind z)+razeSTK = \case+  STKScalar -> error "razeSTK: impossible argument"+  STKR snat@SNat x ->+    case cmpNat (SNat @1) snat of+      LTI -> STKR SNat x+      EQI -> STKR SNat x+      _ -> error "razeSTK: impossible argument"+  STKS ZSS _ -> error "razeSTK: impossible argument"+  STKS (_ :$$ sh) x -> STKS sh x+  STKX ZKX _ -> error "razeSTK: impossible argument"+  STKX (SUnknown _ :!% _) _ -> error "razeSTK: impossible argument"+  STKX (SKnown _ :!% sh) x -> STKX sh x+  STKProduct stk1 stk2 -> STKProduct (razeSTK stk1) (razeSTK stk2)++adSTK :: SingletonTK y -> SingletonTK (ADTensorKind y)+adSTK = \case+  t@(STKScalar @r) -> case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> t+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> t+      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1)+           STKScalar+  STKR sh x -> STKR sh $ adSTK x+  STKS sh x -> STKS sh $ adSTK x+  STKX sh x -> STKX sh $ adSTK x+  STKProduct stk1 stk2 -> STKProduct (adSTK stk1) (adSTK stk2)++lemKnownSTKOfBuild :: SNat k -> SingletonTK y+                     -> Dict KnownSTK (BuildTensorKind k y)+lemKnownSTKOfBuild snat = lemKnownSTK . buildSTK snat++lemKnownSTKOfAD :: SingletonTK y+                  -> Dict KnownSTK (ADTensorKind y)+lemKnownSTKOfAD = lemKnownSTK . adSTK++lemBuildOfAD :: SNat k -> SingletonTK y+             -> BuildTensorKind k (ADTensorKind y)+                :~: ADTensorKind (BuildTensorKind k y)+lemBuildOfAD snat@SNat = \case+  STKScalar -> Refl+  STKR{} -> unsafeCoerceRefl+  STKS{} -> unsafeCoerceRefl+  STKX{} -> unsafeCoerceRefl+  STKProduct stk1 stk2 | Refl <- lemBuildOfAD snat stk1+                       , Refl <- lemBuildOfAD snat stk2 -> Refl++lengthSTK :: SingletonTK x -> Int+lengthSTK STKScalar = 0+lengthSTK (STKR snat _) = fromInteger $ fromSNat snat+lengthSTK (STKS sh _) = shsLength sh+lengthSTK (STKX sh _) = ssxLength sh+lengthSTK (STKProduct sy sz) = lengthSTK sy `max` lengthSTK sz++widthSTK :: SingletonTK y -> Int+widthSTK stk = case stk of+  STKScalar @r -> case testEquality (typeRep @r) (typeRep @Z1) of+    Just Refl -> 0+    _ -> 1+  STKR{} -> 1+  STKS{} -> 1+  STKX{} -> 1+  STKProduct stk1 stk2 -> widthSTK stk1 + widthSTK stk2++-- | This is copied, with modifications, from ox-arrays.+--+-- This is a recipe for converting arrays, not always followed,+-- and a proof a conversion is possible, with some proof obligations+-- delayed to runtime (in ConvXS' and ConvXX', where not only the ranks+-- of the shapes need to agree, but also the dimensions of the input+-- array and of the output shape, which is not all captured in the type).+-- As in ox-arrays, conversions only change the meta-data, not the underlying+-- vector representation of the array.+type role TKConversion nominal nominal+data TKConversion (a :: TK) (b :: TK) where+  ConvId  :: TKConversion a a+  ConvCmp :: TKConversion b c -> TKConversion a b -> TKConversion a c++  ConvRX  :: TKConversion (TKR2 n a) (TKX2 (Replicate n Nothing) a)+  ConvSX  :: TKConversion (TKS2 sh a) (TKX2 (MapJust sh) a)++  ConvXR  :: SingletonTK a -> TKConversion (TKX2 sh a) (TKR2 (Rank sh) a)+  ConvXS  :: TKConversion (TKX2 (MapJust sh) a) (TKS2 sh a)+  ConvXS' :: Rank sh ~ Rank sh'+          => FullShapeTK (TKS2 sh' a)+          -> TKConversion (TKX2 sh a) (TKS2 sh' a)++  ConvXX' :: Rank sh ~ Rank sh'+          => FullShapeTK (TKX2 sh' a)+          -> TKConversion (TKX2 sh a) (TKX2 sh' a)++  ConvRR  :: TKConversion a b -> TKConversion (TKR2 n a) (TKR2 n b)+  ConvSS  :: TKConversion a b -> TKConversion (TKS2 sh a) (TKS2 sh b)+  ConvXX  :: TKConversion a b -> TKConversion (TKX2 sh a) (TKX2 sh b)+  ConvT2  :: TKConversion a a'+          -> TKConversion b b'+          -> TKConversion (TKProduct a b) (TKProduct a' b')++  Conv0X  :: SingletonTK a -> TKConversion a (TKX2 '[] a)+  ConvX0  :: TKConversion (TKX2 '[] a) a++  ConvNest :: SingletonTK (TKX2 sh a)+           -> TKConversion (TKX2 (sh ++ sh') a) (TKX2 sh (TKX2 sh' a))+  ConvUnnest :: TKConversion (TKX2 sh (TKX2 sh' a)) (TKX2 (sh ++ sh') a)++  ConvZip   :: SingletonTK a -> SingletonTK b+            -> TKConversion (TKProduct (TKX2 sh a) (TKX2 sh b))+                            (TKX2 sh (TKProduct a b))+  ConvUnzip :: SingletonTK a -> SingletonTK b+            -> TKConversion (TKX2 sh (TKProduct a b))+                            (TKProduct (TKX2 sh a) (TKX2 sh b))++deriving instance Show (TKConversion a b)++instance Category TKConversion where+  id = ConvId+  (.) = ConvCmp++convertSTK :: TKConversion a b -> SingletonTK a -> SingletonTK b+convertSTK = \cases+  ConvId astk -> astk+  (ConvCmp c1 c2) astk -> convertSTK c1 (convertSTK c2 astk)+  ConvRX (STKR n a) -> STKX (ssxReplicate n) a+  ConvSX (STKS sh a) -> STKX (ssxFromShX $ shxFromShS sh) a+  (ConvXR _stk) (STKX ssx a) -> STKR (ssxRank ssx) a+  ConvXS (STKX ssx a) -> STKS (shsFromSSX ssx) a+  (ConvXS' (FTKS sh _x)) (STKX _ssx2 a) -> STKS sh a+  (ConvXX' (FTKX shx _x)) (STKX _ssx2 a) -> STKX (ssxFromShX shx) a+  (ConvRR c) (STKR n a) -> STKR n (convertSTK c a)+  (ConvSS c) (STKS sh a) -> STKS sh (convertSTK c a)+  (ConvXX c) (STKX ssx a) -> STKX ssx (convertSTK c a)+  (ConvT2 c1 c2) (STKProduct stk1 stk2) ->+    STKProduct (convertSTK c1 stk1) (convertSTK c2 stk2)+  (Conv0X _stk) stk -> STKX ZKX stk+  ConvX0 (STKX ZKX stk) -> stk+  (ConvNest (STKX ssx x)) (STKX shsh' _x) ->+    STKX ssx (STKX (ssxDropSSX ssx shsh') x)+  ConvUnnest (STKX sh (STKX sh' x)) -> STKX (sh `ssxAppend` sh') x+  (ConvZip _ _) (STKProduct (STKX sh a1) (STKX _sh a2)) ->+    STKX sh (STKProduct a1 a2)+  (ConvUnzip _ _) (STKX sh (STKProduct a1 a2)) ->+    STKProduct (STKX sh a1) (STKX sh a2)++convertFTK :: TKConversion a b -> FullShapeTK a -> FullShapeTK b+convertFTK = \cases+  ConvId aftk -> aftk+  (ConvCmp c1 c2) aftk -> convertFTK c1 (convertFTK c2 aftk)+  ConvRX (FTKR shr a) -> FTKX (shxFromShR shr) a+  ConvSX (FTKS sh a) -> FTKX (shxFromShS sh) a+  (ConvXR _stk) (FTKX shx a) -> FTKR (shrFromShX shx) a+  ConvXS (FTKX shx a) -> FTKS (shsFromShX shx) a+  (ConvXS' ftk) _ -> ftk+  (ConvXX' ftk) _ -> ftk+  (ConvRR c) (FTKR shr a) -> FTKR shr (convertFTK c a)+  (ConvSS c) (FTKS sh a) -> FTKS sh (convertFTK c a)+  (ConvXX c) (FTKX shx a) -> FTKX shx (convertFTK c a)+  (ConvT2 c1 c2) (FTKProduct ftk1 ftk2) ->+    FTKProduct (convertFTK c1 ftk1) (convertFTK c2 ftk2)+  (Conv0X _stk) ftk -> FTKX ZSX ftk+  ConvX0 (FTKX ZSX ftk) -> ftk+  (ConvNest @_ @_ @sh' (STKX ssx _x)) (FTKX shsh' x) ->+    FTKX (shxTakeSSX (Proxy @sh') ssx shsh') (FTKX (shxDropSSX ssx shsh') x)+  ConvUnnest (FTKX sh (FTKX sh' x)) -> FTKX (sh `shxAppend` sh') x+  (ConvZip _ _) (FTKProduct (FTKX sh a1) (FTKX _sh a2)) ->+    FTKX sh (FTKProduct a1 a2)+  (ConvUnzip _ _) (FTKX sh (FTKProduct a1 a2)) ->+    FTKProduct (FTKX sh a1) (FTKX sh a2)++buildTKConversion :: SNat k -> FullShapeTK a+                  -> TKConversion a b+                  -> TKConversion (BuildTensorKind k a) (BuildTensorKind k b)+buildTKConversion k aftk c0 = case c0 of+  ConvId -> ConvId+  ConvCmp c1 c2 -> ConvCmp (buildTKConversion k (convertFTK c2 aftk) c1)+                           (buildTKConversion k aftk c2)+  ConvRX | FTKR @n shr xstk <- aftk+         , Refl <- lemRankReplicate (Proxy @n)+         , Refl <- lemRankReplicate (Proxy @(1 + n)) ->+    ConvCmp (ConvXX' (FTKX (SKnown k :$% shxFromShR shr) xstk)) ConvRX+  ConvSX -> ConvSX+  ConvXR stk -> ConvXR stk+  ConvXS -> ConvXS+  ConvXS' ftk -> ConvXS' (buildFTK k ftk)+  ConvXX' ftk -> ConvXX' (buildFTK k ftk)+  ConvRR c -> ConvRR c+  ConvSS c -> ConvSS c+  ConvXX c -> ConvXX c+  ConvT2 c1 c2 | FTKProduct ftk1 ftk2 <- aftk ->+    ConvT2 (buildTKConversion k ftk1 c1) (buildTKConversion k ftk2 c2)+  Conv0X _astk -> case aftk of+    FTKScalar -> ConvSX+    FTKR @n shr x | Refl <- lemRankReplicate (Proxy @n)+                  , Refl <- lemRankReplicate (Proxy @(1 + n)) ->+      ConvCmp (ConvXX (ConvXR (ftkToSTK x)))+              (ConvCmp (ConvNest (STKX (SKnown k :!% ZKX) (ftkToSTK x)))+                       (ConvCmp+                          (ConvXX' (FTKX (SKnown k :$% shxFromShR shr) x))+                          ConvRX))+    FTKS _sh x ->+      ConvCmp (ConvXX ConvXS)+              (ConvCmp (ConvNest (STKX (SKnown k :!% ZKX) (ftkToSTK x)))+                       ConvSX)+    FTKX _ssx x -> ConvNest (STKX (SKnown k :!% ZKX) (ftkToSTK x))+    FTKProduct aftk1 aftk2 ->+      buildTKConversion+        k aftk (ConvCmp (ConvZip (ftkToSTK aftk1) (ftkToSTK aftk2))+                        (ConvT2 (Conv0X (ftkToSTK aftk1))+                                (Conv0X (ftkToSTK aftk2))))+  ConvX0 -> case aftk of+    FTKX ZSX FTKScalar -> ConvXS+    FTKX ZSX (FTKR @n _n x) | Refl <- lemRankReplicate (Proxy @n) ->+      ConvCmp (ConvXR (ftkToSTK x))+              (ConvCmp ConvUnnest (ConvXX ConvRX))+    FTKX ZSX FTKS{} ->+      ConvCmp ConvXS+              (ConvCmp ConvUnnest (ConvXX ConvSX))+    FTKX ZSX FTKX{} -> ConvUnnest+    FTKX ZSX (FTKProduct aftk1 aftk2) ->+      buildTKConversion+        k aftk (ConvCmp (ConvT2 ConvX0 ConvX0)+                        (ConvUnzip (ftkToSTK aftk1) (ftkToSTK aftk2)))+  ConvNest (STKX sh x) -> ConvNest (STKX (SKnown k :!% sh) x)+  ConvUnnest -> ConvUnnest+  ConvZip astk1 astk2 -> ConvZip astk1 astk2+  ConvUnzip astk1 astk2 -> ConvUnzip astk1 astk2+++-- * Full shape tensor kind quasi-singletons++-- | Full shape tensor kind singleton type.+type role FullShapeTK nominal+data FullShapeTK y where+  FTKScalar :: GoodScalar r+            => FullShapeTK (TKScalar r)+  FTKR :: IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)+  FTKS :: ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)+  FTKX :: IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)+  FTKProduct :: FullShapeTK y -> FullShapeTK z+             -> FullShapeTK (TKProduct y z)++deriving instance Show (FullShapeTK y)+deriving instance Eq (FullShapeTK y)++-- | A plausible implementation of `testEquality` on `FullShapeTK`. It does not+-- take into account shape difference in ranked and mixed tensors+-- that `FullShapeTK`, but not `SingletonTK`, captures.+matchingFTK :: FullShapeTK y1 -> FullShapeTK y2 -> Maybe (y1 :~: y2)+matchingFTK ftk1 ftk2 = case (ftk1, ftk2) of+  (FTKScalar @r1, FTKScalar @r2)+    | Just Refl <- testEquality (typeRep @r1) (typeRep @r2) ->+      Just Refl+  (FTKR sh1 x1, FTKR sh2 x2)+    | Just Refl <- matchingFTK x1 x2+    , Just Refl <- testEquality (shrRank sh1) (shrRank sh2) ->  -- weaker!!!+      Just Refl+  (FTKS sh1 x1, FTKS sh2 x2)+    | Just Refl <- matchingFTK x1 x2+    , Just Refl <- testEquality sh1 sh2 ->+      Just Refl+  (FTKX sh1 x1, FTKX sh2 x2)+    | Just Refl <- matchingFTK x1 x2+    , Just Refl <- testEquality (ssxFromShX sh1) (ssxFromShX sh2) ->  -- !!!+      Just Refl+  (FTKProduct x1 y1, FTKProduct x2 y2)+    | Just Refl <- matchingFTK x1 x2, Just Refl <- matchingFTK y1 y2 ->+      Just Refl+  _ -> Nothing++-- | A conversion that is fully determined by the property that it+-- commutes with the `testEquality` implementations.+ftkToSTK :: FullShapeTK y -> SingletonTK y+ftkToSTK = \case+  FTKScalar -> STKScalar+  FTKR sh x -> STKR (shrRank sh) (ftkToSTK x)+  FTKS sh x -> STKS sh (ftkToSTK x)+  FTKX sh x -> STKX (ssxFromShX sh) (ftkToSTK x)+  FTKProduct ftk1 ftk2 -> STKProduct (ftkToSTK ftk1) (ftkToSTK ftk2)++ftkUnit :: FullShapeTK TKUnit+ftkUnit = FTKScalar++buildFTK :: SNat k -> FullShapeTK y -> FullShapeTK (BuildTensorKind k y)+buildFTK snat@SNat = \case+  FTKScalar -> FTKS (snat :$$ ZSS) FTKScalar+  FTKR sh x -> FTKR (sNatValue snat :$: sh) x+  FTKS sh x -> FTKS (snat :$$ sh) x+  FTKX sh x -> FTKX (SKnown snat :$% sh) x+  FTKProduct ftk1 ftk2 -> FTKProduct (buildFTK snat ftk1) (buildFTK snat ftk2)++razeFTK :: forall y k.+           SNat k -> SingletonTK y+        -> FullShapeTK (BuildTensorKind k y)+        -> FullShapeTK y+razeFTK snat@SNat stk ftk = case (stk, ftk) of+  (STKScalar, FTKS (_ :$$ ZSS) FTKScalar) -> FTKScalar+  (STKR{}, FTKR (_ :$: sh) x) -> FTKR sh x+  (STKR{}, FTKR ZSR _) -> error "razeFTK: impossible built tensor kind"+  (STKS{}, FTKS (_ :$$ sh) x) -> FTKS sh x+  (STKX{}, FTKX (_ :$% sh) x) -> FTKX sh x+  (STKProduct stk1 stk2, FTKProduct ftk1 ftk2) ->+    FTKProduct (razeFTK snat stk1 ftk1) (razeFTK snat stk2 ftk2)++adFTK :: FullShapeTK y+      -> FullShapeTK (ADTensorKind y)+adFTK = \case+  t@(FTKScalar @r) -> case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> t+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> t+      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $+           FTKScalar @Z1+  FTKR sh x -> FTKR sh $ adFTK x+  FTKS sh x -> FTKS sh $ adFTK x+  FTKX sh x -> FTKX sh $ adFTK x+  FTKProduct ftk1 ftk2 -> FTKProduct (adFTK ftk1) (adFTK ftk2)++-- A test whether the argument tensor collection is free+-- from any non-differentiable types, such as integers.+differentiableFTK :: FullShapeTK y -> Bool+differentiableFTK = \case+  FTKScalar @r -> case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> True+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> True+      _ -> False+  FTKR _ x -> differentiableFTK x+  FTKS _ x -> differentiableFTK x+  FTKX _ x -> differentiableFTK x+  FTKProduct ftk1 ftk2 -> differentiableFTK ftk1 && differentiableFTK ftk2++type role DummyDualTarget nominal+type DummyDualTarget :: Target+newtype DummyDualTarget y = DummyDualTarget (FullShapeTK y)
+ src/HordeAd/Core/Types.hs view
@@ -0,0 +1,735 @@+{-# LANGUAGE AllowAmbiguousTypes, DerivingVia, ImpredicativeTypes,+             UndecidableInstances, ViewPatterns #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+{-# OPTIONS_GHC -Wno-orphans #-}+-- | An assortment of type families and types fundamental for horde-ad.+module HordeAd.Core.Types+  ( -- * Re-exports and definitions to help express and manipulate type-level natural numbers+    SNat, pattern SNat, pattern SNat'+  , withSNat, sNatValue, proxyFromSNat, valueOf+    -- * Kinds of the parameterized types that determine the structure of a tensor+  , Target, TK (..), TKR, TKS, TKX, TKUnit+    -- * Some fundamental constraints and types related to tensors+  , GoodScalar, GoodScalarConstraint, Differentiable, IfDifferentiable(..)+  , BuildTensorKind, RazeTensorKind, ADTensorKind, ADTensorScalar+    -- * Type families that tensors belong to+  , IntOf, HFunOf, PrimalOf, DualOf, ShareOf, BoolOf+  , IxROf, IxSOf, IxXOf+    -- * The Z1 Num unit type and its instances+  , Z1(..)+    -- * Misc+  , Dict(..), IntegralH(..), RealFloatH(..), Boolean (..), EqH(..), OrdH(..)+  , backpermutePrefixList+  , toLinearIdxR, fromLinearIdxR, toLinearIdxS, fromLinearIdxS+  , toLinearIdxX, fromLinearIdxX+    -- * Feature requests for ox-arrays+  , Take, Drop, UnMapSucc+  , listsTake, listsDrop, listsSplitAt, ixrTake, ixrDrop, ixrSplitAt+  , shrTake, shrDrop, shrSplitAt+  , listrSplitAt, ixsDrop, shsTake, shsDrop+  , shxTake, shxDrop, ixxTake, ixxDrop'+  , listsTakeLen, listsDropLen, shsDropLen+  , permRInverse, ssxPermutePrefix, shxPermutePrefix+  , shCastSX+  , shsFromIxS, shsFromListS+  , normalizePermutationHack, backpermCycle, permCycle+  , permUnShift1+  , ssxTakeIx+  ) where++import Prelude++import Control.DeepSeq (NFData (..))+import Data.Array.Internal.RankedS qualified as RS+import Data.Boolean (Boolean (..))+import Data.Coerce (coerce)+import Data.Default+import Data.Functor.Const+import Data.Int (Int64)+import Data.Kind (Type)+import Data.List (dropWhileEnd, sort)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:) (Refl))+import Data.Vector.Storable qualified as V+import Foreign.C (CInt)+import Foreign.Storable (Storable (..))+import GHC.Exts (IsList (..))+import GHC.Generics (Generic)+import GHC.TypeLits+  ( KnownNat+  , Nat+  , SNat+  , fromSNat+  , pattern SNat+  , sameNat+  , type (+)+  , type (-)+  , withSomeSNat+  )+import System.Random+import Type.Reflection (Typeable)+import Unsafe.Coerce (unsafeCoerce)++import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (shxFromShS)+import Data.Array.Nested.Mixed qualified as Mixed+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation (DropLen, PermR, TakeLen)+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (Dict (..), Tail, fromSNat', unsafeCoerceRefl)+import Data.Array.Strided.Orthotope (NumElt (..), fromO, toO)++-- * Definitions to help express and manipulate type-level natural numbers++instance NFData (SNat n) where+  rnf _ = ()++withSNat :: Int -> (forall n. KnownNat n => (SNat n -> r)) -> r+withSNat i f = withSomeSNat (fromIntegral i) $ \case+  Just snat@SNat -> f snat+  Nothing -> error $ "withSNat: negative argument: " ++ show i++sNatValue :: forall n. SNat n -> Int+{-# INLINE sNatValue #-}+sNatValue = fromInteger . fromSNat++proxyFromSNat :: SNat n -> Proxy n+proxyFromSNat SNat = Proxy++{-# INLINE valueOf #-}+valueOf :: forall n r. (KnownNat n, Num r) => r+valueOf = fromInteger $ fromSNat (SNat @n)++pattern SNat' :: forall n m. KnownNat n => (KnownNat m, n ~ m) => SNat m+pattern SNat' <- (matchSNat (Proxy @n) -> Just (Refl :: n :~: m))+  where SNat' = SNat++matchSNat :: forall n m proxy. KnownNat n+          => proxy n -> SNat m -> Maybe (n :~: m)+matchSNat p m@SNat = sameNat p m+++-- * Kinds of the functors that determine the structure of a tensor type++-- | The parameterized carrier types of the algebras+-- that are the targets of interpretation -- of AST+-- and instantiation of the @Ops@ tensor classes. E.g., the type+-- that represents tensor kinds in concrete arrays is one such target.+-- AST itself is another. Dual numbers is yet another target.+type Target = TK -> Type++-- | The type of tensor kinds constraining the shapes+-- of (nested pairs of) tensors.+type data TK =+    TKScalar Type+  | TKR2 Nat TK+  | TKS2 [Nat] TK+  | TKX2 [Maybe Nat] TK+  | TKProduct TK TK++type TKR n r = TKR2 n (TKScalar r)+type TKS sh r = TKS2 sh (TKScalar r)+type TKX sh r = TKX2 sh (TKScalar r)++type TKUnit = TKScalar Z1+++-- * Some fundamental constraints and types++type GoodScalarConstraint r =+  ( Show r, Ord r, Num r, Typeable r, IfDifferentiable r, Default r+  , NFData r, Nested.PrimElt r, Nested.KnownElt r, Nested.NumElt r+  , forall sh. Show (Nested.Mixed sh r), forall sh. Ord (Nested.Mixed sh r)+  , forall sh. NFData (Nested.Mixed sh r))++-- Attempted optimization via storing one pointer to a class dictionary+-- in existential datatypes instead of six pointers. No effect, strangely.+-- As a side effect, this avoids ImpredicativeTypes.+-- Also, the constraint can be represented by a single Dict.+--+-- | The constraint that signifies a scalar type, e.g., a float or an integers,+-- is a well-behaved cell content of tensors supported by horde-ad.+class GoodScalarConstraint r => GoodScalar r+instance GoodScalarConstraint r => GoodScalar r++-- | The constraint for scalars that can be non-trivially differentiated,+-- e.g., floating numbers, but not integers.+type Differentiable r =+  (RealFloatH r, Nested.FloatElt r, RealFrac r, RealFloat r, Random r)++-- We white-list all types on which we permit differentiation (e.g., SGD)+-- to work. This is for technical typing purposes and imposes updates+-- (and buggy omissions) when new scalar types are added, but it has+-- the advantage of giving more control and visiblity. Sadly the list+-- is repeated in several other places.+class IfDifferentiable r where+  ifDifferentiable :: (Differentiable r => a) -> a -> a++instance {-# OVERLAPPABLE #-} IfDifferentiable r where+  ifDifferentiable _ a = a++-- The white-listed differentiable types.+instance IfDifferentiable Double where+  ifDifferentiable ra _ = ra+instance IfDifferentiable Float where+  ifDifferentiable ra _ = ra++type family BuildTensorKind k tk where+  BuildTensorKind k (TKScalar r) = TKS '[k] r+  BuildTensorKind k (TKR2 n r) = TKR2 (1 + n) r+  BuildTensorKind k (TKS2 sh r) = TKS2 (k : sh) r+  BuildTensorKind k (TKX2 sh r) = TKX2 (Just k : sh) r+  BuildTensorKind k (TKProduct y z) =+    TKProduct (BuildTensorKind k y) (BuildTensorKind k z)++-- This is an inverse of BuildTensorKind.+-- This could be more efficient+--   RazeTensorKind (TKS2 '[m] (TKScalar r)) = TKScalar r+-- but then we'd lose the simplifying property that razing does not+-- change the tensor kind variant, which is important, e.g.,+-- when rewriting AstFromS t and trying to use AstFromS of the razed t.+type family RazeTensorKind tk where+  RazeTensorKind (TKR2 n r) = TKR2 (n - 1) r+  RazeTensorKind (TKS2 sh r) = TKS2 (Tail sh) r+  RazeTensorKind (TKX2 sh r) = TKX2 (Tail sh) r+  RazeTensorKind (TKProduct y z) =+    TKProduct (RazeTensorKind y) (RazeTensorKind z)++type family ADTensorKind tk where+  ADTensorKind (TKScalar r) = TKScalar (ADTensorScalar r)+  ADTensorKind (TKR2 n r) = TKR2 n (ADTensorKind r)+  ADTensorKind (TKS2 sh r) = TKS2 sh (ADTensorKind r)+  ADTensorKind (TKX2 sh r) = TKX2 sh (ADTensorKind r)+  ADTensorKind (TKProduct y z) =+    TKProduct (ADTensorKind y) (ADTensorKind z)++type family ADTensorScalar r where+  ADTensorScalar Double = Double+  ADTensorScalar Float = Float+  ADTensorScalar t = Z1+++-- * Type families that tensors belong to++-- This is used only in indexing and similar contexts.+-- If used as size or shape, it would give more expressiveness,+-- but would lead to irregular tensors, especially after vectorization,+-- and would prevent statically known shapes.+-- | The (arbitrariliy chosen) type of scalars to be used for indexing.+type IntOf (f :: Target) = PrimalOf f (TKScalar Int64)++-- | The type family is defined in order to give a special instance+-- for AST that preserves sharing and, even more importantly, keeps+-- the computation of dervative functions lazy. See the definition+-- of 'HordeAd.Core.Ast.AstLambda' and the code that processes it,+-- maintaining laziness.+--+-- The type family can't easily be made injective, because the @ADVal f@+-- instance is independent of @f@.+type family HFunOf (f :: Target)+                   (x :: TK) (z :: TK) :: Type++type family PrimalOf (f :: Target) :: Target++type family DualOf (f :: Target) :: Target++type family ShareOf (f :: Target) :: Target++type family BoolOf (t :: Target) :: Type++-- | Ranked index, that is, a sized list of individual indices into+-- ranked tensors.+--+-- Thanks to the OverloadedLists mechanism, values of this type+-- and of 'IxSOf' and 'IxXOf' can be written using the normal+-- list notation. However, such values, if not+-- explicitly typed, do not inform the compiler about the length+-- of the list until runtime. That means that some errors are hidden+-- and also extra type applications may be needed to satisfy the compiler.+-- Therefore, there is a real trade-off between @[2]@ and @(2 :.: ZIR).+type IxROf (f :: Target) n = IxR n (IntOf f)++-- | Shaped index, that is, a sized list of individual indices into+-- shaped tensors.+--+-- The values of this type are additionally decorated with a shape type @sh@,+-- which indirectly determines in what tensor operations+-- (e.g., what indexing) the values can appear. The type has no direct+-- relation to the runtime payload of the list, except that the list+-- has the same length as the shape. In particular, the individual+-- indices can be "out of bounds" with respect to the shape type.+type IxSOf (f :: Target) (sh :: [Nat]) = IxS sh (IntOf f)++-- | Mixed index, that is, a sized list of individual indices+-- into tensors with mixed ranked-shaped typing.+type IxXOf (f :: Target) (sh :: [Maybe Nat]) = IxX sh (IntOf f)+++-- * The Z1 Num unit type and its instances++data Z1 = Z1+ deriving (Eq, Ord, Show)++instance NFData Z1 where+  rnf Z1 = ()++instance Storable Z1 where+  sizeOf _ = 0+  alignment _ = 1+  peek _ = return Z1+  poke _ _ = return ()++instance Num Z1 where+  Z1 + Z1 = Z1+  Z1 * Z1 = Z1+  negate Z1 = Z1+  abs Z1 = Z1+  signum Z1 = Z1+  fromInteger _ = Z1++instance Default Z1 where+  def = Z1++instance Nested.PrimElt Z1+newtype instance Mixed.Mixed sh Z1 = M_NilZ1 (Mixed.Mixed sh (Mixed.Primitive Z1)) deriving (Eq, Ord, Generic)  -- no content, orthotope optimises this (via Vector)+newtype instance Mixed.MixedVecs s sh Z1 = MV_NilZ1 (V.MVector s Z1)  -- no content, MVector optimises this+deriving via Nested.Primitive Z1 instance Nested.Elt Z1+deriving via Nested.Primitive Z1 instance Nested.KnownElt Z1++instance NumElt Z1 where+  numEltAdd _ arr1 _arr2 = arr1+  numEltSub _ arr1 _arr2 = arr1+  numEltMul _ arr1 _arr2 = arr1+  numEltNeg _ arr = arr+  numEltAbs _ arr = arr+  numEltSignum _ arr = arr+  numEltSum1Inner _ arr = fromO (RS.index (toO arr) 0)+  numEltProduct1Inner _ arr = fromO (RS.index (toO arr) 0)+  numEltSumFull _ _arr = Z1+  numEltProductFull _ _arr = Z1+  numEltMinIndex snat _arr = replicate (sNatValue snat) 0+  numEltMaxIndex snat _arr = replicate (sNatValue snat) 0+  numEltDotprodInner _ arr1 _arr2 = fromO (RS.index (toO arr1) 0)+++-- * Misc++-- TODO: move all these somewhere++-- | A variant of Integral for horde-ad without the problematic operations:+-- 'mod' and 'rem' that are very slow, pair-producing operations that don't fit+-- the AST GADT format and 'toInteger' that doesn't make sense with AST+-- without an environment to look up variables in.+class Num a => IntegralH a where+  quotH, remH :: a -> a -> a++instance IntegralH Int64 where+  quotH = quot+  remH = rem++instance IntegralH CInt where+  quotH = quot+  remH = rem++-- | The standard 'RealFloat' brings in a lot of operations we are not+-- interested in and 'RealFrac' and 'Real' that we can't faithfully implement+-- (e.g., 'toRational' doesn't make sense with AST without an environment).+class (Floating a) => RealFloatH a where+  atan2H :: a -> a -> a++instance RealFloatH Float where+  atan2H = atan2++instance RealFloatH Double where+  atan2H = atan2++{- TODO: these would be better, but everything then overlaps with everything:+instance {-# OVERLAPPABLE #-} Integral r => IntegralH r where+  quotH = quot+  remH = rem++instance {-# OVERLAPPABLE #-} (Floating r, RealFloat r) => RealFloatH r where+  atan2H = atan2+-}++infix 4 ==., /=.+class Boolean (BoolOf f) => EqH (f :: Target) y where+  (==.), (/=.) :: f y -> f y -> BoolOf f+  u /=. v = notB (u ==. v)++infix 4 <., <=., >=., >.+class Boolean (BoolOf f) => OrdH (f :: Target) y where+  (<.), (<=.), (>.), (>=.) :: f y -> f y -> BoolOf f+  u <. v = notB (v <=. u)+  u >. v = notB (u <=. v)+  u >=. v = v <=. u++backpermutePrefixList :: PermR -> [i] -> [i]+backpermutePrefixList p l = map (l !!) p ++ drop (length p) l++-- I can't switch to ixxFromLinear and ixxToLinear from ox-arrays+-- even just because IntegralH is not available in ox-arrays.+--+-- | Given a multidimensional index, get the corresponding linear+-- index into the buffer. Note that the index doesn't need to be pointing+-- at a scalar. It may point at the start of a larger tensor instead.+--+-- If any of the dimensions is 0 or if rank is 0, the result will be 0,+-- which is fine, that's pointing at the start of the empty buffer.+-- Note that the resulting 0 may be a complex term.+--+-- Warning: @fromInteger@ of type @j@ cannot be used.+toLinearIdxR :: forall m n j. Num j+             => (Int -> j) -> ShR (m + n) Int -> IxR m j -> j+toLinearIdxR fromInt = \sh idx -> go sh idx (fromInt 0)+  where+    -- Additional argument: index, in the @m - m1@ dimensional array so far,+    -- of the @m - m1 + n@ dimensional tensor pointed to by the current+    -- @m - m1@ dimensional index prefix.+    go :: ShR (m1 + n) Int -> IxR m1 j -> j -> j+    go sh ZIR tensidx = fromInt (shrSize sh) * tensidx+    go (n :$: sh) (i :.: idx) tensidx = go sh idx (fromInt n * tensidx + i)++-- | Given a linear index into the buffer, get the corresponding+-- multidimensional index. Here we require an index pointing at a scalar.+--+-- If any of the dimensions is 0, the linear index has to be 0+-- (which we can't assert, because j may be a term and so == lies)+-- and a fake index with correct length but lots of zeroes is produced,+-- because it doesn't matter, because it's going to point at the start+-- of the empty buffer anyway.+--+-- Warning: @fromInteger@ of type @j@ cannot be used.+fromLinearIdxR :: forall n j. IntegralH j+               => (Int -> j) -> ShR n Int -> j -> IxR n j+fromLinearIdxR fromInt = \sh lin -> snd (go sh lin)+  where+    -- Returns (linear index into array of sub-tensors,+    -- multi-index within sub-tensor).+    go :: ShR n1 Int -> j -> (j, IxR n1 j)+    go ZSR n = (n, ZIR)+    go (k :$: sh) _ | signum k == 0 =+      (fromInt 0, fromInt 0 :.: zeroOfR fromInt sh)+    go (n :$: sh) lin =+      let (tensLin, idxInTens) = go sh lin+          tensLin' = tensLin `quotH` fromInt n+          i = tensLin `remH` fromInt n+      in (tensLin', i :.: idxInTens)++-- | The zero index in this shape (not dependent on the actual integers).+zeroOfR :: Num j => (Int -> j) -> ShR n i -> IxR n j+zeroOfR _ ZSR = ZIR+zeroOfR fromInt (_ :$: sh) = fromInt 0 :.: zeroOfR fromInt sh++-- | Given a multidimensional index, get the corresponding linear+-- index into the buffer. Note that the index doesn't need to be pointing+-- at a scalar. It may point at the start of a larger tensor instead.+--+-- If any of the dimensions is 0 or if rank is 0, the result will be 0,+-- which is fine, that's pointing at the start of the empty buffer.+-- Note that the resulting 0 may be a complex term.+--+-- Warning: @fromInteger@ of type @j@ cannot be used.+toLinearIdxS :: forall sh1 sh2 j. Num j+             => (Int -> j) -> ShS (sh1 ++ sh2) -> IxS sh1 j -> j+toLinearIdxS fromInt = \sh idx -> go sh idx (fromInt 0)+  where+    -- Additional argument: index, in the @m - m1@ dimensional array so far,+    -- of the @m - m1 + n@ dimensional tensor pointed to by the current+    -- @m - m1@ dimensional index prefix.+    go :: forall sh3. ShS (sh3 ++ sh2) -> IxS sh3 j -> j -> j+    go sh ZIS tensidx = fromInt (shsSize sh) * tensidx+    go ((:$$) n sh) (i :.$ idx) tensidx =+      go sh idx (fromInt (sNatValue n) * tensidx + i)+    go _ _ _ = error "toLinearIdxS: impossible pattern needlessly required"++-- | Given a linear index into the buffer, get the corresponding+-- multidimensional index. Here we require an index pointing at a scalar.+--+-- If any of the dimensions is 0, the linear index has to be 0+-- (which we can't assert, because j may be a term and so == lies)+-- and a fake index with correct length but lots of zeroes is produced,+-- because it doesn't matter, because it's going to point at the start+-- of the empty buffer anyway.+--+-- Warning: @fromInteger@ of type @j@ cannot be used.+fromLinearIdxS :: forall sh j. IntegralH j+               => (Int -> j) -> ShS sh -> j -> IxS sh j+fromLinearIdxS fromInt = \sh lin -> snd (go sh lin)+  where+    -- Returns (linear index into array of sub-tensors,+    -- multi-index within sub-tensor).+    go :: ShS sh1 -> j -> (j, IxS sh1 j)+    go ZSS n = (n, ZIS)+    go ((:$$) k sh) _ | sNatValue k == 0 =+      (fromInt 0, fromInt 0 :.$ zeroOfS fromInt sh)+    go ((:$$) n sh) lin =+      let (tensLin, idxInTens) = go sh lin+          tensLin' = tensLin `quotH` fromInt (sNatValue n)+          i = tensLin `remH` fromInt (sNatValue n)+      in (tensLin', i :.$ idxInTens)++-- | The zero index in this shape (not dependent on the actual integers).+zeroOfS :: Num j => (Int -> j) -> ShS sh -> IxS sh j+zeroOfS _ ZSS = ZIS+zeroOfS fromInt ((:$$) _ sh) = fromInt 0 :.$ zeroOfS fromInt sh++toLinearIdxX :: forall sh1 sh2 j. Num j+             => (Int -> j) -> IShX (sh1 ++ sh2) -> IxX sh1 j -> j+toLinearIdxX fromInt = \sh idx -> go sh idx (fromInt 0)+  where+    -- Additional argument: index, in the @m - m1@ dimensional array so far,+    -- of the @m - m1 + n@ dimensional tensor pointed to by the current+    -- @m - m1@ dimensional index prefix.+    go :: forall sh3. IShX (sh3 ++ sh2) -> IxX sh3 j -> j -> j+    go sh ZIX tensidx = fromInt (shxSize sh) * tensidx+    go ((:$%) n sh) (i :.% idx) tensidx =+      go sh idx (fromInt (fromSMayNat' n) * tensidx + i)+    go _ _ _ = error "toLinearIdxX: impossible pattern needlessly required"++fromLinearIdxX :: forall sh j. IntegralH j+               => (Int -> j) -> IShX sh -> j -> IxX sh j+fromLinearIdxX fromInt = \sh lin -> snd (go sh lin)+  where+    -- Returns (linear index into array of sub-tensors,+    -- multi-index within sub-tensor).+    go :: IShX sh1 -> j -> (j, IxX sh1 j)+    go ZSX n = (n, ZIX)+    go ((:$%) k sh) _ | fromSMayNat' k == 0 =+      (fromInt 0, fromInt 0 :.% zeroOfX fromInt sh)+    go ((:$%) n sh) lin =+      let (tensLin, idxInTens) = go sh lin+          tensLin' = tensLin `quotH` fromInt (fromSMayNat' n)+          i = tensLin `remH` fromInt (fromSMayNat' n)+      in (tensLin', i :.% idxInTens)++-- | The zero index in this shape (not dependent on the actual integers).+zeroOfX :: Num j => (Int -> j) -> IShX sh -> IxX sh j+zeroOfX _ ZSX = ZIX+zeroOfX fromInt ((:$%) _ sh) = fromInt 0 :.% zeroOfX fromInt sh+++-- * Shopping list for ox-arrays++-- All of the below should have better types and/or implementations,+-- just as in ox-arrays, and should have variants for all 10 kinds+-- of shape things.++-- ** Conversions and related++-- The constraint is erroneously reported as redundant.+shCastSX :: Rank sh ~ Rank sh' => StaticShX sh' -> ShS sh -> IShX sh'+shCastSX ssx sh = shxCast' ssx (shxFromShS sh)++-- ** Permutation-related operations++-- - Permutation.permInverse for ShS would be helpful in addition to+--   for StaticShX (the proof does not convert (easily)).+--   Though, frankly, the proof is often useless,+--   due to how bad GHC is at reasoning (no (++) congruence, no (:~:)+--   transitivity, etc., see astGatherCase.AstTransposeS+--   and astTransposeAsGatherS), so it's easier to postulate the thing+--   GHC needs in one or two clauses than to write a dozen bread crumbs+--   to lead GHC to grudgingly use the proof.++permRInverse :: PermR -> PermR+permRInverse perm = map snd $ sort $ zip perm [0 .. length perm - 1]++-- TODO: port to shaped permutations and then remove the Hack suffix+normalizePermutationHack :: Permutation.PermR -> Permutation.PermR+normalizePermutationHack perm =+  map fst $ dropWhileEnd (uncurry (==)) $ zip perm [0 ..]++-- TODO: can this be defined for Permutation.Perm using @Mod@?+-- A representation of a cycle backpermutation that moves elements+-- to indexes one less (the the left, to the back).+backpermCycle :: Int -> Permutation.PermR+backpermCycle 0 = []+backpermCycle 1 = []+backpermCycle n = [k `mod` n | k <- [1 .. n]]++-- TODO: can this be defined for Permutation.Perm using @Mod@?+-- A representation of a cycle permutation that is reverse to @backpermCycle@.+permCycle :: Int -> Permutation.PermR+permCycle 0 = []+permCycle 1 = []+permCycle n = [k `mod` n | k <- [-1, 0 .. n - 2]]++type family UnMapSucc is where+  UnMapSucc '[] = '[]+  UnMapSucc (i : is) = i - 1 : UnMapSucc is++-- The inverse of permShift1. Morally:+-- permUnShift1 :: Permutation.Perm (0 : Permutation.MapSucc l)+--              -> Permutation.Perm l+permUnShift1 :: Permutation.Perm (0 : l)+             -> Permutation.Perm (UnMapSucc l)+permUnShift1 (Permutation.PCons _ permRest) =+  Permutation.permFromList+    (permUnMapSucc (Permutation.permToList' permRest)) unsafeCoerce+ where+  permUnMapSucc :: [Int] -> [Int]+  permUnMapSucc [] = []+  permUnMapSucc (i : ns) = i - 1 : permUnMapSucc ns++-- TODO:+_withPermShift1 :: forall is r. -- Permutation.IsPermutation is+                  Permutation.Perm is+               -> (Permutation.IsPermutation (0 : Permutation.MapSucc is) =>+                   Permutation.Perm (0 : Permutation.MapSucc is) -> r)+               -> r+_withPermShift1 _perm _f = undefined  -- f (Permutation.permShift1 perm)++ssxPermutePrefix :: Permutation.Perm is -> StaticShX sh+                 -> StaticShX (Permutation.PermutePrefix is sh)+ssxPermutePrefix = undefined++shxPermutePrefix :: Permutation.Perm is -> ShX sh i+                 -> ShX (Permutation.PermutePrefix is sh) i+shxPermutePrefix = undefined++-- ** Takes and drops; this is postponed until we decide how to handle this; in particular, IIRC, unary nats for ranks is a pre-requisite for the most promising approach. There may be an ox-arrays branch for that. However, this requires lots of effort, so it's probably future work.++type family Take (n :: Nat) (xs :: [k]) :: [k] where+  Take 0 xs = '[]+  Take n (x ': xs) = x ': Take (n - 1) xs++type family Drop (n :: Nat) (xs :: [k]) :: [k] where+  Drop 0 xs = xs+  Drop n (x ': xs) = Drop (n - 1) xs++listsTake :: forall len sh i.+             (KnownShS sh, KnownNat len, KnownShS (Take len sh))+          => ListS sh (Const i) -> ListS (Take len sh) (Const i)+listsTake ix = fromList $ take (valueOf @len) $ toList ix++listsDrop :: forall len sh i.+             (KnownShS sh, KnownNat len, KnownShS (Drop len sh))+          => ListS sh (Const i) -> ListS (Drop len sh) (Const i)+listsDrop ix = fromList $ drop (valueOf @len) $ toList ix++listsSplitAt+  :: (KnownShS sh, KnownNat len, KnownShS (Drop len sh), KnownShS (Take len sh))+  => ListS sh (Const i)+  -> (ListS (Take len sh) (Const i), ListS (Drop len sh) (Const i))+listsSplitAt ix = (listsTake ix, listsDrop ix)++ixrTake :: forall m n i. (KnownNat m, KnownNat n)+        => IxR (m + n) i -> IxR m i+ixrTake (IxR ix) = IxR $ listrTake ix++ixrDrop :: forall m n i. (KnownNat m, KnownNat n)+        => IxR (m + n) i -> IxR n i+ixrDrop (IxR ix) = IxR $ listrDrop ix++ixrSplitAt :: (KnownNat m, KnownNat n)+           => IxR (m + n) i -> (IxR m i, IxR n i)+ixrSplitAt ix = (ixrTake ix, ixrDrop ix)++shrTake :: forall m n i. (KnownNat n, KnownNat m)+        => ShR (m + n) i -> ShR m i+shrTake (ShR ix) = ShR $ listrTake ix++shrDrop :: forall m n i. (KnownNat m, KnownNat n)+          => ShR (m + n) i -> ShR n i+shrDrop (ShR ix) = ShR $ listrDrop ix++shrSplitAt :: (KnownNat m, KnownNat n)+           => ShR (m + n) i -> (ShR m i, ShR n i)+shrSplitAt ix = (shrTake ix, shrDrop ix)++listrTake :: forall len n i. (KnownNat n, KnownNat len)+          => ListR (len + n) i -> ListR len i+listrTake ix = fromList $ take (valueOf @len) $ toList ix++listrDrop :: forall len n i. (KnownNat len, KnownNat n)+          => ListR (len + n) i -> ListR n i+listrDrop ix = fromList $ drop (valueOf @len) $ toList ix++listrSplitAt :: (KnownNat m, KnownNat n)+             => ListR (m + n) i -> (ListR m i, ListR n i)+listrSplitAt ix = (listrTake ix, listrDrop ix)++ixsDrop :: forall len sh i. (KnownShS sh, KnownNat len, KnownShS (Drop len sh))+        => IxS sh i -> IxS (Drop len sh) i+ixsDrop (IxS ix) = IxS $ listsDrop ix++-- TODO+shsTake :: forall len sh. (KnownNat len, KnownShS sh)+        => ShS sh -> ShS (Take len sh)+shsTake sh0 = fromList2 $ take (valueOf @len) $ toList sh0+ where+  fromList2 topl = ShS (go (knownShS @sh) topl)+    where  -- TODO: induction over (unary) SNat?+      go :: forall sh'. ShS sh' -> [Int] -> ListS (Take len sh') SNat+      go _ [] = gcastWith (unsafeCoerceRefl :: len :~: 0) $ gcastWith (unsafeCoerceRefl :: sh' :~: '[]) ZS+      go (sn :$$ sh) (i : is)+        | i == fromSNat' sn = unsafeCoerce $ sn ::$ go sh is+        | otherwise = error $ "shsTake: Value does not match typing (type says "+                                ++ show (fromSNat' sn) ++ ", list contains " ++ show i ++ ")"+      go _ _ = error $ "shsTake: Mismatched list length (type says "+                         ++ show (shsLength (knownShS @sh)) ++ ", list has length "+                         ++ show (length topl) ++ ")"++-- TODO+shsDrop :: forall len sh. (KnownNat len, KnownShS sh)+        => ShS sh -> ShS (Drop len sh)+shsDrop sh0 = fromList2 $ drop (valueOf @len) $ toList sh0+ where+  fromList2 topl = ShS (go (knownShS @sh) $ replicate (valueOf @len) (-1) ++ topl)+    where  -- TODO: induction over (unary) SNat?+      go :: forall sh'. ShS sh' -> [Int] -> ListS (Drop len sh') SNat+      go _ [] = gcastWith (unsafeCoerceRefl :: len :~: 0) $ gcastWith (unsafeCoerceRefl :: sh' :~: '[]) ZS+      go (sn :$$ sh) (i : is)+        | i == -1 = unsafeCoerce $ go sh is+        | i == fromSNat' sn = unsafeCoerce $ sn ::$ go sh is+        | otherwise = error $ "shsDrop: Value does not match typing (type says "+                                ++ show (fromSNat' sn) ++ ", list contains " ++ show i ++ ")"+      go _ _ = error $ "shsDrop: Mismatched list length (type says "+                         ++ show (shsLength (knownShS @sh)) ++ ", list has length "+                         ++ show (length topl) ++ ")"++shxTake :: forall len sh. (KnownNat len, KnownShX sh, KnownShX (Take len sh))+        => IShX sh -> IShX (Take len sh)+shxTake sh0 = fromList $ take (valueOf @len) $ toList sh0++shxDrop :: forall len sh. (KnownNat len, KnownShX sh, KnownShX (Drop len sh))+        => IShX sh -> IShX (Drop len sh)+shxDrop sh0 = fromList $ drop (valueOf @len) $ toList sh0++ixxTake :: forall len sh i. (KnownNat len, KnownShX sh, KnownShX (Take len sh))+        => IxX sh i -> IxX (Take len sh) i+ixxTake sh0 = fromList $ take (valueOf @len) $ toList sh0++ixxDrop' :: forall len sh i. (KnownNat len, KnownShX sh, KnownShX (Drop len sh))+         => IxX sh i -> IxX (Drop len sh) i+ixxDrop' sh0 = fromList $ drop (valueOf @len) $ toList sh0++listsTakeLen :: forall f g sh1 sh2.+                ListS sh1 f -> ListS sh2 g -> ListS (TakeLen sh1 sh2) g+listsTakeLen ZS _ = ZS+listsTakeLen (_ ::$ sh1) (n ::$ sh2) = n ::$ listsTakeLen sh1 sh2+listsTakeLen (_ ::$ _) ZS = error "listsTakeLen: list too short"++listsDropLen :: forall f g sh1 sh2.+                ListS sh1 f -> ListS sh2 g -> ListS (DropLen sh1 sh2) g+listsDropLen ZS sh = sh+listsDropLen (_ ::$ sh1) (_ ::$ sh2) = listsDropLen sh1 sh2+listsDropLen (_ ::$ _) ZS = error "listsDropLen: list too short"++shsDropLen :: Permutation.Perm is -> ShS sh -> ShS (DropLen is sh)+shsDropLen = coerce (listsDropLenPerm @SNat)++-- This is only needed as a workaround for other ops not provided.++listxTake :: forall f g sh sh'. ListX (sh ++ sh') f -> ListX sh g -> ListX sh f+listxTake _ ZX = ZX+listxTake (i ::% long') ((::%) @_ @_ @sh2 _ short) =+  i ::% listxTake @f @g @sh2 @sh' long' short++ssxTakeIx :: forall sh sh' i. StaticShX (sh ++ sh') -> IxX sh i -> StaticShX sh+ssxTakeIx = coerce (listxTake @(Nested.SMayNat () SNat) @(Const i) @_ @sh')
+ src/HordeAd/Core/Unwind.hs view
@@ -0,0 +1,531 @@+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+-- | Types and functions needed to define general tensor operations+-- that work for any tensor kind, including nested (product) arrays+-- and an assortment of such operations.+module HordeAd.Core.Unwind+  ( addTarget, multTarget, sum0Target, dot0Target+  , replTarget, defTarget, concreteTarget+  , toADTensorKindShared, fromADTensorKindShared+  ) where++import Prelude++import Data.Default+import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))+import GHC.TypeLits (type (+))+import Type.Reflection (typeRep)++import Data.Array.Nested (MapJust, Replicate, type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.OpsTensor++-- * Winding and unwinding++-- | This captures the normal form of type family UnWind and also+-- corresponds to the portion of ox-arrays that has Num defined.+type role RepW nominal nominal+data RepW target y where+  WTKScalar :: GoodScalar r+            => target (TKScalar r)+            -> RepW target (TKScalar r)+  WTKR :: GoodScalar r+       => target (TKR n r)+       -> RepW target (TKR n r)+  WTKS :: GoodScalar r+       => target (TKS sh r)+       -> RepW target (TKS sh r)+  WTKX :: GoodScalar r+       => target (TKX sh r)+       -> RepW target (TKX sh r)+  WTKProduct :: RepW target x -> RepW target z+             -> RepW target (TKProduct x z)++-- | This captures the normal form of type family UnWind for full shape+-- singletons.+type role FullShapeTKW nominal+data FullShapeTKW y where+  WFTKScalar :: GoodScalar r+             => FullShapeTKW (TKScalar r)+  WFTKR :: GoodScalar r+        => IShR n -> FullShapeTKW (TKR n r)+  WFTKS :: GoodScalar r+        => ShS sh -> FullShapeTKW (TKS sh r)+  WFTKX :: GoodScalar r+        => IShX sh -> FullShapeTKW (TKX sh r)+  WFTKProduct :: FullShapeTKW y -> FullShapeTKW z+              -> FullShapeTKW (TKProduct y z)++addRepW :: forall y target. BaseTensor target+        => RepW target y -> RepW target y -> RepW target y+addRepW a b = case (a, b) of+  (WTKScalar ta, WTKScalar tb) -> WTKScalar $ ta + tb+  (WTKR ta, WTKR tb) -> WTKR $ ta + tb+  (WTKS ta, WTKS tb) -> WTKS $ ta + tb+  (WTKX ta, WTKX tb) -> WTKX $ ta + tb+  (WTKProduct ta1 ta2, WTKProduct tb1 tb2) ->+    WTKProduct (addRepW ta1 tb1) (addRepW ta2 tb2)++multRepW :: forall y target. BaseTensor target+         => RepW target y -> RepW target y -> RepW target y+multRepW a b = case (a, b) of+  (WTKScalar ta, WTKScalar tb) -> WTKScalar $ ta * tb+  (WTKR ta, WTKR tb) -> WTKR $ ta * tb+  (WTKS ta, WTKS tb) -> WTKS $ ta * tb+  (WTKX ta, WTKX tb) -> WTKX $ ta * tb+  (WTKProduct ta1 ta2, WTKProduct tb1 tb2) ->+    WTKProduct (multRepW ta1 tb1) (multRepW ta2 tb2)++sum0RepW :: forall y target. (BaseTensor target, ConvertTensor target)+         => FullShapeTKW y -> RepW target y+         -> target (TKScalar Double)+sum0RepW ftk a = case (ftk, a) of+  (_, WTKScalar @r ta) ->+    ifDifferentiable @r (kcast ta) 0+  (WFTKR sh, WTKR @r ta) | SNat <- shrRank sh ->+    ifDifferentiable @r (kcast $ kfromR $ rsum0 ta) 0+  (WFTKS sh, WTKS @r ta) ->+    withKnownShS sh $+    ifDifferentiable @r (kcast $ kfromS $ ssum0 ta) 0+  (WFTKX sh, WTKX @r ta) ->+    withKnownShX (ssxFromShX sh) $+    ifDifferentiable @r (kcast $ kfromX $ xsum0 ta) 0+  (WFTKProduct ftk1 ftk2, WTKProduct ta1 ta2) ->+    sum0RepW ftk1 ta1 + sum0RepW ftk2 ta2++dot0RepW :: forall y target. (BaseTensor target, ConvertTensor target)+         => FullShapeTKW y -> RepW target y -> RepW target y+         -> target (TKScalar Double)+dot0RepW ftk a b = case (ftk, a, b) of+  (_, WTKScalar @r ta, WTKScalar tb) ->+    ifDifferentiable @r (kcast $ ta * tb) 0+  (WFTKR sh, WTKR @r ta, WTKR tb) | SNat <- shrRank sh ->+    ifDifferentiable @r (kcast $ kfromR $ rdot0 ta tb) 0+  (WFTKS sh, WTKS @r ta, WTKS tb) ->+    withKnownShS sh $+    ifDifferentiable @r (kcast $ kfromS $ sdot0 ta tb) 0+  (WFTKX sh, WTKX @r ta, WTKX tb) ->+    withKnownShX (ssxFromShX sh) $+    ifDifferentiable @r (kcast $ kfromX $ xdot0 ta tb) 0+  (WFTKProduct ftk1 ftk2, WTKProduct ta1 ta2, WTKProduct tb1 tb2) ->+    dot0RepW ftk1 ta1 tb1 + dot0RepW ftk2 ta2 tb2++replRepW :: forall y target. BaseTensor target+         => (forall r. GoodScalar r => r)+         -> FullShapeTKW y -> RepW target y+replRepW r = \case+  WFTKScalar -> WTKScalar $ kconcrete r+  WFTKR sh -> WTKR $ rrepl sh r+  WFTKS sh -> WTKS $ sconcrete $ Nested.sreplicateScal sh r+  WFTKX sh -> WTKX $ xrepl sh r+  WFTKProduct ftk1 ftk2 ->+    WTKProduct (replRepW r ftk1) (replRepW r ftk2)++defRepW :: forall y target. BaseTensor target+        => FullShapeTKW y -> RepW target y+defRepW = \case+  WFTKScalar -> WTKScalar $ kconcrete def+  WFTKR sh -> WTKR $ rrepl sh def+  WFTKS sh -> WTKS $ sconcrete $ Nested.sreplicateScal sh def+  WFTKX sh -> WTKX $ xrepl sh def+  WFTKProduct ftk1 ftk2 ->+    WTKProduct (defRepW ftk1) (defRepW ftk2)++concreteRepW+  :: forall y target. (ConvertTensor Concrete, ConvertTensor target)+  => (forall r. GoodScalar r => Concrete (TKScalar r) -> target (TKScalar r))+  -> (forall r sh. GoodScalar r => Concrete (TKS sh r) -> target (TKS sh r))+  -> (forall x z. FullShapeTK z -> target x -> target z)+  -> RepW Concrete y -> RepW target y+{-# INLINE concreteRepW #-}+concreteRepW concreteK concreteS fromS w = case w of+  WTKScalar v -> WTKScalar $ concreteK v+  WTKR v -> WTKR $+    let sh' = Nested.rshape $ unConcrete v+    in withShsFromShR sh' $ \(sh :: ShS sh) ->+      withKnownShS sh $+      fromS (FTKR sh' FTKScalar)+      $ concreteS (sfromR @_ @sh v)+  WTKS v -> WTKS $ concreteS v+  WTKX v -> WTKX $+    let sh' = Nested.mshape $ unConcrete v+    in withShsFromShX sh' $ \(sh :: ShS sh) ->+      withKnownShS sh $+      fromS (FTKX sh' FTKScalar)+      $ concreteS (sfromX @_ @sh v)+  WTKProduct v1 v2 ->+    WTKProduct (concreteRepW concreteK concreteS fromS v1)+               (concreteRepW concreteK concreteS fromS v2)++toADTensorKindW+  :: forall y target. BaseTensor target+  => RepW target y -> FullShapeTKW y -> RepW target (ADTensorKind y)+toADTensorKindW t = \case+  WFTKScalar @r -> case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> t+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> t+      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $+           WTKScalar $ kconcrete Z1+  WFTKR @r sh -> case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> t+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> t+      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $+           WTKR $ rrepl @_ @_ @target sh Z1+  WFTKS @r sh -> case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> t+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> t+      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $+           WTKS $ sconcrete $ Nested.sreplicateScal sh Z1+  WFTKX @r sh -> case testEquality (typeRep @r) (typeRep @Double) of+    Just Refl -> t+    _ -> case testEquality (typeRep @r) (typeRep @Float) of+      Just Refl -> t+      _ -> gcastWith (unsafeCoerceRefl :: ADTensorScalar r :~: Z1) $+           WTKX $ xrepl @_ @_ @target sh Z1+  WFTKProduct ftk1 ftk2 -> case t of+    WTKProduct t1 t2 ->+      WTKProduct (toADTensorKindW t1 ftk1) (toADTensorKindW t2 ftk2)++fromADTensorKindW+  :: forall y target. BaseTensor target+  => SingletonTK y -> RepW target (ADTensorKind y) -> RepW target y+fromADTensorKindW stk t = case (stk, t) of+  (STKScalar @r1, WTKScalar @r2 _) ->+    case testEquality (typeRep @r1) (typeRep @r2) of+      Just Refl -> t+      _ -> replRepW 0 WFTKScalar+  (STKR _ (STKScalar @r1), WTKR @r2 v) ->+    case testEquality (typeRep @r1) (typeRep @r2) of+      Just Refl -> t+      _ -> replRepW 0 (WFTKR (rshape v))+  (STKS sh (STKScalar @r1), WTKS @r2 _) ->+    case testEquality (typeRep @r1) (typeRep @r2) of+      Just Refl -> t+      _ -> replRepW 0 (WFTKS sh)+  (STKX _ (STKScalar @r1), WTKX @r2 v) ->+    case testEquality (typeRep @r1) (typeRep @r2) of+      Just Refl -> t+      _ -> replRepW 0 (WFTKX (xshape v))+  (STKProduct stk1 stk2, WTKProduct t1 t2) ->+    WTKProduct (fromADTensorKindW stk1 t1) (fromADTensorKindW stk2 t2)+  _ -> error "fromADTensorKindW: impossible SingletonTK"++type family UnWind y where+  UnWind (TKScalar r) =+    TKScalar r+  UnWind (TKR2 n (TKScalar r)) =+    TKR2 n (TKScalar r)+  UnWind (TKR2 n (TKR2 m x)) =+    UnWind (TKR2 (n + m) x)+  UnWind (TKR2 n (TKS2 sh2 x)) =+    UnWind (TKX2 (Replicate n Nothing ++ MapJust sh2) x)+  UnWind (TKR2 n (TKX2 sh2 x)) =+    UnWind (TKX2 (Replicate n Nothing ++ sh2) x)+  UnWind (TKR2 n (TKProduct y z)) =+    TKProduct (UnWind (TKR2 n y)) (UnWind (TKR2 n z))+  UnWind (TKS2 sh1 (TKScalar r)) =+    TKS2 sh1 (TKScalar r)+  UnWind (TKS2 sh1 (TKR2 m x)) =+    UnWind (TKX2 (MapJust sh1 ++ Replicate m Nothing) x)+  UnWind (TKS2 sh1 (TKS2 sh2 x)) =+    UnWind (TKS2 (sh1 ++ sh2) x)+  UnWind (TKS2 sh1 (TKX2 sh2 x)) =+    UnWind (TKX2 (MapJust sh1 ++ sh2) x)+  UnWind (TKS2 sh1 (TKProduct y z)) =+    TKProduct (UnWind (TKS2 sh1 y)) (UnWind (TKS2 sh1 z))+  UnWind (TKX2 sh1 (TKScalar r)) =+    TKX2 sh1 (TKScalar r)+  UnWind (TKX2 sh1 (TKR2 m x)) =+    UnWind (TKX2 (sh1 ++ Replicate m Nothing) x)+  UnWind (TKX2 sh1 (TKS2 sh2 x)) =+    UnWind (TKX2 (sh1 ++ MapJust sh2) x)+  UnWind (TKX2 sh1 (TKX2 sh2 x)) =+    UnWind (TKX2 (sh1 ++ sh2) x)+  UnWind (TKX2 sh1 (TKProduct y z)) =+    TKProduct (UnWind (TKX2 sh1 y)) (UnWind (TKX2 sh1 z))+  UnWind (TKProduct y z) =+    TKProduct (UnWind y) (UnWind z)++unWindSTK :: SingletonTK y -> SingletonTK (UnWind y)+unWindSTK = \case+  stk@STKScalar -> stk+  stk@(STKR _ STKScalar) -> stk+  STKR (SNat @n) (STKR (SNat @m) stk2) ->+    unWindSTK $ STKR (SNat @(n + m)) stk2+  STKR n (STKS sh2 stk2) ->+    unWindSTK+    $ STKX (ssxReplicate n `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2+  STKR n (STKX sh2 stk2) ->+    unWindSTK $ STKX (ssxReplicate n `ssxAppend` sh2) stk2+  STKR n@SNat (STKProduct y z) ->+    unWindSTK $ STKProduct (STKR n y) (STKR n z)+  stk@(STKS _ STKScalar) -> stk+  STKS sh1 (STKR m stk2) ->+    unWindSTK+    $ STKX (ssxFromShX (shxFromShS sh1) `ssxAppend` ssxReplicate m) stk2+  STKS sh1 (STKS sh2 stk2) ->+    unWindSTK $ STKS (sh1 `shsAppend` sh2) stk2+  STKS sh1 (STKX sh2 stk2) ->+    unWindSTK $ STKX (ssxFromShX (shxFromShS sh1) `ssxAppend` sh2) stk2+  STKS sh1 (STKProduct y z)->+    unWindSTK $ STKProduct (STKS sh1 y) (STKS sh1 z)+  stk@(STKX _ STKScalar) -> stk+  STKX sh1 (STKR m stk2) ->+    unWindSTK $ STKX (sh1 `ssxAppend` ssxReplicate m) stk2+  STKX sh1 (STKS sh2 stk2) ->+    unWindSTK $ STKX (sh1 `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2+  STKX sh1 (STKX sh2 stk2) ->+    unWindSTK $ STKX (sh1 `ssxAppend` sh2) stk2+  STKX sh1 (STKProduct y z) ->+    unWindSTK $ STKProduct (STKX sh1 y) (STKX sh1 z)+  STKProduct y z -> STKProduct (unWindSTK y) (unWindSTK z)++unWindFTK :: FullShapeTK y -> FullShapeTKW (UnWind y)+unWindFTK = \case+  FTKScalar -> WFTKScalar+  FTKR sh FTKScalar -> WFTKR sh+  FTKR sh1 (FTKR sh2 ftk2) ->+    unWindFTK $ FTKR (sh1 `shrAppend` sh2) ftk2+  FTKR sh1 (FTKS sh2 ftk2) ->+    unWindFTK+    $ FTKX (shxFromShR sh1 `shxAppend` shxFromShS sh2) ftk2+  FTKR sh1 (FTKX sh2 ftk2) ->+    unWindFTK $ FTKX (shxFromShR sh1 `shxAppend` sh2) ftk2+  FTKR sh1 (FTKProduct y z) ->+    unWindFTK $ FTKProduct (FTKR sh1 y) (FTKR sh1 z)+  FTKS sh FTKScalar -> WFTKS sh+  FTKS sh1 (FTKR sh2 ftk2) ->+    unWindFTK+    $ FTKX (shxFromShS sh1 `shxAppend` shxFromShR sh2) ftk2+  FTKS sh1 (FTKS sh2 ftk2) ->+    unWindFTK $ FTKS (sh1 `shsAppend` sh2) ftk2+  FTKS sh1 (FTKX sh2 ftk2) ->+    unWindFTK $ FTKX (shxFromShS sh1 `shxAppend` sh2) ftk2+  FTKS sh1 (FTKProduct y z) ->+    unWindFTK $ FTKProduct (FTKS sh1 y) (FTKS sh1 z)+  FTKX sh FTKScalar -> WFTKX sh+  FTKX sh1 (FTKR sh2 ftk2) ->+    unWindFTK $ FTKX (sh1 `shxAppend` shxFromShR sh2) ftk2+  FTKX sh1 (FTKS sh2 ftk2) ->+    unWindFTK $ FTKX (sh1 `shxAppend` shxFromShS sh2) ftk2+  FTKX sh1 (FTKX sh2 ftk2) ->+    unWindFTK $ FTKX (sh1 `shxAppend` sh2) ftk2+  FTKX sh1 (FTKProduct y z) ->+    unWindFTK $ FTKProduct (FTKX sh1 y) (FTKX sh1 z)+  FTKProduct y z -> WFTKProduct (unWindFTK y) (unWindFTK z)++-- This uses tunpairConv so to preserve sharing, @target@ either has+-- to have a `ShareTensor` instance or the argument has to be duplicable.+-- Only the argument of the first call, not of recursive calls,+-- is assumed to be duplicable. In the AST case, this creates+-- a tower of projections for product, but if it's balanced,+-- that's of logarithmic length, so maybe even better than sharing+-- excessively, which is hard for technical typing reasons.+unWindTarget :: ConvertTensor target+             => SingletonTK y -> target y -> RepW target (UnWind y)+unWindTarget stk t = case stk of+  STKScalar -> WTKScalar t+  STKR SNat STKScalar -> WTKR t+  STKR (SNat @n) (STKR (SNat @m) stk2) | Dict <- lemKnownSTK stk2 ->+    unWindTarget (STKR (SNat @(n + m)) stk2) (runNest t)+  STKR n@SNat (STKS sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh2 $+    unWindTarget (STKX (ssxReplicate n+                        `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2)+                 (runNestS t)+  STKR n@SNat (STKX sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $+    unWindTarget (STKX (ssxReplicate n `ssxAppend` sh2) stk2)+                 (runNestX t)+  STKR n@SNat (STKProduct stk1 stk2) ->+    unWindTarget (STKProduct (STKR n stk1) (STKR n stk2)) (runzip t)+  STKS sh1 STKScalar -> withKnownShS sh1 $ WTKS t+  STKS sh1 (STKR m@(SNat @m) stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh1 $+    unWindTarget (STKX (ssxFromShX (shxFromShS sh1)+                        `ssxAppend` ssxReplicate m) stk2) (sunNestR @_ @_ @m t)+  STKS sh1 (STKS sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh1 $ withKnownShS sh2 $+    unWindTarget (STKS (sh1 `shsAppend` sh2) stk2) (sunNest t)+  STKS sh1 (STKX sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $ withKnownShS sh1 $+    unWindTarget (STKX (ssxFromShX (shxFromShS sh1) `ssxAppend` sh2) stk2)+                 (sunNestX t)+  STKS sh1 (STKProduct stk1 stk2)->+    unWindTarget (STKProduct (STKS sh1 stk1) (STKS sh1 stk2)) (sunzip t)+  STKX sh1 STKScalar -> withKnownShX sh1 $ WTKX t+  STKX sh1 (STKR m@(SNat @m) stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh1 $+    unWindTarget (STKX (sh1 `ssxAppend` ssxReplicate m) stk2)+                 (xunNestR @_ @_ @m t)+  STKX sh1 (STKS sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh1 $ withKnownShS sh2 $+    unWindTarget (STKX (sh1 `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2)+                 (xunNestS t)+  STKX sh1 (STKX sh2 stk2) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh1 $ withKnownShX sh2 $+    unWindTarget (STKX (sh1 `ssxAppend` sh2) stk2) (xunNest t)+  STKX sh1 (STKProduct stk1 stk2) ->+    unWindTarget (STKProduct (STKX sh1 stk1) (STKX sh1 stk2)) (xunzip t)+  STKProduct stk1 stk2 ->+    let (t1, t2) = tunpairConv t+    in WTKProduct (unWindTarget stk1 t1) (unWindTarget stk2 t2)++windTarget :: ConvertTensor target+           => SingletonTK y -> RepW target (UnWind y) -> target y+windTarget stk t = case (stk, t) of+  (STKScalar, WTKScalar v) -> v+  (STKR _ STKScalar, WTKR v) -> v+  (STKR n@(SNat @n) (STKR (SNat @m) stk2), _)+   | Dict <- lemKnownSTK stk2 ->+    rnest n $ windTarget (STKR (SNat @(n + m)) stk2) t+  (STKR n (STKS sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh2 $+    rnestS n+    $ windTarget (STKX (ssxReplicate n+                        `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2) t+  (STKR n (STKX sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $+    rnestX n+    $ windTarget (STKX (ssxReplicate n `ssxAppend` sh2) stk2) t+  (STKR n@SNat (STKProduct stk1 stk2), _) | Dict <- lemKnownSTK stk1+                                          , Dict <- lemKnownSTK stk2 ->+    rzip $ windTarget (STKProduct (STKR n stk1) (STKR n stk2)) t+  (STKS _ STKScalar, WTKS v) -> v+  (STKS sh1 (STKR m@SNat stk2), _) | Dict <- lemKnownSTK stk2 ->+    snestR sh1+    $ windTarget (STKX (ssxFromShX (shxFromShS sh1)+                        `ssxAppend` ssxReplicate m) stk2) t+  (STKS sh1 (STKS sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh2 $+    snest sh1 $ windTarget (STKS (shsAppend sh1 sh2) stk2) t+  (STKS sh1 (STKX sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $+    snestX sh1 $ windTarget (STKX (ssxFromShX (shxFromShS sh1)+                                   `ssxAppend` sh2) stk2) t+  (STKS sh1 (STKProduct stk1 stk2), _) | Dict <- lemKnownSTK stk1+                                       , Dict <- lemKnownSTK stk2 ->+    szip $ windTarget (STKProduct (STKS sh1 stk1) (STKS sh1 stk2)) t+  (STKX _ STKScalar, WTKX v) -> v+  (STKX sh1 (STKR m@SNat stk2), _) | Dict <- lemKnownSTK stk2 ->+    xnestR sh1+    $ windTarget (STKX (sh1 `ssxAppend` ssxReplicate m) stk2) t+  (STKX sh1 (STKS sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShS sh2 $+    xnestS sh1+    $ windTarget (STKX (sh1 `ssxAppend` ssxFromShX (shxFromShS sh2)) stk2) t+  (STKX sh1 (STKX sh2 stk2), _) | Dict <- lemKnownSTK stk2 ->+    withKnownShX sh2 $+    xnest sh1 $ windTarget (STKX (ssxAppend sh1 sh2) stk2) t+  (STKX sh1 (STKProduct stk1 stk2), _) | Dict <- lemKnownSTK stk1+                                       , Dict <- lemKnownSTK stk2 ->+    xzip $ windTarget (STKProduct (STKX sh1 stk1) (STKX sh1 stk2)) t+  (STKProduct stk1 stk2, WTKProduct t1 t2) ->+    tpairConv (windTarget stk1 t1) (windTarget stk2 t2)+++-- * Operations defined using unwinding++-- | Add two (nested pairs of) tensors. Requires duplicable arguments+-- or a `ShareTensor` instance.+addTarget :: (BaseTensor target, ConvertTensor target)+          => SingletonTK y -> target y -> target y -> target y+addTarget stk a b =+  let a2 = unWindTarget stk a+      b2 = unWindTarget stk b+  in windTarget stk $ addRepW a2 b2++-- | Multiply two (nested pairs of) tensors. Requires duplicable arguments+-- or a `ShareTensor` instance.+multTarget :: (BaseTensor target, ConvertTensor target)+           => SingletonTK y -> target y -> target y -> target y+multTarget stk a b =+  let a2 = unWindTarget stk a+      b2 = unWindTarget stk b+  in windTarget stk $ multRepW a2 b2++-- | Sum all dimensions of each component and then sum it all.+-- Requires duplicable arguments or a `ShareTensor` instance.+sum0Target :: (BaseTensor target, ConvertTensor target)+           => FullShapeTK y -> target y+           -> target (TKScalar Double)+sum0Target ftk a =+  let a2 = unWindTarget (ftkToSTK ftk) a+  in sum0RepW (unWindFTK ftk) a2++-- | Dot product each component and then sum it all.+-- Requires duplicable arguments or a `ShareTensor` instance.+dot0Target :: (BaseTensor target, ConvertTensor target)+           => FullShapeTK y -> target y -> target y+           -> target (TKScalar Double)+dot0Target ftk a b =+  let a2 = unWindTarget (ftkToSTK ftk) a+      b2 = unWindTarget (ftkToSTK ftk) b+  in dot0RepW (unWindFTK ftk) a2 b2++-- | Replicate a scalar along the given full shape singleton.+replTarget :: forall y target. (BaseTensor target, ConvertTensor target)+           => (forall r. GoodScalar r => r)+           -> FullShapeTK y -> target y+replTarget r ftk =+  windTarget (ftkToSTK ftk) $ replRepW r (unWindFTK ftk)++-- | Replicate the default value along the given full shape singleton.+defTarget :: forall y target. (BaseTensor target, ConvertTensor target)+          => FullShapeTK y -> target y+defTarget ftk =+  windTarget (ftkToSTK ftk) $ defRepW (unWindFTK ftk)++concreteTarget+  :: forall y target. (ConvertTensor Concrete, ConvertTensor target)+  => (forall r. GoodScalar r => Concrete (TKScalar r) -> target (TKScalar r))+  -> (forall r sh. GoodScalar r => Concrete (TKS sh r) -> target (TKS sh r))+  -> (forall x z. FullShapeTK z -> target x -> target z)+  -> SingletonTK y -> Concrete y+  -> target y+concreteTarget concreteK concreteS fromS stk v =+  windTarget stk+  $ concreteRepW concreteK concreteS fromS+  $ unWindTarget stk v++lemUnWindOfAD :: SingletonTK y+              -> UnWind (ADTensorKind y) :~: ADTensorKind (UnWind y)+lemUnWindOfAD _ = unsafeCoerceRefl++-- | Convert a tensor into a tensor with only trivial non-differentiable+-- scalars. The `ShareTensor` constraint is needed, despite what GHC says,+-- in order not to require duplicable arguments.+toADTensorKindShared+  :: (BaseTensor target, ConvertTensor target, ShareTensor target)+  => FullShapeTK y -> target y+  -> target (ADTensorKind y)+toADTensorKindShared ftk a | Refl <- lemUnWindOfAD (ftkToSTK ftk) =+  windTarget (adSTK $ ftkToSTK ftk)+  $ toADTensorKindW (unWindTarget (ftkToSTK ftk) a) (unWindFTK ftk)++-- | Convert a tensor with only trivial non-differentiable scalars+-- into a tensor with the non-differentiable scalars given by the singleton+-- and with zero values at the non-differentiable types. The `ShareTensor`+-- constraint is needed, despite what GHC says, in order not to require+-- duplicable arguments.+fromADTensorKindShared+  :: (BaseTensor target, ConvertTensor target, ShareTensor target)+  => SingletonTK y -> target (ADTensorKind y)+  -> target y+fromADTensorKindShared stk a | Refl <- lemUnWindOfAD stk =+  windTarget stk+  $ fromADTensorKindW (unWindSTK stk) $ unWindTarget (adSTK stk) a
+ src/HordeAd/External/CommonRankedOps.hs view
@@ -0,0 +1,267 @@+{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}+-- | Commonly used ranked operations on tensors.+module HordeAd.External.CommonRankedOps+  ( module HordeAd.External.CommonRankedOps+  ) where++import Prelude++import Control.Exception.Assert.Sugar+import Data.List.NonEmpty qualified as NonEmpty+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:) (Refl))+import GHC.Exts (IsList (..))+import GHC.TypeLits (KnownNat, sameNat)++import Data.Array.Nested.Types (unsafeCoerceRefl)+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape++import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.OpsTensor++-- This is not only ranked, so move it once we have CommonAnyOps.hs.+assumeEquality :: forall a b r. (a ~ b => r) -> r+assumeEquality = gcastWith (unsafeCoerceRefl :: a :~: b)++rminimum :: forall target n r.+            ( BaseTensor target, ConvertTensor target, LetTensor target+            , GoodScalar r )+         => target (TKR n r) -> target (TKR 0 r)+rminimum t =+  tlet (rflatten t) $ \tf ->+    rindex0 tf (tprimalPart (kfromR (rminIndex tf)) :.: ZIR)++rmaximum :: forall target n r.+            ( BaseTensor target, ConvertTensor target, LetTensor target+            , GoodScalar r )+         => target (TKR n r) -> target (TKR 0 r)+rmaximum t =+  tlet (rflatten t) $ \tf ->+    rindex0 tf (tprimalPart (kfromR (rmaxIndex tf)) :.: ZIR)++rfromIndex0 :: forall r target.+               (BaseTensor target, ConvertTensor target, GoodScalar r)+            => IntOf target -> target (TKR 0 r)+rfromIndex0 = rfromIntegral . rfromK . tfromPrimal STKScalar++rfromIndex1 :: forall n r target.+               ( KnownNat n , BaseTensor target+               , BaseTensor (PrimalOf target), ConvertTensor (PrimalOf target)+               , GoodScalar r )+            => IxROf target n -> target (TKR 1 r)+rfromIndex1 = case sameNat (Proxy @n) (Proxy @0) of+  Just Refl -> const $ rconcrete $ Nested.rfromListPrimLinear (0 :$: ZSR) []+  _ -> rfromIntegral . rfromPrimal . rfromList . NonEmpty.fromList+       . map rfromK . toList++{-+rint64FromIndex1 :: forall n target.+                    ( KnownNat n+                    , BaseTensor target, BaseTensor (PrimalOf target) )+                 => IxROf target n -> target (TKR Int64 1)+rint64FromIndex1 = case sameNat (Proxy @n) (Proxy @0) of+  Just Refl -> const $ rconcrete $ Nested.rfromListPrimLinear (0 :$: ZSR) []+  _ -> rfromPrimal . rfromList . NonEmpty.fromList . indexToList++rint64ToIndex1 :: forall n target.+                  ( KnownNat n+                  , BaseTensor target, BaseTensor (PrimalOf target) )+               => target (TKR Int64 1) -> IxROf target n+rint64ToIndex1 v = listToIndex $ runravelToList $ rprimalPart v++tletIx :: ( KnownNat n, KnownNat m, GoodScalar r+          , BaseTensor target, BaseTensor (PrimalOf target)+          , LetTensor target )+       => IxROf target n -> (IxROf target n -> target (TKR m r))+       -> target (TKR m r)+tletIx ix0 f = tlet (rint64FromIndex1 ix0) $ \ixT -> f $ rint64ToIndex1 ixT+-}++relu, reluLeaky+  :: forall target n r.+     (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+  => target (TKR n r) -> target (TKR n r)+relu v0 = tlet v0 $ \v ->+  let oneIfGtZero =+        rmap0N (\x -> ifH (x <=. rscalar 0) (rscalar 0.0) (rscalar 1.0)) v+  in oneIfGtZero * v++reluLeaky v0 = tlet v0 $ \v ->+  let oneIfGtZero =+        rmap0N (\x -> ifH (x <=. rscalar 0) (rscalar 0.01) (rscalar 1.0)) v+  in oneIfGtZero * v++logistic :: forall target r n.+            ( BaseTensor target, LetTensor target, BaseTensor (PrimalOf target)+            , KnownNat n, GoodScalar r, Differentiable r )+         => target (TKR n r) -> target (TKR n r)+logistic d0 = tlet d0 $ \d ->  -- used in rprimalPart and in tdualPart+  let one = rrepl (rshape d) 1+      y0 = recip (one + exp (- rprimalPart d))+  in ttletPrimal y0 $ \y ->+       rfromPrimal y + rfromDual (rScale @target (y * (one - y)) $ rdualPart d)++-- Optimized and more clearly written @u ** 2@. It's not clear if this is+-- currently faster than @u ** 2@ and in which pipelines, but it's different,+-- so useful as a test.+square :: forall target r n.+          ( BaseTensor target, LetTensor target+          , KnownNat n, Num (PrimalOf target (TKR n r)), GoodScalar r )+       => target (TKR n r) -> target (TKR n r)+square d = let u = rprimalPart @target d+               u' = rdualPart @target d+           in tD knownSTK (u * u) (rScale @target (2 * u) u')++squaredDifference+  :: forall target n r.+     ( BaseTensor target, LetTensor target+     , KnownNat n, Num (PrimalOf target (TKR n r)), GoodScalar r )+  => PrimalOf target (TKR n r) -> target (TKR n r) -> target (TKR n r)+squaredDifference targ res = square @target $ res - rfromPrimal @target targ++lossCrossEntropyV+  :: ( BaseTensor target, ConvertTensor target+     , KnownNat n, GoodScalar r, Differentiable r )+  => target (TKR n r) -> target (TKR n r) -> target (TKScalar r)+lossCrossEntropyV targ res = kfromR $ negate $ log res `rdot0` targ++-- | Note that this is equivalent to a composition of softMax and cross entropy+-- only when @expected@ is one-hot. Otherwise, results vary wildly. In our+-- rendering of the MNIST data all labels are one-hot.+lossSoftMaxCrossEntropyR+  :: forall target n r.+     ( BaseTensor target, ConvertTensor target, LetTensor target+     , BaseTensor (PrimalOf target), ConvertTensor (PrimalOf target)+     , LetTensor (PrimalOf target)+     , KnownNat n, GoodScalar r, Differentiable r )+  => PrimalOf target (TKR n r) -> target (TKR n r) -> target (TKScalar r)+lossSoftMaxCrossEntropyR expected d' = tlet d' $ \d ->+  -- The following protects from underflows, overflows and exploding gradients+  -- and is required by QuickCheck tests to avoid NaNs, etc., for argument+  -- values we don't fully control.+  -- See https://github.com/tensorflow/tensorflow/blob/5a566a7701381a5cf7f70fce397759483764e482/tensorflow/core/kernels/sparse_softmax_op.cc#L106+  -- and https://github.com/tensorflow/tensorflow/blob/5a566a7701381a5cf7f70fce397759483764e482/tensorflow/core/kernels/xent_op.h+  let u = rprimalPart d+      expU' = exp (u - rreplicate0N (rshape u) (rminimum u))+  in ttletPrimal expU' $ \expU ->+    let softMaxU0 =+          let sumExpU = rsum0 expU+              recipSum = recip sumExpU+          in rreplicate0N (rshape u) recipSum * expU+    in ttletPrimal softMaxU0 $ \softMaxU -> kfromR $+      tD knownSTK+         (negate $ log softMaxU `rdot0` expected)+           -- TODO: avoid: log . exp+         (rdualPart $ rfromPrimal (softMaxU - expected) `rdot0` d)++-- | No padding; remaining areas ignored.+maxPool1 :: ( BaseTensor target, ConvertTensor target, LetTensor target+            , GoodScalar r )+         => Int -> Int -> target (TKR 1 r) -> target (TKR 1 r)+maxPool1 ksize stride v =+  let slices = [rslice i ksize v | i <- [0, stride .. rwidth v - ksize]]+  in rfromList $ NonEmpty.fromList $ map rmaximum slices++softMax1 :: ( BaseTensor target, LetTensor target+            , KnownNat n, GoodScalar r, Differentiable r )+         => target (TKR n r) -> target (TKR n r)+softMax1 d =+  let expU0 = exp d+  in tlet expU0 $ \expU -> rreplicate0N (rshape d) (recip $ rsum0 expU) * expU++-- | Unpadded full convolution,+--   where the output size is the same as the input size.+--+-- BTW, the indexing lower bounds in the code are spurious,+-- so they get simplified away in the resulting AST program.+conv2dUnpadded+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicez shK1 arrA [iImg, 0, iBh, iBw]+          arrKt = slicez shK1 arrK [iCout, 0, 0, 0]+      in rdot0 arrAt arrKt+    _ -> error "conv2dUnpadded: impossible pattern needlessly required"++-- | Full convolution with custom padding,+--   where the output size depends on the input size, kernel size and padding.+conv2dCustomPadded+  :: (ADReady target, GoodScalar r)+  => (Int, Int) -> target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dCustomPadded (nPh, nPw) arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      nBh = nAh + 2 * nPh - nKh + 1+      nBw = nAw + 2 * nPw - nKw + 1+      shB = [nImgs, nCoutK, nBh, nBw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let iFh = iBh - fromIntegral nPh+          iFw = iBw - fromIntegral nPw+          arrAt = slicez shK1 arrA [iImg, 0, iFh, iFw]+          arrKt = slicez shK1 arrK [iCout, 0, 0, 0]+      in rdot0 arrAt arrKt+    _ -> error "conv2dCustomPadded: impossible pattern needlessly required"++-- | Full convolution with just enough padding to ensure all output points+--   are affected by the same number of input points,+--   where the output size shrinks depending on the input size and kernel size.+--   Also no input points are ever ignored, though some are read less often.+conv2dShrinking+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dShrinking arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh - nKh, nAw - nKw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicez shK1 arrA [iImg, 0, iBh, iBw]+          arrKt = slicez shK1 arrK [iCout, 0, 0, 0]+      in rdot0 arrAt arrKt+    _ -> error "conv2dShrinking: impossible pattern needlessly required"++-- | Slice a section out of a tensor,+--   given a base offset and shape of the section.+--+--   If the slice extends out side the source array then the corresponding+--   elements are set to zero.+slicez+  :: (ADReady target, GoodScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez shOut d ixBase =+  rbuild shOut $ \ixResult -> rindex0 d (ixrZipWith (+) ixBase ixResult)++maxPool2dUnpadded+  :: (ADReady target, GoodScalar r)+  => Int -> Int -> target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded ksize stride arr =+  let [batch_size, channels, h, w] = rshape arr+      shOut = [batch_size, channels, h `div` stride, w `div` stride]+      shK1 = [1, 1, ksize, ksize]+  in rbuild shOut $ \case+    [iImg, iChan, iBh, iBw] ->+      rmaximum $ slicez shK1 arr [ iImg, iChan+                                 , fromIntegral stride * iBh+                                 , fromIntegral stride * iBw ]+    _ -> error "maxPool2dUnpadded: impossible pattern needlessly required"++xfromIndex0 :: forall r target.+               (BaseTensor target, ConvertTensor target, GoodScalar r)+            => IntOf target -> target (TKX '[] r)+xfromIndex0 = xfromIntegral . xfromK . tfromPrimal STKScalar
+ src/HordeAd/External/CommonShapedOps.hs view
@@ -0,0 +1,229 @@+{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Commonly used operations on shaped tensors.+module HordeAd.External.CommonShapedOps+  ( module HordeAd.External.CommonShapedOps+  ) where++import Prelude++import Data.List.NonEmpty qualified as NonEmpty+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:) (Refl))+import Data.Type.Ord (Compare)+import GHC.Exts (IsList (..))+import GHC.TypeLits+  (Div, KnownNat, SomeNat (..), sameNat, someNatVal, type (-), type (<=))++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Convert (ixrFromIxS, ixsFromIxR')+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (unsafeCoerceRefl)++import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.OpsTensor++sminimum :: forall r sh target. (ADReady target, GoodScalar r, KnownShS sh)+         => target (TKS sh r) -> target (TKS '[] r)+sminimum t | SNat <- shsProduct (knownShS @sh) =+  tlet (sflatten t) $ \tf ->+    sindex0 tf (tprimalPart (kfromS (sminIndex tf)) :.$ ZIS)++smaximum :: forall r sh target. (ADReady target, GoodScalar r, KnownShS sh)+         => target (TKS sh r) -> target (TKS '[] r)+smaximum t | SNat <- shsProduct (knownShS @sh) =+  tlet (sflatten t) $ \tf ->+    sindex0 tf (tprimalPart (kfromS (smaxIndex tf)) :.$ ZIS)++sfromIndex0 :: forall r target. (ADReady target, GoodScalar r)+            => IntOf target -> target (TKS '[] r)+sfromIndex0 = sfromR . rfromIntegral . rfromPrimal . rfromK++sfromIndex1 :: forall r sh target.+               (ADReady target, GoodScalar r, KnownShS sh)+            => IxSOf target sh -> target (TKS '[Rank sh] r)+sfromIndex1 | SNat <- shsRank (knownShS @sh) =+  case sameNat (Proxy @(Rank sh)) (Proxy @0) of+    Just Refl -> const $ sconcrete $ Nested.sfromListPrimLinear knownShS []+    _ -> sfromR . rfromIntegral . rfromPrimal . rfromList+         . NonEmpty.fromList . map rfromK . toList++{-+sletIx :: forall r sh n target.+          (ADReady target, GoodScalar r, KnownShS sh, KnownNat n)+       => IxROf target n -> (IxROf target n -> target (TKS sh r))+       -> target (TKS sh r)+sletIx ix0 f = tlet (sfromR @target @Int64 @'[n]+                     $ rint64FromIndex1 ix0) $ \ixT ->+                 f $ rint64ToIndex1 $ rfromS @target ixT+-}++reluS, reluLeakyS+  :: forall target sh r.+     (KnownShS sh, ADReady target, GoodScalar r, Differentiable r)+  => target (TKS sh r) -> target (TKS sh r)+reluS v0 = tlet v0 $ \v ->+  let oneIfGtZero =+        smap0N (\x -> ifH (x <=. sscalar 0) (sscalar 0.0) (sscalar 1.0)) v+  in oneIfGtZero * v++reluLeakyS v0 = tlet v0 $ \v ->+  let oneIfGtZero =+        smap0N (\x -> ifH (x <=. sscalar 0) (sscalar 00.01) (sscalar 01.0)) v+  in oneIfGtZero * v++logisticS :: forall target r sh.+             ( BaseTensor target, LetTensor target, BaseTensor (PrimalOf target)+             , KnownShS sh, GoodScalar r, Differentiable r )+          => target (TKS sh r) -> target (TKS sh r)+logisticS d0 = tlet d0 $ \d ->  -- used in rprimalPart and in sdualPart+  let one = srepl 1+      y0 = recip (one + exp (- sprimalPart d))+  in ttletPrimal y0 $ \y ->+       sfromPrimal y + sfromDual (sScale @target (y * (one - y)) $ sdualPart d)++-- Optimized and more clearly written @u ** 2@. It's not clear if this is+-- currently faster than @u ** 2@ and in which pipelines, but it's different,+-- so useful as a test.+squareS :: forall target r sh.+           ( KnownShS sh, BaseTensor target, LetTensor target+           , Num (PrimalOf target (TKS sh r)), GoodScalar r )+        => target (TKS sh r) -> target (TKS sh r)+squareS d = let u = sprimalPart d+                u' = sdualPart d+            in tD knownSTK (u * u) (sScale @target (2 * u) u')++squaredDifferenceS+  :: forall target sh r.+     ( KnownShS sh, BaseTensor target, LetTensor target+     , Num (PrimalOf target (TKS sh r)), GoodScalar r )+  => PrimalOf target (TKS sh r) -> target (TKS sh r) -> target (TKS sh r)+squaredDifferenceS targ res = squareS $ res - sfromPrimal targ++lossCrossEntropyVS :: ( KnownShS sh, GoodScalar r, Differentiable r+                      , BaseTensor target, ConvertTensor target )+                   => target (TKS sh r)+                   -> target (TKS sh r)+                   -> target (TKScalar r)+lossCrossEntropyVS targ res = kfromS $ negate $ log res `sdot0` targ++-- | Note that this is equivalent to a composition of softMax and cross entropy+-- only when @expected@ is one-hot. Otherwise, results vary wildly. In our+-- rendering of the MNIST data all labels are one-hot.+lossSoftMaxCrossEntropyS+  :: forall target sh r.+     ( ADReady target, ADReady (PrimalOf target), GoodScalar r, KnownShS sh+     , Differentiable r )+  => PrimalOf target (TKS sh r) -> target (TKS sh r) -> target (TKScalar r)+lossSoftMaxCrossEntropyS expected d' = tlet d' $ \d ->+  -- The following protects from underflows, overflows and exploding gradients+  -- and is required by QuickCheck tests to avoid NaNs, etc., for argument+  -- values we don't fully control.+  -- See https://github.com/tensorflow/tensorflow/blob/5a566a7701381a5cf7f70fce397759483764e482/tensorflow/core/kernels/sparse_softmax_op.cc#L106+  -- and https://github.com/tensorflow/tensorflow/blob/5a566a7701381a5cf7f70fce397759483764e482/tensorflow/core/kernels/xent_op.h+  let softMaxU0 =+        let u = sprimalPart d+            expU' = exp (u - sreplicate0N (sminimum u))+        in tlet expU' $ \expU ->+          let sumExpU = ssum0 expU+              recipSum = recip sumExpU+          in sreplicate0N recipSum * expU+  in ttletPrimal softMaxU0 $ \softMaxU -> kfromS $+    tD knownSTK+       (negate $ log softMaxU `sdot0` expected)+         -- TODO: avoid: log . exp+       (sdualPart $ sfromPrimal (softMaxU - expected) `sdot0` d)++-- | No padding; remaining areas ignored.+maxPool1S :: forall ksize stride m target r.+             ( ADReady target, GoodScalar r+             , KnownNat ksize, KnownNat stride, KnownNat m )+          => target (TKS '[m] r) -> target (TKS '[m] r)+maxPool1S v =+  let l = [0, valueOf @stride .. swidth v - valueOf @ksize]+      maxOfSlice i =+        case someNatVal $ toInteger i of+          Just (SomeNat @i _proxy) ->+            gcastWith (unsafeCoerceRefl :: Compare i m :~: LT) $+            gcastWith (unsafeCoerceRefl :: Compare ksize (m - i) :~: LT) $+            smaximum $ sslice @i @(m - i - ksize) @ksize SNat SNat SNat v+          Nothing -> error "maxPool1S: impossible someNatVal error"+  in sfromList $ NonEmpty.fromList $ map maxOfSlice l++softMax1S :: forall target sh r.+             ( KnownShS sh, BaseTensor target, LetTensor target+             , GoodScalar r, Differentiable r )+          => target (TKS sh r) -> target (TKS sh r)+softMax1S d =+  let expU0 = exp d+  in tlet expU0 $ \expU -> sreplicate0N (recip $ ssum0 expU) * expU++-- | Unpadded full convolution,+--   where the output size is the same as the input size.+conv2dUnpaddedS+  :: forall nCoutK nCinpK nKh nKw nImgs nCinpA nAh nAw+            target r shB shK1.+     ( KnownNat nCoutK, KnownNat nCinpK, KnownNat nKh, KnownNat nKw+     , KnownNat nImgs, KnownNat nAh, KnownNat nAw+     , ADReady target, GoodScalar r+     , nCinpA ~ nCinpK+     , shB ~ '[nImgs, nCoutK, nAh, nAw]+     , shK1 ~ '[1, nCinpA, nKh, nKw]+     )+  => target (TKS '[nCoutK, nCinpK, nKh, nKw] r)+  -> target (TKS '[nImgs, nCinpA, nAh, nAw] r)+  -> target (TKS shB r)+conv2dUnpaddedS arrK arrA =+  sbuild @(Rank shB) $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezS @shK1 arrA [iImg, 0, iBh, iBw]+          arrKt = slicezS arrK [iCout, 0, 0, 0]+      in sdot0 arrAt arrKt+    _ -> error "conv2dUnpaddedS: impossible pattern needlessly required"++-- | Slice a section out of a tensor,+--   given a base offset and shape of the section.+--+--   If the slice extends out side the source array then the corresponding+--   elements are set to zero.+slicezS+  :: forall shOut sh target r.+     ( KnownShS sh, KnownShS shOut, KnownShS (Take (Rank sh) shOut)+     , KnownNat (Rank sh)+     , Rank shOut ~ Rank sh, ADReady target, GoodScalar r )+  => target (TKS sh r) -> IxSOf target sh -> target (TKS shOut r)+slicezS d ixBase =+  gcastWith (unsafeCoerceRefl+             :: Rank (Take (Rank shOut) shOut) :~: Rank shOut) $+  gcastWith (unsafeCoerceRefl :: Drop (Rank sh) shOut :~: '[]) $+  sbuild @(Rank shOut)+  $ \ixResult ->+      sindex0 d (ixsFromIxR' knownShS $ ixrZipWith (+) (ixrFromIxS ixBase) (ixrFromIxS ixResult))+      -- TODO: this doesn't work, because ixsZipWith has too strict a type:+      -- sbuild @(Rank shOut) $ \ixResult -> sindex0 d (ixsZipWith (+) ixBase ixResult)++maxPool2dUnpaddedS+  :: forall ksize stride batch_size channels h w target r shOut shK1.+     ( KnownNat ksize, KnownNat stride, KnownNat batch_size, KnownNat channels+     , KnownNat h, KnownNat w+     , 1 <= stride  -- wrongly reported as redundant due to plugins+     , ADReady target, GoodScalar r+     , shOut ~ '[batch_size, channels, h `Div` stride, w `Div` stride]+     , shK1 ~ '[1, 1, ksize, ksize]+     )+  => target (TKS '[batch_size, channels, h, w] r)+  -> target (TKS shOut r)+maxPool2dUnpaddedS arr =+  let stride = valueOf @stride :: Int+  in sbuild @(Rank shOut) $ \case+    [iImg, iChan, iBh, iBw] ->+      smaximum $ slicezS @shK1 arr [ iImg, iChan+                                   , fromIntegral stride * iBh+                                   , fromIntegral stride * iBw ]+    _ -> error "maxPool2dUnpaddedS: impossible pattern needlessly required"
+ src/HordeAd/External/Optimizer.hs view
@@ -0,0 +1,87 @@+-- | A couple of gradient descent scheme implementations.+module HordeAd.External.Optimizer+  ( sgd, sgdSTK, updateWithGradient+  , sgdAdam, sgdAdamArgs, updateWithGradientAdam+  , StateAdam, initialStateAdam, defaultArgsAdam+  ) where++import Prelude++import HordeAd.Core.CarriersADVal+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Delta+import HordeAd.Core.OpsADVal+import HordeAd.Core.OpsConcrete ()+import HordeAd.Core.TensorKind+import HordeAd.External.OptimizerTools++-- | Stochastic Gradient Descent.+sgdSTK :: forall a x z.+          SingletonTK x+       -> Double  -- ^ gamma (learning_rate?)+       -> (a -> ADVal Concrete x -> ADVal Concrete z)+       -> [a]  -- ^ training data+       -> Concrete x  -- ^ initial parameters+       -> (Concrete x, Concrete z)+sgdSTK stk gamma f trainingData parameters0 = go trainingData parameters0 where+  zftk = tftkG stk $ unConcrete parameters0+  deltaInputs :: Delta Concrete x+  deltaInputs = generateDeltaInputs zftk+  go :: [a] -> Concrete x -> (Concrete x, Concrete z)+  go [] parameters = (parameters, undefined)+  go (a : rest) !parameters =+    let inputs :: ADVal Concrete x+        inputs = dDnotShared parameters deltaInputs+        (gradients, valueNew) = crevOnADInputs Nothing (f a) zftk inputs+        parametersNew = updateWithGradient gamma stk parameters gradients+    in if null rest+       then (parametersNew, valueNew)+       else go rest parametersNew++sgd :: forall a x z. KnownSTK x+    => Double  -- ^ gamma (learning_rate?)+    -> (a -> ADVal Concrete x -> ADVal Concrete z)+    -> [a]  -- ^ training data+    -> Concrete x  -- ^ initial parameters+    -> (Concrete x, Concrete z)+sgd = sgdSTK knownSTK++-- We inline (possibly causing a binary blowup) until we are able to work around+-- https://gitlab.haskell.org/ghc/ghc/-/issues/23798+-- and specialize.+-- | An implementation of the Adam gradient descent.+sgdAdam+  :: forall a x z . KnownSTK x+  => (a -> ADVal Concrete x -> ADVal Concrete z)+  -> [a]+  -> Concrete x+  -> StateAdam x+  -> (Concrete x, StateAdam x)+{-# INLINE sgdAdam #-}+sgdAdam = sgdAdamArgs defaultArgsAdam++sgdAdamArgs+  :: forall a x z. KnownSTK x+  => ArgsAdam+  -> (a -> ADVal Concrete x -> ADVal Concrete z)+  -> [a]+  -> Concrete x+  -> StateAdam x+  -> (Concrete x, StateAdam x)+{-# INLINE sgdAdamArgs #-}+sgdAdamArgs argsAdam f trainingData !parameters0 !stateAdam0 =+  go trainingData parameters0 stateAdam0+ where+  zftk = tftkG knownSTK $ unConcrete parameters0+  deltaInputs :: Delta Concrete x+  deltaInputs = generateDeltaInputs zftk+  go :: [a] -> Concrete x -> StateAdam x -> (Concrete x, StateAdam x)+  go [] parameters stateAdam = (parameters, stateAdam)+  go (a : rest) !parameters !stateAdam =+    let inputs :: ADVal Concrete x+        inputs = dDnotShared parameters deltaInputs+        gradients = fst $ crevOnADInputs Nothing (f a) zftk inputs+        (parametersNew, stateAdamNew) =+          updateWithGradientAdam+            argsAdam stateAdam knownSTK parameters gradients+    in go rest parametersNew stateAdamNew
+ src/HordeAd/External/OptimizerTools.hs view
@@ -0,0 +1,243 @@+-- | Tools for implementing (and debugging the use of) gradient descent schemes.+module HordeAd.External.OptimizerTools+  ( updateWithGradient+--  , gradientIsNil, minimumGradient, maximumGradient+  , ArgsAdam(..), defaultArgsAdam+  , StateAdam(..), initialStateAdam+  , updateWithGradientAdam+  ) where++import Prelude++import Data.Type.Equality (gcastWith, (:~:))++import Data.Array.Nested.Types (unsafeCoerceRefl)+import Data.Array.Nested qualified as Nested++import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Ops+import HordeAd.Core.OpsConcrete ()+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++updateWithGradient :: forall y.+                      Double -> SingletonTK y+                   -> Concrete y -> Concrete (ADTensorKind y)+                   -> Concrete y+updateWithGradient gamma stk p@(Concrete params)+                             g@(Concrete gradient) = case stk of+  STKScalar @r -> Concrete $+    ifDifferentiable @r+      (gcastWith (unsafeCoerceRefl :: y :~: ADTensorKind y) $+       params - realToFrac gamma * gradient)+      params+  STKR _ (STKScalar @r) -> Concrete $+    ifDifferentiable @r+      (gcastWith (unsafeCoerceRefl :: y :~: ADTensorKind y) $+       params - Nested.rreplicateScal (Nested.rshape params)+                                      (realToFrac gamma)+                * gradient)+      params+  STKS _ (STKScalar @r) -> Concrete $+    ifDifferentiable @r+      (gcastWith (unsafeCoerceRefl :: y :~: ADTensorKind y) $+       params - Nested.sreplicateScal (Nested.sshape params)+                                      (realToFrac gamma)+                * gradient)+      params+  STKX _ (STKScalar @r) -> Concrete $+    ifDifferentiable @r+      (gcastWith (unsafeCoerceRefl :: y :~: ADTensorKind y) $+       params - Nested.mreplicateScal (Nested.mshape params)+                                      (realToFrac gamma)+                * gradient)+      params+  STKProduct stk1 stk2 ->+    tpair (updateWithGradient gamma stk1 (tproject1 p) (tproject1 g))+          (updateWithGradient gamma stk2 (tproject2 p) (tproject2 g))+  _ -> error "updateWithGradient: only non-nested arrays supported"+         -- we could support nested arrays, but some new unwinding code+         -- would be needed here and for Adam, with no existing use case++{- there are mostly for debugging+gradientIsNil :: (Eq r) => HVector Concrete -> Bool+gradientIsNil (HVector Concrete gradient0 gradientR) =+  V.all (== 0) gradient0+  && V.all isTensorDummyD gradientR++minimumGradient :: (Ord r) => HVector Concrete -> r+minimumGradient (HVector Concrete gradient0 gradientR) =+  min (if V.null gradient0 then 0 else LA.minElement gradient0)+      (if V.null gradientR then 0+       else V.minimum (V.map OR.minimumA gradientR))++maximumGradient :: (Ord r) => HVector Concrete -> r+maximumGradient (HVector Concrete gradient0 gradientR) =+  max (if V.null gradient0 then 0 else LA.maxElement gradient0)+      (if V.null gradientR then 0+       else V.maximum (V.map OR.maximumA gradientR))+-}++data ArgsAdam = ArgsAdam+  { alpha   :: Double+  , betaOne :: Double+  , betaTwo :: Double+  , epsilon :: Double+  }++-- The defaults taken from+-- https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam+defaultArgsAdam :: ArgsAdam+defaultArgsAdam = ArgsAdam+  { alpha = 0.001+  , betaOne = 0.9+  , betaTwo = 0.999+  , epsilon = 1e-7+  }++type family Triplify y where+  Triplify (TKScalar r) =+    TKProduct (TKProduct (TKScalar r) (TKScalar r)) (TKScalar r)+  Triplify (TKR n r) = TKProduct (TKProduct (TKR n r) (TKR n r)) (TKR n r)+  Triplify (TKS sh r) = TKProduct (TKProduct (TKS sh r) (TKS sh r)) (TKS sh r)+  Triplify (TKX sh r) = TKProduct (TKProduct (TKX sh r) (TKX sh r)) (TKX sh r)+  Triplify (TKProduct x z) = TKProduct (Triplify x) (Triplify z)++unzip3Rep+  :: SingletonTK y -> Concrete (Triplify y)+  -> (Concrete y, Concrete y, Concrete y)+unzip3Rep stk (Concrete t) = case stk of+  STKScalar ->+    (Concrete $ fst $ fst t, Concrete $ snd $ fst t, Concrete $ snd t)+  STKR _ STKScalar ->+    (Concrete $ fst $ fst t, Concrete $ snd $ fst t, Concrete $ snd t)+  STKS _ STKScalar ->+    (Concrete $ fst $ fst t, Concrete $ snd $ fst t, Concrete $ snd t)+  STKX _ STKScalar ->+    (Concrete $ fst $ fst t, Concrete $ snd $ fst t, Concrete $ snd t)+  STKProduct stk1 stk2 ->+    let !(!a1, !b1, !c1) = unzip3Rep stk1 $ Concrete $ fst t+        !(!a2, !b2, !c2) = unzip3Rep stk2 $ Concrete $ snd t+    in ( Concrete (unConcrete a1, unConcrete a2)+       , Concrete (unConcrete b1, unConcrete b2)+       , Concrete (unConcrete c1, unConcrete c2))+  _ -> error "unzip3Rep: only non-nested arrays supported"++type role StateAdam nominal+data StateAdam y = StateAdam+  { tAdam :: Int  -- iteration count+  , mAdam :: Concrete y+  , vAdam :: Concrete y+  }++-- TODO: introduce and use something like TensorOrZero+initialStateAdam :: FullShapeTK y -> StateAdam y+initialStateAdam ftk =+  StateAdam { tAdam = 0+            , mAdam = tdefTarget ftk+            , vAdam = tdefTarget ftk+            }++updateWithGradientAdam+  :: ArgsAdam -> StateAdam y -> SingletonTK y+  -> Concrete y -> Concrete (ADTensorKind y)+  -> (Concrete y, StateAdam y)+updateWithGradientAdam ArgsAdam{..} StateAdam{..} stk0 paramsR gradientR =+  let mAdamR = mAdam+      vAdamR = vAdam+      tAdamNew = tAdam + 1+      oneMinusBeta1 = 1 - betaOne+      oneMinusBeta2 = 1 - betaTwo+      updateR :: (Fractional r, Nested.FloatElt r, Nested.PrimElt r)+              => Nested.Ranked n r -> Nested.Ranked n r+              -> Nested.Ranked n r -> Nested.Ranked n r+              -> (Nested.Ranked n r, Nested.Ranked n r, Nested.Ranked n r)+      updateR mA vA p g =+        let sh = Nested.rshape g+            mANew = Nested.rreplicateScal sh (realToFrac betaOne) * mA+                    + Nested.rreplicateScal sh (realToFrac oneMinusBeta1) * g+            vANew = Nested.rreplicateScal sh (realToFrac betaTwo) * vA+                    + Nested.rreplicateScal sh (realToFrac oneMinusBeta2)+                      * (g * g)+            alphat = alpha * sqrt (1 - betaTwo ^ tAdamNew)+                             / (1 - betaOne ^ tAdamNew)+        in ( mANew+           , vANew+           , p - (Nested.rreplicateScal sh (realToFrac alphat) * mANew)+                 / (sqrt vANew+                    + Nested.rreplicateScal sh (realToFrac epsilon)) )+      updateProd :: forall y2.+                    SingletonTK y2+                 -> Concrete y2 -> Concrete y2+                 -> Concrete y2 -> Concrete (ADTensorKind y2)+                 -> Concrete (Triplify y2)+      updateProd stk (Concrete mA) (Concrete vA)+                     (Concrete p) (Concrete g) = case stk of+        STKScalar @r ->+          ifDifferentiable @r+            (gcastWith (unsafeCoerceRefl :: y2 :~: ADTensorKind y2) $+             let !(!mAN, !vAN, !pN) =+                   updateR (Nested.rscalar mA)+                           (Nested.rscalar vA)+                           (Nested.rscalar p)+                           (Nested.rscalar g)+             in Concrete+                (( Nested.runScalar mAN+                 , Nested.runScalar vAN )+                , Nested.runScalar pN ))+            (Concrete ((mA, vA), p))+        STKR SNat (STKScalar @r) ->+          ifDifferentiable @r+            (gcastWith (unsafeCoerceRefl :: y2 :~: ADTensorKind y2) $+             let !(!mAN, !vAN, !pN) = updateR mA vA p g+             in Concrete ((mAN, vAN), pN))+            (Concrete ((mA, vA), p))+        STKS sh (STKScalar @r) ->+          ifDifferentiable @r+            (gcastWith (unsafeCoerceRefl :: y2 :~: ADTensorKind y2) $+             let !(!mAN, !vAN, !pN) =+                   updateR (Nested.stoRanked mA)+                           (Nested.stoRanked vA)+                           (Nested.stoRanked p)+                           (Nested.stoRanked g)+             in Concrete+                ( ( Nested.rcastToShaped mAN sh+                  , Nested.rcastToShaped vAN sh )+                , Nested.rcastToShaped pN sh ))+            (Concrete ((mA, vA), p))+        STKX _ (STKScalar @r) ->+          ifDifferentiable @r+            (gcastWith (unsafeCoerceRefl :: y2 :~: ADTensorKind y2) $+             let !(!mAN, !vAN, !pN) =+                   updateR (Nested.mtoRanked mA)+                           (Nested.mtoRanked vA)+                           (Nested.mtoRanked p)+                           (Nested.mtoRanked g)+             in Concrete+                ( ( Nested.mreshape (Nested.mshape mA)+                    $ Nested.rtoMixed mAN+                  , Nested.mreshape (Nested.mshape vA)+                    $ Nested.rtoMixed vAN )+                , Nested.mreshape (Nested.mshape p)+                  $ Nested.rtoMixed pN ))+            (Concrete ((mA, vA), p))+        STKProduct stk1 stk2 ->+          let !a1 = unConcrete $ updateProd stk1+                      (Concrete $ fst mA) (Concrete $ fst vA)+                      (Concrete $ fst p) (Concrete $ fst g)+              !a2 = unConcrete $ updateProd stk2+                      (Concrete $ snd mA) (Concrete $ snd vA)+                      (Concrete $ snd p) (Concrete $ snd g)+          in Concrete (a1, a2)+        _ -> error "updateWithGradientAdam: only non-nested arrays supported"+               -- we could support nested arrays, but some new unwinding code+               -- woudl be needed, with no existing use case+      (!mAdamRNew, !vAdamRNew, !paramsRNew) =+        unzip3Rep stk0 $ updateProd stk0 mAdamR vAdamR paramsR gradientR+  in ( paramsRNew+     , StateAdam+         { tAdam = tAdamNew+         , mAdam = mAdamRNew+         , vAdam = vAdamRNew+         }+     )
+ src/HordeAd/OpsTensor.hs view
@@ -0,0 +1,1133 @@+{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | The tensor operations intended for the casual library user.+--+-- The less user-friendly+-- prototypes of most of these operation can be found in "HordeAd.Core.Ops"+-- where some additional rarely used operations reside.+-- All these operations, together with instances of numerical classes+-- such as @Num@, @Fractional@, @IntegralH@, @RealFloatH@, @EqH@ and others+-- (see class instances of type 'HordeAd.Core.Ast.AstTensor' for the full list),+-- are a major part of the high-level API of the horde-ad library,+-- which is relatively orthogonal to the other major part,+-- the differentiation interface exposed in "HordeAd.ADEngine".+module HordeAd.OpsTensor+  ( -- * Shape manipulation+    rshape, rlength, rsize, rwidth+  , sshape, slength, ssize, swidth+  , xshape, xlength, xsize, xwidth+  , tsize, tftk+    -- * Constructing arrays from concrete values, lists and vectors+  , rconcrete, rscalar, rrepl, ringestData, rfromListLinear+  , sconcrete, sscalar, srepl, singestData, sfromListLinear+  , xconcrete, xscalar, xrepl, xingestData, xfromListLinear+  , kconcrete+  , rfromList, rfromVector, rfromVector0N, rfromList0N, runravelToList+  , sfromList, sfromVector, sfromVector0N, sfromList0N, sunravelToList+  , xfromList, xfromVector, xfromVector0N, xfromList0N, xunravelToList+    -- * Main array operations+  , tunit, tlet, ifH, minH, maxH, tpair, tproject1, tproject2+  , rsum, rsum0, rdot0, rdot1In, rmatvecmul, rmatmul2, rreplicate, rreplicate0N+  , ssum, ssum0, sdot0, sdot1In, smatvecmul, smatmul2, sreplicate, sreplicate0N+  , xsum, xsum0, xdot0, xdot1In, xmatvecmul, xmatmul2, xreplicate, xreplicate0N+  , rindex, (!), rindex0, roneHot, rscatter, rscatter1, rgather, rgather1+  , sindex, (!$), sindex0, soneHot, sscatter, sscatter1, sgather, sgather1+  , xindex, xindex0, xoneHot, xscatter, xscatter1, xgather, xgather1+  , rtr, rtranspose, rflatten, rreshape+  , str, stranspose, sflatten, sreshape+  , xtr, xtranspose, xflatten, xreshape+   -- * Auxiliary array operations+  , rfloor, rfromIntegral, rcast, rminIndex, rmaxIndex, riota+  , sfloor, sfromIntegral, scast, sminIndex, smaxIndex, siota+  , xfloor, xfromIntegral, xcast, xminIndex, xmaxIndex, xiota+  , kfloor, kfromIntegral, kcast+  , rappend, rconcat, rslice, runcons, rreverse+  , sappend, sslice, suncons, sreverse+  , xappend, xappend0, xconcat, xslice, xuncons, xreverse+    -- * Array operations derived from @build@+  , rbuild, rbuild1, rmap, rmap1, rmap0N, rzipWith, rzipWith1, rzipWith0N+  , rzipWith3, rzipWith31, rzipWith30N, rzipWith4, rzipWith41, rzipWith40N+  , sbuild, sbuild1, smap, smap1, smap0N, szipWith, szipWith1, szipWith0N+  , szipWith3, szipWith31, szipWith30N, szipWith4, szipWith41, szipWith40N+  , xbuild, xbuild1+    -- * Array operations derived from @mapAccum@+  , rfold, rscan, sfold, sscan, xfold, xscan, tmapAccumR, tmapAccumL+    -- * Array operations producing derivatives+  , kgrad, rvjp, rjvp, svjp, sjvp+    -- * Operations dealing with dual numbers+  , rprimalPart, rdualPart, rfromPrimal, rfromDual, rScale+  , sprimalPart, sdualPart, sfromPrimal, sfromDual, sScale+  , xprimalPart, xdualPart, xfromPrimal, xfromDual, xScale+  , kprimalPart, kdualPart, kfromPrimal, kfromDual, kScale+    -- * Array operations that utilize unwinding of nested arrays+  , treplTarget, tdefTarget, taddTarget, tmultTarget, tsum0Target, tdot0Target+    -- * Minimal re-exports to make this module a higher level replacement for "HordeAd.Core.Ops"+  , ADReady+  , LetTensor, BaseTensor+  ) where++import Prelude++import Data.Int (Int64)+import Data.List.NonEmpty (NonEmpty)+import Data.List.NonEmpty qualified as NonEmpty+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality (gcastWith, (:~:))+import Data.Vector.Generic qualified as V+import Data.Vector.Strict qualified as Data.Vector+import GHC.TypeLits+  (KnownNat, OrderingI (..), cmpNat, type (+), type (-), type (<=))++import Data.Array.Nested (type (++))+import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Permutation qualified as Permutation+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape+import Data.Array.Nested.Types (Init, unsafeCoerceRefl)++import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor+import HordeAd.Core.Ops+import HordeAd.Core.TensorKind+import HordeAd.Core.Types++rconcrete :: (GoodScalar r, BaseTensor target)+          => Nested.Ranked n r -> target (TKR n r)+rconcrete = trconcrete+rscalar :: (GoodScalar r, BaseTensor target)+        => r -> target (TKR 0 r)+rscalar r = rconcrete $ Nested.rscalar r+rrepl :: forall n r target. (GoodScalar r, BaseTensor target)+      => IShR n -> r -> target (TKR n r)+rrepl sh a = tconcrete (FTKR sh FTKScalar)+                       (Concrete $ Nested.rreplicateScal sh a)+ringestData :: forall n r target. (GoodScalar r, BaseTensor target)+            => IShR n -> [r] -> target (TKR n r)+ringestData sh l =+  tconcrete (FTKR sh FTKScalar) (Concrete $ Nested.rfromListPrimLinear sh l)+-- | Create a tensor from a list of individual underlying scalar values.+--+-- An operation with the same name is used in @ox-arrays@ to represent+-- pretty-printed concrete arrays. The types agree, so the same representation+-- results in analogous horde-ad array when this operation is used instead.+rfromListLinear :: forall n r target. (GoodScalar r, BaseTensor target)+                => IShR n -> NonEmpty r -> target (TKR n r)+rfromListLinear sh = ringestData sh . NonEmpty.toList++sconcrete :: (GoodScalar r, BaseTensor target)+          => Nested.Shaped sh r -> target (TKS sh r)+sconcrete = tsconcrete+sscalar :: (GoodScalar r, BaseTensor target)+        => r -> target (TKS '[] r)+sscalar r = sconcrete $ Nested.sscalar r+srepl :: (KnownShS sh, GoodScalar r, BaseTensor target)+      => r -> target (TKS sh r)+srepl = sconcrete . Nested.sreplicateScal knownShS+singestData :: (KnownShS sh, GoodScalar r, BaseTensor target)+            => [r] -> target (TKS sh r)+singestData l = sconcrete $ Nested.sfromListPrimLinear knownShS l+-- | Create a tensor from a list of individual underlying scalar values.+--+-- An operation with the same name is used in @ox-arrays@ to represent+-- pretty-printed concrete arrays. The types agree, so the same representation+-- results in analogous horde-ad array when this operation is used instead.+sfromListLinear :: forall sh r target. (GoodScalar r, BaseTensor target)+                => ShS sh -> NonEmpty r -> target (TKS sh r)+sfromListLinear sh = sconcrete . Nested.sfromListPrimLinear sh . NonEmpty.toList++xconcrete :: (GoodScalar r, BaseTensor target)+          => Nested.Mixed sh r -> target (TKX sh r)+xconcrete = txconcrete+xscalar :: (GoodScalar r, BaseTensor target)+        => r -> target (TKX '[] r)+xscalar r = xconcrete $ Nested.mscalar r+xrepl :: forall sh r target. (GoodScalar r, BaseTensor target)+      => IShX sh -> r -> target (TKX sh r)+xrepl sh a = tconcrete (FTKX sh FTKScalar)+                       (Concrete $ Nested.mreplicateScal sh a)+xingestData :: forall sh r target. (GoodScalar r, BaseTensor target)+            => IShX sh -> [r] -> target (TKX sh r)+xingestData sh l =+  tconcrete (FTKX sh FTKScalar) (Concrete $ Nested.mfromListPrimLinear sh l)+xfromListLinear :: forall sh r target. (GoodScalar r, BaseTensor target)+                => IShX sh -> NonEmpty r -> target (TKX sh r)+xfromListLinear sh = xingestData sh . NonEmpty.toList++kconcrete :: (GoodScalar r, BaseTensor target)+          => r -> target (TKScalar r)+kconcrete = tkconcrete++-- | Create a tensor from a list treated as the outermost dimension,+-- going through strict boxed vectors, because laziness is risky with+-- impurity, e.g., it easily perturbs results of fragile tests.+rfromList :: (KnownNat n, KnownSTK x, BaseTensor target)+          => NonEmpty (target (TKR2 n x)) -> target (TKR2 (1 + n) x)+rfromList = trfromVector . V.fromList . NonEmpty.toList+-- | Create a tensor from a non-empty strict boxed vector treated+-- as the outermost dimension.+rfromVector :: (KnownNat n, KnownSTK x, BaseTensor target)+            => Data.Vector.Vector (target (TKR2 n x))+            -> target (TKR2 (1 + n) x)+rfromVector = trfromVector+rfromVector0N :: forall n x target. (KnownSTK x, BaseTensor target)+              => IShR n -> Data.Vector.Vector (target (TKR2 0 x))+              -> target (TKR2 n x)+rfromVector0N = trfromVector0N+rfromList0N :: forall n x target. (KnownSTK x, BaseTensor target)+            => IShR n -> [target (TKR2 0 x)]+            -> target (TKR2 n x)+rfromList0N sh = trfromVector0N sh . V.fromList+-- | Unravel a tensor into a list of its immediate subtensors.+--+-- Warning: during computation, sharing between the elements+-- of the resulting list is likely to be lost, so it needs to be ensured+-- by explicit sharing, e.g., 'ttlet'.+runravelToList :: forall n x target.+                  (KnownSTK x, KnownNat n, BaseTensor target)+               => target (TKR2 (1 + n) x) -> [target (TKR2 n x)]+runravelToList = trunravelToList++-- | Create a tensor from a list treated as the outermost dimension,+-- going through strict boxed vectors, because laziness is risky with+-- impurity, e.g., it easily perturbs results of fragile tests.+sfromList :: (KnownNat n, KnownShS sh, KnownSTK x, BaseTensor target)+          => NonEmpty (target (TKS2 sh x)) -> target (TKS2 (n ': sh) x)+sfromList = tsfromVector . V.fromList . NonEmpty.toList+-- | Create a tensor from a non-empty strict boxed vector treated+-- as the outermost dimension.+sfromVector :: (KnownNat n, KnownShS sh, KnownSTK x, BaseTensor target)+            => Data.Vector.Vector (target (TKS2 sh x))+            -> target (TKS2 (n ': sh) x)+sfromVector = tsfromVector+sfromVector0N :: (KnownShS sh, KnownSTK x, BaseTensor target)+              => Data.Vector.Vector (target (TKS2 '[] x))+              -> target (TKS2 sh x)+sfromVector0N = tsfromVector0N+sfromList0N :: (KnownShS sh, KnownSTK x, BaseTensor target)+            => [target (TKS2 '[] x)]+            -> target (TKS2 sh x)+sfromList0N = tsfromVector0N . V.fromList+-- | Unravel a tensor into a list of its immediate subtensors.+--+-- Warning: during computation, sharing between the elements+-- of the resulting list is likely to be lost, so it needs to be ensured+-- by explicit sharing, e.g., 'ttlet'.+sunravelToList :: (KnownNat n, KnownShS sh, KnownSTK x, BaseTensor target)+               => target (TKS2 (n ': sh) x) -> [target (TKS2 sh x)]+sunravelToList = tsunravelToList++-- | Create a tensor from a list treated as the outermost dimension,+-- going through strict boxed vectors, because laziness is risky with+-- impurity, e.g., it easily perturbs results of fragile tests.+xfromList :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)+          => NonEmpty (target (TKX2 sh x)) -> target (TKX2 (Just n ': sh) x)+xfromList = txfromVector . V.fromList . NonEmpty.toList+-- | Create a tensor from a non-empty strict boxed vector treated+-- as the outermost dimension.+xfromVector :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)+            => Data.Vector.Vector (target (TKX2 sh x))+            -> target (TKX2 (Just n ': sh) x)+xfromVector = txfromVector+xfromVector0N :: forall sh x target. (KnownSTK x, BaseTensor target)+              => IShX sh -> Data.Vector.Vector (target (TKX2 '[] x))+              -> target (TKX2 sh x)+xfromVector0N = txfromVector0N+xfromList0N :: forall sh x target. (KnownSTK x, BaseTensor target)+            => IShX sh -> [target (TKX2 '[] x)]+            -> target (TKX2 sh x)+xfromList0N sh = txfromVector0N sh . V.fromList+-- | Unravel a tensor into a list of its immediate subtensors.+--+-- Warning: during computation, sharing between the elements+-- of the resulting list is likely to be lost, so it needs to be ensured+-- by explicit sharing, e.g., 'ttlet'.+xunravelToList :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)+               => target (TKX2 (Just n ': sh) x) -> [target (TKX2 sh x)]+xunravelToList = txunravelToList++tunit :: BaseTensor target+      => target TKUnit+tunit = kconcrete Z1++tlet :: forall x z target. LetTensor target+     => target x -> (target x -> target z) -> target z+tlet = ttlet++-- | The operation is potentially strict in all arguments.+ifH :: (KnownSTK y, Boolean (BoolOf target), BaseTensor target)+    => BoolOf target -> target y -> target y -> target y+ifH = tcond knownSTK+minH :: (KnownSTK y, OrdH target y, BaseTensor target)+     => target y -> target y -> target y+minH u v = ifH (u <=. v) u v+maxH :: (KnownSTK y, OrdH target y, BaseTensor target)+     => target y -> target y -> target y+maxH u v = ifH (u >=. v) u v++rsum :: (KnownNat n, KnownSTK x, BaseTensor target)+     => target (TKR2 (1 + n) x) -> target (TKR2 n x)+rsum = trsum+rsum0 :: (KnownNat n, KnownSTK x, BaseTensor target)+      => target (TKR2 n x) -> target (TKR2 0 x)+rsum0 = trsum0+rdot0 :: ( KnownNat n, GoodScalar r, BaseTensor target)+      => target (TKR n r) -> target (TKR n r) -> target (TKR 0 r)+rdot0 = trdot0+rdot1In :: (KnownNat n, GoodScalar r, BaseTensor target)+        => target (TKR (1 + n) r) -> target (TKR (1 + n) r)+        -> target (TKR n r)+rdot1In = trdot1In+rmatvecmul :: (GoodScalar r, BaseTensor target)+           => target (TKR 2 r) -> target (TKR 1 r) -> target (TKR 1 r)+rmatvecmul = trmatvecmul+rmatmul2 :: (GoodScalar r, BaseTensor target)+         => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)+rmatmul2 = trmatmul2+-- | Copy the given tensor along the new, outermost dimension.+rreplicate :: (KnownNat n, KnownSTK x, BaseTensor target)+           => Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)+rreplicate = trreplicate+rreplicate0N :: (KnownNat n, KnownSTK x, BaseTensor target)+             => IShR n -> target (TKR2 0 x) -> target (TKR2 n x)+rreplicate0N = trreplicate0N++ssum :: (KnownNat n, KnownShS sh, KnownSTK x, BaseTensor target)+     => target (TKS2 (n ': sh) x) -> target (TKS2 sh x)+ssum = tssum+ssum0 :: (KnownShS sh, KnownSTK x, BaseTensor target)+      => target (TKS2 sh x) -> target (TKS2 '[] x)+ssum0 = tssum0+sdot0 :: (KnownShS sh, GoodScalar r, BaseTensor target)+      => target (TKS sh r) -> target (TKS sh r) -> target (TKS '[] r)+sdot0 = tsdot0+sdot1In :: (KnownShS sh, KnownNat n, GoodScalar r, BaseTensor target)+        => target (TKS (sh ++ '[n]) r) -> target (TKS (sh ++ '[n]) r)+        -> target (TKS sh r)+sdot1In @sh @n = tsdot1In @_ @sh (SNat @n)+smatvecmul :: (KnownNat m, KnownNat n, GoodScalar r, BaseTensor target)+           => target (TKS '[m, n] r) -> target (TKS '[n] r)+           -> target (TKS '[m] r)+smatvecmul = tsmatvecmul+smatmul2 :: ( KnownNat m, KnownNat n, KnownNat p+            , GoodScalar r, BaseTensor target )+         => target (TKS '[m, n] r) -> target (TKS '[n, p] r)+         -> target (TKS '[m, p] r)+smatmul2 = tsmatmul2+sreplicate :: (KnownNat k, KnownShS sh, KnownSTK x, BaseTensor target)+           => target (TKS2 sh x) -> target (TKS2 (k ': sh) x)+sreplicate = tsreplicate SNat knownShS+sreplicate0N :: (KnownShS sh, KnownSTK x, BaseTensor target)+             => target (TKS2 '[] x) -> target (TKS2 sh x)+sreplicate0N = tsreplicate0N knownShS++xsum :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)+     => target (TKX2 (Just n ': sh) x) -> target (TKX2 sh x)+xsum = txsum+xsum0 :: (KnownShX sh, KnownSTK x, BaseTensor target, ConvertTensor target)+      => target (TKX2 sh x) -> target (TKX2 '[] x)+xsum0 = txsum0+xdot0 :: ( KnownShX sh, GoodScalar r+         , BaseTensor target, ConvertTensor target )+      => target (TKX sh r) -> target (TKX sh r) -> target (TKX '[] r)+xdot0 = txdot0+xdot1In :: (KnownShX sh, KnownNat n, GoodScalar r, BaseTensor target)+        => target (TKX (sh ++ '[Just n]) r)+        -> target (TKX (sh ++ '[Just n]) r)+        -> target (TKX sh r)+xdot1In @sh @n = txdot1In @_ @sh (SNat @n)+xmatvecmul :: forall mm mn r target.+              (GoodScalar r, BaseTensor target, ConvertTensor target)+           => Nested.SMayNat Int SNat mm -> Nested.SMayNat Int SNat mn+           -> target (TKX '[mm, mn] r) -> target (TKX '[mn] r)+           -> target (TKX '[mm] r)+xmatvecmul = txmatvecmul+xmatmul2 :: ( KnownNat m, KnownNat n, KnownNat p+            , GoodScalar r, BaseTensor target, ConvertTensor target )+         => target (TKX '[Just m, Just n] r)+         -> target (TKX '[Just n, Just p] r)+         -> target (TKX '[Just m, Just p] r)+xmatmul2 = txmatmul2+xreplicate :: (KnownNat k, KnownShX sh, KnownSTK x, BaseTensor target)+           => target (TKX2 sh x) -> target (TKX2 (Just k ': sh) x)+xreplicate = txreplicate SNat knownShX+xreplicate0N :: (KnownShX sh, KnownSTK x, BaseTensor target)+             => IShX sh -> target (TKX2 '[] x) -> target (TKX2 sh x)+xreplicate0N = txreplicate0N++-- | First index is for outermost dimension; empty index means identity,+-- if index is out of bounds, the result is defined and is @def@, which is @0@.+-- The same holds for all operations with indexes.+rindex, (!) :: (KnownNat m, KnownNat n, KnownSTK x, BaseTensor target)+            => target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)+rindex = trindex+infixl 9 !+(!) = rindex  -- prefix form better when type applications are necessary+rindex0 :: (KnownNat m, KnownSTK x, BaseTensor target)+        => target (TKR2 m x) -> IxROf target m -> target (TKR2 0 x)+rindex0 = trindex0+roneHot :: ( KnownNat m, KnownNat n, KnownSTK x+           , BoolOf (PrimalOf target) ~ BoolOf target+           , EqH (PrimalOf target) (TKScalar Int64), BaseTensor target)+        => IShR m -> target (TKR2 n x) -> IxROf target m+        -> target (TKR2 (m + n) x)+roneHot = troneHot+rscatter :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x, BaseTensor target)+         => IShR (p + n) -> target (TKR2 (m + n) x)+         -> (IxROf target m -> IxROf target p)+         -> target (TKR2 (p + n) x)+rscatter = trscatter+-- | Build a tensor by adding up tensors of rank @n@ taken from+-- the second argument and inserted in a zero tensor+-- at indexes of length @p@.+-- The semantics of the operation permits index out of bounds+-- and then no tensor is added at such an index.+rscatter1 :: (KnownNat n, KnownNat p, KnownSTK x, BaseTensor target)+          => IShR (p + n) -> target (TKR2 (1 + n) x)+          -> (IntOf target -> IxROf target p)+          -> target (TKR2 (p + n) x)+rscatter1 = trscatter1+rgather :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x, BaseTensor target)+        => IShR (m + n) -> target (TKR2 (p + n) x)+        -> (IxROf target m -> IxROf target p)+        -> target (TKR2 (m + n) x)+rgather = trgather+-- | Build a tensor by collecting tensors of rank @n@ obtained by indexing+-- in the second argument at the given indexes of length @p@.+-- The semantics of the operation permits index out of bounds+-- and the result of such indexing is @def@, which is @0@.+rgather1 :: (KnownNat n, KnownNat p, KnownSTK x, BaseTensor target)+         => Int -> target (TKR2 (p + n) x)+         -> (IntOf target -> IxROf target p)+         -> target (TKR2 (1 + n) x)+rgather1 = trgather1++sindex, (!$) :: (KnownShS shm, KnownShS shn, KnownSTK x, BaseTensor target)+             => target (TKS2 (shm ++ shn) x) -> IxSOf target shm+             -> target (TKS2 shn x)+sindex = tsindex+infixl 9 !$+(!$) = sindex  -- prefix form better when type applications are necessary+sindex0 :: (KnownShS sh1, KnownSTK x, BaseTensor target)+        => target (TKS2 sh1 x) -> IxSOf target sh1+        -> target (TKS2 '[] x)+sindex0 = tsindex0+soneHot :: ( KnownShS sh1, KnownShS sh2, KnownSTK x+           , BoolOf (PrimalOf target) ~ BoolOf target+           , EqH (PrimalOf target) (TKScalar Int64), BaseTensor target )+        => target (TKS2 sh2 x) -> IxSOf target sh1+        -> target (TKS2 (sh1 ++ sh2) x)+soneHot = tsoneHot+sscatter+  :: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x, BaseTensor target)+  => target (TKS2 (shm ++ shn) x)+  -> (IxSOf target shm -> IxSOf target shp)+  -> target (TKS2 (shp ++ shn) x)+sscatter @shm @shn @shp = tsscatter @_ @shm @shn @shp+sscatter1+  :: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x, BaseTensor target)+  => target (TKS2 (n2 ': shn) x)+  -> (IntOf target -> IxSOf target shp)+  -> target (TKS2 (shp ++ shn) x)+sscatter1 = tsscatter1+sgather+  :: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x, BaseTensor target)+  => target (TKS2 (shp ++ shn) x)+  -> (IxSOf target shm -> IxSOf target shp)+  -> target (TKS2 (shm ++ shn) x)+sgather @shm @shn @shp = tsgather @_ @shm @shn @shp+sgather1+  :: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x, BaseTensor target)+  => target (TKS2 (shp ++ shn) x)+  -> (IntOf target -> IxSOf target shp)+  -> target (TKS2 (n2 ': shn) x)+sgather1 = tsgather1++xindex :: (KnownShX sh1, KnownShX sh2, KnownSTK x, BaseTensor target)+       => target (TKX2 (sh1 ++ sh2) x) -> IxXOf target sh1+       -> target (TKX2 sh2 x)+xindex = txindex+xindex0 :: (KnownShX sh1, KnownSTK x, BaseTensor target)+        => target (TKX2 sh1 x) -> IxXOf target sh1+        -> target (TKX2 '[] x)+xindex0 = txindex0+xoneHot :: ( KnownShX sh1, KnownShX sh2, KnownSTK x+           , BoolOf (PrimalOf target) ~ BoolOf target+           , EqH (PrimalOf target) (TKScalar Int64)+           , BaseTensor target, ConvertTensor target )+        => IShX sh1 -> target (TKX2 sh2 x) -> IxXOf target sh1+        -> target (TKX2 (sh1 ++ sh2) x)+xoneHot = txoneHot+xscatter :: ( KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x+            , BaseTensor target )+         => IShX (shp ++ shn) -> target (TKX2 (shm ++ shn) x)+         -> (IxXOf target shm -> IxXOf target shp)+         -> target (TKX2 (shp ++ shn) x)+xscatter @shm @shn @shp = txscatter @_ @shm @shn @shp+xscatter1 :: ( KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x+             , BaseTensor target )+          => IShX (shp ++ shn) -> target (TKX2 (Just n2 ': shn) x)+          -> (IntOf target -> IxXOf target shp)+          -> target (TKX2 (shp ++ shn) x)+xscatter1 = txscatter1+xgather :: ( KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x+           , BaseTensor target )+        => IShX (shm ++ shn)+        -> target (TKX2 (shp ++ shn) x)+        -> (IxXOf target shm -> IxXOf target shp)+        -> target (TKX2 (shm ++ shn) x)+xgather @shm @shn @shp = txgather @_ @shm @shn @shp+xgather1 :: ( KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x+            , BaseTensor target )+         => SNat n2 -> target (TKX2 (shp ++ shn) x)+         -> (IntOf target -> IxXOf target shp)+         -> target (TKX2 (Just n2 ': shn) x)+xgather1 = txgather1++-- | Transpose according to the permutation.+rtranspose :: forall n x target. (KnownSTK x, BaseTensor target)+           => Permutation.PermR -> target (TKR2 n x) -> target (TKR2 n x)+rtranspose = trtranspose+-- | Change the shape of the tensor to the given one.+rreshape :: forall n m x target. (KnownSTK x, BaseTensor target)+         => IShR m -> target (TKR2 n x) -> target (TKR2 m x)+rreshape = trreshape+stranspose :: ( Permutation.KnownPerm perm, Permutation.IsPermutation perm+              , Rank perm <= Rank sh, KnownSTK x, BaseTensor target )+           => target (TKS2 sh x)+           -> target (TKS2 (Permutation.PermutePrefix perm sh) x)+stranspose @perm = tstranspose (Permutation.makePerm @perm)+sreshape :: ( Nested.Product sh ~ Nested.Product sh2, KnownShS sh2+            , KnownSTK x, BaseTensor target )+         => target (TKS2 sh x) -> target (TKS2 sh2 x)+sreshape = tsreshape knownShS+xtranspose :: ( Permutation.KnownPerm perm, Permutation.IsPermutation perm+              , Rank perm <= Rank sh, KnownSTK x, BaseTensor target )+           => target (TKX2 sh x)+           -> target (TKX2 (Permutation.PermutePrefix perm sh) x)+xtranspose @perm = txtranspose (Permutation.makePerm @perm)+xreshape :: forall sh sh2 x target. (KnownSTK x, BaseTensor target)+         => IShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)+xreshape = txreshape++rfloor :: ( GoodScalar r, RealFrac r, GoodScalar r2, Integral r2+          , BaseTensor target )+       => target (TKR n r) -> target (TKR n r2)+rfloor = trfloor+rfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2, BaseTensor target)+              => target (TKR n r1) -> target (TKR n r2)+rfromIntegral = trfromIntegral+rcast :: ( RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2+         , BaseTensor target )+      => target (TKR n r1) -> target (TKR n r2)+rcast = trcast+rminIndex, rmaxIndex  -- partial+  :: forall n r r2 target. (GoodScalar r, GoodScalar r2, BaseTensor target)+  => target (TKR (1 + n) r) -> target (TKR n r2)+rminIndex = trminIndex+rmaxIndex = trmaxIndex+riota :: (GoodScalar r, BaseTensor target)+      => Int -> target (TKR 1 r)  -- from 0 to n - 1+riota = triota++sfloor :: ( GoodScalar r, RealFrac r, GoodScalar r2, Integral r2+          , BaseTensor target )+       => target (TKS sh r) -> target (TKS sh r2)+sfloor = tsfloor+sfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2, BaseTensor target)+              => target (TKS sh r1) -> target (TKS sh r2)+sfromIntegral = tsfromIntegral+scast :: ( RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2+         , BaseTensor target )+      => target (TKS sh r1) -> target (TKS sh r2)+scast = tscast+sminIndex, smaxIndex  -- partial+  :: forall n sh r r2 target. (GoodScalar r, GoodScalar r2, BaseTensor target)+  => target (TKS (n ': sh) r) -> target (TKS (Init (n ': sh)) r2)+sminIndex = tsminIndex+smaxIndex = tsmaxIndex+siota :: (KnownNat n, GoodScalar r, BaseTensor target)+      => target (TKS '[n] r)  -- from 0 to n - 1+siota = tsiota++xfloor :: ( GoodScalar r, RealFrac r, GoodScalar r2, Integral r2+          , BaseTensor target )+       => target (TKX sh r) -> target (TKX sh r2)+xfloor = txfloor+xfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2, BaseTensor target)+              => target (TKX sh r1) -> target (TKX sh r2)+xfromIntegral = txfromIntegral+xcast :: ( RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2+         , BaseTensor target )+      => target (TKX sh r1) -> target (TKX sh r2)+xcast = txcast+xminIndex, xmaxIndex  -- partial+  :: forall mn sh r r2 target. (GoodScalar r, GoodScalar r2, BaseTensor target)+  => target (TKX (mn ': sh) r) -> target (TKX (Init (mn ': sh)) r2)+xminIndex = txminIndex+xmaxIndex = txmaxIndex+xiota :: (KnownNat n, GoodScalar r, BaseTensor target)+      => target (TKX '[Just n] r)  -- from 0 to n - 1+xiota = txiota++kfloor :: ( GoodScalar r, RealFrac r, GoodScalar r2, Integral r2+          , BaseTensor target )+       => target (TKScalar r) -> target (TKScalar r2)+kfloor = tkfloor+kfromIntegral :: (GoodScalar r1, Integral r1, GoodScalar r2, BaseTensor target)+              => target (TKScalar r1) -> target (TKScalar r2)+kfromIntegral = tkfromIntegral+kcast :: ( RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2+         , BaseTensor target )+      => target (TKScalar r1) -> target (TKScalar r2)+kcast = tkcast++-- | Append two arrays along the outermost dimension.+-- All dimensions, except the outermost, must be the same.+rappend :: forall n x target. (KnownSTK x, BaseTensor target)+        => target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)+        -> target (TKR2 (1 + n) x)+rappend = trappend+rconcat :: forall n x target. (KnownSTK x, BaseTensor target)+        => NonEmpty (target (TKR2 (1 + n) x)) -> target (TKR2 (1 + n) x)+rconcat = foldr1 rappend+-- | Extract a slice of an array along the outermost dimension.+-- The extracted slice must fall within the dimension.+rslice :: forall n x target. (KnownSTK x, BaseTensor target)+       => Int -> Int -> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)+rslice = trslice+runcons :: (KnownNat n, KnownSTK x, BaseTensor target)+        => target (TKR2 (1 + n) x)+        -> Maybe (target (TKR2 n x), target (TKR2 (1 + n) x))+runcons v = case rshape v of+              len :$: _ -> Just (v ! [0], rslice 1 (len - 1) v)+-- | Reverse elements of the outermost dimension.+rreverse :: forall n x target. (KnownSTK x, BaseTensor target)+         => target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)+rreverse = trreverse++sappend :: forall m n sh x target. (KnownSTK x, BaseTensor target)+        => target (TKS2 (m ': sh) x) -> target (TKS2 (n ': sh) x)+        -> target (TKS2 ((m + n) ': sh) x)+sappend = tsappend+sslice :: forall i n k sh x target. (KnownSTK x, BaseTensor target)+       => SNat i -> SNat n -> SNat k+       -> target (TKS2 (i + n + k ': sh) x) -> target (TKS2 (n ': sh) x)+sslice = tsslice+suncons :: (KnownNat n, KnownShS sh, KnownSTK x, BaseTensor target)+        => target (TKS2 (n ': sh) x)+        -> Maybe (target (TKS2 sh x), target (TKS2 (n - 1 ': sh) x))+suncons @n v = case cmpNat (Proxy @1) (Proxy @n) of+ EQI -> Just ( v !$ (0 :.$ ZIS)+             , sslice @1 @(n - 1) @0 SNat SNat SNat v )+ LTI -> Just ( v !$ (0 :.$ ZIS)+             , sslice @1 @(n - 1) @0 SNat SNat SNat v )+ _ -> Nothing+sreverse :: forall n sh x target. (KnownSTK x, BaseTensor target)+         => target (TKS2 (n ': sh) x) -> target (TKS2 (n ': sh) x)+sreverse = tsreverse++xappend :: forall m n sh x target. (KnownSTK x, BaseTensor target)+        => target (TKX2 (Just m ': sh) x) -> target (TKX2 (Just n ': sh) x)+        -> target (TKX2 (Just (m + n) ': sh) x)+xappend = txappend+xappend0 :: forall sh x target.+            (KnownSTK x, BaseTensor target, ConvertTensor target)+         => target (TKX2 (Nothing ': sh) x) -> target (TKX2 (Nothing ': sh) x)+         -> target (TKX2 (Nothing ': sh) x)+xappend0 a b = case xshape a of+  mmsnat :$% sh ->+    withSNat (fromSMayNat' mmsnat) $ \msnat ->+    withSNat (shxLength $ xshape b) $ \nsnat ->+    let sh0 = Nested.SUnknown () :!% ssxFromShX sh+        sha = Nested.SKnown msnat :!% ssxFromShX sh+        shb = Nested.SKnown nsnat :!% ssxFromShX sh+    in withKnownShX (ssxFromShX sh) $+       xmcast sh0 $ xappend (xmcast sha a) (xmcast shb b)+xconcat :: forall sh x target.+           (KnownSTK x, BaseTensor target, ConvertTensor target)+        => NonEmpty (target (TKX2 (Nothing ': sh) x))+        -> target (TKX2 (Nothing ': sh) x)+xconcat = foldr1 xappend0+xslice :: forall i n k sh x target. (KnownSTK x, BaseTensor target)+       => SNat i -> SNat n -> SNat k+       -> target (TKX2 (Just (i + n + k) ': sh) x)+       -> target (TKX2 (Just n ': sh) x)+xslice = txslice+xuncons :: (KnownNat n, KnownShX sh, KnownSTK x, BaseTensor target)+        => target (TKX2 (Just n ': sh) x)+        -> Maybe (target (TKX2 sh x), target (TKX2 (Just (n - 1) ': sh) x))+xuncons @n v = case cmpNat (Proxy @1) (Proxy @n) of+  EQI -> Just ( v `xindex` (0 :.% ZIX)+              , xslice @1 @(n - 1) @0 SNat SNat SNat v )+  LTI -> Just ( v `xindex` (0 :.% ZIX)+              , xslice @1 @(n - 1) @0 SNat SNat SNat v )+  _ -> Nothing+xreverse :: forall mn sh x target. (KnownSTK x, BaseTensor target)+         => target (TKX2 (mn ': sh) x) -> target (TKX2 (mn ': sh) x)+xreverse = txreverse++rbuild1 :: (KnownNat n, KnownSTK x, BaseTensor target)+        => Int  -- ^ width of the outermost dimension of the created tensor+        -> (IntOf target -> target (TKR2 n x))  -- ^ the function to build with+        -> target (TKR2 (1 + n) x)+rbuild1 = trbuild1+rmap :: (KnownNat m, KnownNat n, KnownSTK x, KnownSTK x2, BaseTensor target)+     => (target (TKR2 n x) -> target (TKR2 n x2))  -- ^ the function to map with+     -> target (TKR2 (m + n) x)  -- ^ the tensor to map over+     -> target (TKR2 (m + n) x2)+rmap f v = rbuild (rshape v) (\ix -> f (v ! ix))+rmap1 :: (KnownNat n, KnownSTK x, KnownSTK x2, BaseTensor target)+      => (target (TKR2 n x) -> target (TKR2 n x2))+            -- ^ the function to map with+      -> target (TKR2 (1 + n) x)  -- ^ the tensor to map over+      -> target (TKR2 (1 + n) x2)+rmap1 f u = rbuild1 (rwidth u) (\i -> f (u ! [i]))+rmap0N :: (KnownNat n, KnownSTK x, KnownSTK x1, BaseTensor target)+       => (target (TKR2 0 x1) -> target (TKR2 0 x))+            -- ^ the function to map with+       -> target (TKR2 n x1)  -- ^ the tensor to map over+       -> target (TKR2 n x)+rmap0N = trmap0N+rzipWith :: ( KnownNat m, KnownNat n1, KnownNat n2, KnownNat n, KnownSTK x+            , KnownSTK x1, KnownSTK x2, BaseTensor target )+         => IShR (m + n)  -- ^ the shape of the resulting tensor+         -> (target (TKR2 n1 x1) -> target (TKR2 n2 x2) -> target (TKR2 n x))+              -- ^ the function to zip with+         -> target (TKR2 (m + n1) x1)  -- ^ the first tensor to zip over+         -> target (TKR2 (m + n2) x2)  -- ^ the second tensor to zip over+         -> target (TKR2 (m + n) x)+rzipWith sh f u v = rbuild sh (\ix -> f (u ! ix) (v ! ix))+rzipWith1 :: ( KnownNat n1, KnownNat n2, KnownNat n, KnownSTK x+             , KnownSTK x1, KnownSTK x2, BaseTensor target)+          => (target (TKR2 n1 x1) -> target (TKR2 n2 x2) -> target (TKR2 n x))+               -- ^ the function to zip with+          -> target (TKR2 (1 + n1) x1)  -- ^ the first tensor to zip over+          -> target (TKR2 (1 + n2) x2)  -- ^ the second tensor to zip over+          -> target (TKR2 (1 + n) x)+rzipWith1 f u v = rbuild1 (rwidth u) (\i -> f (u ! [i]) (v ! [i]))+rzipWith0N :: ( KnownNat n, KnownSTK x, KnownSTK x1, KnownSTK x2+              , BaseTensor target )+           => (target (TKR2 0 x1) -> target (TKR2 0 x2) -> target (TKR2 0 x))+                -- ^ the function to zip with+           -> target (TKR2 n x1)  -- ^ the first tensor to zip over+           -> target (TKR2 n x2)  -- ^ the second tensor to zip over+           -> target (TKR2 n x)+rzipWith0N  = trzipWith0N+rzipWith3 :: ( KnownNat m, KnownNat n1, KnownNat n2, KnownNat n3+             , KnownNat n, KnownSTK x+             , KnownSTK x1, KnownSTK x2, KnownSTK x3, BaseTensor target )+          => IShR (m + n)  -- ^ the shape of the resulting tensor+          -> (target (TKR2 n1 x1) -> target (TKR2 n2 x2) -> target (TKR2 n3 x3)+              -> target (TKR2 n x))  -- ^ the function to zip with+          -> target (TKR2 (m + n1) x1)  -- ^ the first tensor to zip over+          -> target (TKR2 (m + n2) x2)  -- ^ the second tensor to zip over+          -> target (TKR2 (m + n3) x3)  -- ^ the third tensor to zip over+          -> target (TKR2 (m + n) x)+rzipWith3 sh f u v w = rbuild sh (\ix -> f (u ! ix) (v ! ix) (w ! ix))+rzipWith31 :: ( KnownNat n1, KnownNat n2, KnownNat n3, KnownNat n, KnownSTK x+              , KnownSTK x1, KnownSTK x2, KnownSTK x3, BaseTensor target )+           => (target (TKR2 n1 x1) -> target (TKR2 n2 x2) -> target (TKR2 n3 x3)+               -> target (TKR2 n x))  -- ^ the function to zip with+           -> target (TKR2 (1 + n1) x1)  -- ^ the first tensor to zip over+           -> target (TKR2 (1 + n2) x2)  -- ^ the second tensor to zip over+           -> target (TKR2 (1 + n3) x3)  -- ^ the third tensor to zip over+           -> target (TKR2 (1 + n) x)+rzipWith31 f u v w =+  rbuild1 (rwidth u) (\i -> f (u ! [i]) (v ! [i]) (w ! [i]))+rzipWith30N :: ( KnownNat n, KnownSTK x+               , KnownSTK x1, KnownSTK x2, KnownSTK x3, BaseTensor target )+            => (target (TKR2 0 x1) -> target (TKR2 0 x2) -> target (TKR2 0 x3)+                -> target (TKR2 0 x))  -- ^ the function to zip with+            -> target (TKR2 n x1)  -- ^ the first tensor to zip over+            -> target (TKR2 n x2)  -- ^ the second tensor to zip over+            -> target (TKR2 n x3)  -- ^ the third tensor to zip over+            -> target (TKR2 n x)+rzipWith30N f u v w =+  rbuild (rshape v) (\ix -> f (rindex0 u ix) (rindex0 v ix) (rindex0 w ix))+rzipWith4 :: ( KnownNat m+             , KnownNat n1, KnownNat n2, KnownNat n3, KnownNat n4+             , KnownNat n, KnownSTK x+             , KnownSTK x1, KnownSTK x2, KnownSTK x3, KnownSTK x4+             , BaseTensor target )+          => IShR (m + n)  -- ^ the shape of the resulting tensor+          -> (target (TKR2 n1 x1) -> target (TKR2 n2 x2)+              -> target (TKR2 n3 x3) -> target (TKR2 n4 x4)+              -> target (TKR2 n x))  -- ^ the function to zip with+          -> target (TKR2 (m + n1) x1)  -- ^ the first tensor to zip over+          -> target (TKR2 (m + n2) x2)  -- ^ the second tensor to zip over+          -> target (TKR2 (m + n3) x3)  -- ^ the third tensor to zip over+          -> target (TKR2 (m + n4) x4)  -- ^ the fourth tensor to zip over+          -> target (TKR2 (m + n) x)+rzipWith4 sh f u v w x =+  rbuild sh (\ix -> f (u ! ix) (v ! ix) (w ! ix) (x ! ix))+rzipWith41 :: ( KnownNat n1, KnownNat n2, KnownNat n3, KnownNat n4+              , KnownNat n, KnownSTK x+              , KnownSTK x1, KnownSTK x2, KnownSTK x3, KnownSTK x4+              , BaseTensor target )+           => (target (TKR2 n1 x1) -> target (TKR2 n2 x2)+               -> target (TKR2 n3 x3) -> target (TKR2 n4 x4)+               -> target (TKR2 n x))  -- ^ the function to zip with+           -> target (TKR2 (1 + n1) x1)  -- ^ the first tensor to zip over+           -> target (TKR2 (1 + n2) x2)  -- ^ the second tensor to zip over+           -> target (TKR2 (1 + n3) x3)  -- ^ the third tensor to zip over+           -> target (TKR2 (1 + n4) x4)  -- ^ the fourth tensor to zip over+           -> target (TKR2 (1 + n) x)+rzipWith41 f u v w x =+  rbuild1 (rwidth u) (\i -> f (u ! [i]) (v ! [i]) (w ! [i]) (x ! [i]))+rzipWith40N :: ( KnownNat n, KnownSTK x+               , KnownSTK x1, KnownSTK x2, KnownSTK x3, KnownSTK x4+               , BaseTensor target )+            => (target (TKR2 0 x1) -> target (TKR2 0 x2)+                -> target (TKR2 0 x3) -> target (TKR2 0 x4)+                -> target (TKR2 0 x))  -- ^ the function to zip with+            -> target (TKR2 n x1)  -- ^ the first tensor to zip over+            -> target (TKR2 n x2)  -- ^ the second tensor to zip over+            -> target (TKR2 n x3)  -- ^ the third tensor to zip over+            -> target (TKR2 n x4)  -- ^ the fourth tensor to zip over+            -> target (TKR2 n x)+rzipWith40N f u v w x =+  rbuild (rshape v) (\ix -> f (rindex0 u ix) (rindex0 v ix) (rindex0 w ix)+                              (rindex0 x ix))++sbuild1 :: (KnownNat k, KnownShS sh, KnownSTK x, BaseTensor target)+        => (IntOf target -> target (TKS2 sh x))  -- ^ the function to build with+        -> target (TKS2 (k ': sh) x)+sbuild1 = tsbuild1+smap :: ( KnownShS (Take m sh), KnownShS (Drop m sh), KnownShS sh+        , KnownSTK x, KnownSTK x2+        , BaseTensor target )+     => (target (TKS2 (Drop m sh) x) -> target (TKS2 (Drop m sh) x2))+          -- ^ the function to map with+     -> target (TKS2 sh x)  -- ^ the tensor to map over+     -> target (TKS2 sh x2)+smap @m @sh f v = gcastWith (unsafeCoerceRefl+                             :: sh :~: Take m sh ++ Drop m sh)+                  $ sbuild (\ix -> f (v !$ ix))+smap1 :: (KnownNat n, KnownShS sh, KnownSTK x, KnownSTK x2, BaseTensor target)+      => (target (TKS2 sh x) -> target (TKS2 sh x2))+           -- ^ the function to map with+      -> target (TKS2 (n ': sh) x)  -- ^ the tensor to map over+      -> target (TKS2 (n ': sh) x2)+smap1 f u = sbuild1 (\i -> f (u !$ (i :.$ ZIS)))+smap0N :: (KnownShS sh, KnownSTK x1, KnownSTK x, BaseTensor target)+       => (target (TKS2 '[] x1) -> target (TKS2 '[] x))+            -- ^ the function to map with+       -> target (TKS2 sh x1)  -- ^ the tensor to map over+       -> target (TKS2 sh x)+smap0N = tsmap0N+szipWith :: ( KnownShS (Drop m sh1), KnownShS (Drop m sh2), KnownShS (Take m sh)+            , KnownSTK x, KnownSTK x1, KnownSTK x2, KnownShS sh+            , sh1 ~ Take m sh ++ Drop m sh1+            , sh2 ~ Take m sh ++ Drop m sh2, BaseTensor target )+         => (target (TKS2 (Drop m sh1) x1) -> target (TKS2 (Drop m sh2) x2)+             -> target (TKS2 (Drop m sh) x))  -- ^ the function to zip with+         -> target (TKS2 sh1 x1)  -- ^ the first tensor to zip over+         -> target (TKS2 sh2 x2)  -- ^ the second tensor to zip over+         -> target (TKS2 sh x)+szipWith f u v = sbuild (\ix -> f (u !$ ix) (v !$ ix))+szipWith1 :: ( KnownNat n, KnownShS sh1, KnownShS sh2, KnownShS sh+             , KnownSTK x, KnownSTK x1, KnownSTK x2, BaseTensor target )+          => (target (TKS2 sh1 x1) -> target (TKS2 sh2 x2)+              -> target (TKS2 sh x))  -- ^ the function to zip with+          -> target (TKS2 (n ': sh1) x1)  -- ^ the first tensor to zip over+          -> target (TKS2 (n ': sh2) x2)  -- ^ the second tensor to zip over+          -> target (TKS2 (n ': sh) x)+szipWith1 f u v = sbuild1 (\i -> f (u !$ (i :.$ ZIS))+                                   (v !$ (i :.$ ZIS)))+szipWith0N :: ( KnownShS sh, KnownSTK x, KnownSTK x1, KnownSTK x2+              , BaseTensor target )+           => (target (TKS2 '[] x1) -> target (TKS2 '[] x2)+               -> target (TKS2 '[] x))  -- ^ the function to zip with+           -> target (TKS2 sh x1)  -- ^ the first tensor to zip over+           -> target (TKS2 sh x2)  -- ^ the second tensor to zip over+           -> target (TKS2 sh x)+szipWith0N = tszipWith0N+szipWith3 :: ( KnownShS (Drop m sh1), KnownShS (Drop m sh2)+             , KnownShS (Drop m sh3), KnownShS (Take m sh), KnownShS sh+             , KnownSTK x, KnownSTK x1, KnownSTK x2, KnownSTK x3+             , sh1 ~ Take m sh ++ Drop m sh1+             , sh2 ~ Take m sh ++ Drop m sh2+             , sh3 ~ Take m sh ++ Drop m sh3, BaseTensor target )+          => (target (TKS2 (Drop m sh1) x1) -> target (TKS2 (Drop m sh2) x2)+              -> target (TKS2 (Drop m sh3) x3)+              -> target (TKS2 (Drop m sh) x))  -- ^ the function to zip with+          -> target (TKS2 sh1 x1)  -- ^ the first tensor to zip over+          -> target (TKS2 sh2 x2)  -- ^ the second tensor to zip over+          -> target (TKS2 sh3 x3)  -- ^ the third tensor to zip over+          -> target (TKS2 sh x)+szipWith3 f u v w = sbuild (\ix -> f (u !$ ix) (v !$ ix) (w !$ ix))+szipWith31 :: ( KnownNat n+              , KnownShS sh1, KnownShS sh2, KnownShS sh3, KnownShS sh+              , KnownSTK x, KnownSTK x1, KnownSTK x2, KnownSTK x3+              , BaseTensor target )+           => (target (TKS2 sh1 x1) -> target (TKS2 sh2 x2)+               -> target (TKS2 sh3 x3)+               -> target (TKS2 sh x))  -- ^ the function to zip with+           -> target (TKS2 (n ': sh1) x1)  -- ^ the first tensor to zip over+           -> target (TKS2 (n ': sh2) x2)  -- ^ the second tensor to zip over+           -> target (TKS2 (n ': sh3) x3)  -- ^ the third tensor to zip over+           -> target (TKS2 (n ': sh) x)+szipWith31 f u v w = sbuild1 (\i -> f (u !$ (i :.$ ZIS))+                                      (v !$ (i :.$ ZIS))+                                      (w !$ (i :.$ ZIS)))+szipWith30N :: ( KnownShS sh, KnownSTK x, KnownSTK x1, KnownSTK x2, KnownSTK x3+               , BaseTensor target )+            => (target (TKS2 '[] x1) -> target (TKS2 '[] x2)+                -> target (TKS2 '[] x3)+                -> target (TKS2 '[] x))  -- ^ the function to zip with+            -> target (TKS2 sh x1)  -- ^ the first tensor to zip over+            -> target (TKS2 sh x2)  -- ^ the second tensor to zip over+            -> target (TKS2 sh x3)  -- ^ the third tensor to zip over+            -> target (TKS2 sh x)+szipWith30N @sh f u v w =+  gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])+  $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)+  $ sbuild @(Rank sh) (\ix -> f (sindex0 u ix)+                                (sindex0 v ix)+                                (sindex0 w ix))+szipWith4 :: ( KnownShS (Drop m sh1), KnownShS (Drop m sh2)+             , KnownShS (Drop m sh3), KnownShS (Drop m sh4)+             , KnownShS (Take m sh), KnownShS sh+             , KnownSTK x, KnownSTK x1, KnownSTK x2, KnownSTK x3, KnownSTK x4+             , sh1 ~ Take m sh ++ Drop m sh1+             , sh2 ~ Take m sh ++ Drop m sh2+             , sh3 ~ Take m sh ++ Drop m sh3+             , sh4 ~ Take m sh ++ Drop m sh4, BaseTensor target )+          => (target (TKS2 (Drop m sh1) x1) -> target (TKS2 (Drop m sh2) x2)+              -> target (TKS2 (Drop m sh3) x3) -> target (TKS2 (Drop m sh4) x4)+              -> target (TKS2 (Drop m sh) x))  -- ^ the function to zip with+          -> target (TKS2 sh1 x1)  -- ^ the first tensor to zip over+          -> target (TKS2 sh2 x2)  -- ^ the second tensor to zip over+          -> target (TKS2 sh3 x3)  -- ^ the third tensor to zip over+          -> target (TKS2 sh4 x4)  -- ^ the fourth tensor to zip over+          -> target (TKS2 sh x)+szipWith4 f u v w x =+  sbuild (\ix -> f (u !$ ix) (v !$ ix) (w !$ ix) (x !$ ix))+szipWith41 :: ( KnownNat n+              , KnownShS sh1, KnownShS sh2, KnownShS sh3, KnownShS sh4+              , KnownShS sh+              , KnownSTK x, KnownSTK x1, KnownSTK x2, KnownSTK x3, KnownSTK x4+              , BaseTensor target )+           => (target (TKS2 sh1 x1) -> target (TKS2 sh2 x2)+               -> target (TKS2 sh3 x3) -> target (TKS2 sh4 x4)+               -> target (TKS2 sh x))  -- ^ the function to zip with+           -> target (TKS2 (n ': sh1) x1)  -- ^ the first tensor to zip over+           -> target (TKS2 (n ': sh2) x2)  -- ^ the second tensor to zip over+           -> target (TKS2 (n ': sh3) x3)  -- ^ the third tensor to zip over+           -> target (TKS2 (n ': sh4) x4)  -- ^ the fourth tensor to zip over+           -> target (TKS2 (n ': sh) x)+szipWith41 f u v w x = sbuild1 (\i -> f (u !$ (i :.$ ZIS))+                                        (v !$ (i :.$ ZIS))+                                        (w !$ (i :.$ ZIS))+                                        (x !$ (i :.$ ZIS)))+szipWith40N :: ( KnownShS sh, KnownSTK x+               , KnownSTK x1, KnownSTK x2, KnownSTK x3, KnownSTK x4+               , BaseTensor target )+            => (target (TKS2 '[] x1) -> target (TKS2 '[] x2)+                -> target (TKS2 '[] x3) -> target (TKS2 '[] x4)+                -> target (TKS2 '[] x))  -- ^ the function to zip with+            -> target (TKS2 sh x1)  -- ^ the first tensor to zip over+            -> target (TKS2 sh x2)  -- ^ the second tensor to zip over+            -> target (TKS2 sh x3)  -- ^ the third tensor to zip over+            -> target (TKS2 sh x4)  -- ^ the fourth tensor to zip over+            -> target (TKS2 sh x)+szipWith40N @sh f u v w x =+  gcastWith (unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])+  $ gcastWith (unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)+  $ sbuild @(Rank sh) (\ix -> f (sindex0 u ix)+                                (sindex0 v ix)+                                (sindex0 w ix)+                                (sindex0 x ix))++xbuild1 :: (KnownNat k, KnownShX sh, KnownSTK x, BaseTensor target)+        => (IntOf target -> target (TKX2 sh x))  -- ^ the function to build with+        -> target (TKX2 (Just k ': sh) x)+xbuild1 = txbuild1+-- xmap and other special cases of build can be defined by the user.++rfold+  :: ( KnownNat n, KnownNat m, KnownSTK xn, KnownSTK xm+     , BaseTensor target, LetTensor target )+  => (forall f. ADReady f => f (TKR2 n xn) -> f (TKR2 m xm) -> f (TKR2 n xn))+       -- ^ the function to fold with+  -> target (TKR2 n xn)  -- ^ the initial accumulator+  -> target (TKR2 (1 + m) xm)  -- ^ the inputs+  -> target (TKR2 n xn)+{-# INLINE rfold #-}+rfold f acc0 es =+  withSNat (rwidth es) $ \k -> tfold k knownSTK knownSTK f acc0 es+rscan+  :: ( KnownNat n, KnownNat m, KnownSTK xn, KnownSTK xm+     , BaseTensor target, LetTensor target )+  => (forall f. ADReady f => f (TKR2 n xn) -> f (TKR2 m xm) -> f (TKR2 n xn))+       -- ^ the function to fold with+  -> target (TKR2 n xn)  -- ^ the initial accumulator+  -> target (TKR2 (1 + m) xm)  -- ^ the inputs+  -> target (TKR2 (1 + n) xn)+{-# INLINE rscan #-}+rscan f acc0 es =+  withSNat (rwidth es) $ \k -> tscan k knownSTK knownSTK f acc0 es+sfold+  :: ( KnownNat k, KnownShS sh, KnownShS shm, KnownSTK xn, KnownSTK xm+     , BaseTensor target, LetTensor target )+  => (forall f. ADReady f+      => f (TKS2 sh xn) -> f (TKS2 shm xm) -> f (TKS2 sh xn))+       -- ^ the function to fold with+  -> target (TKS2 sh xn)  -- ^ the initial accumulator+  -> target (TKS2 (k ': shm) xm)  -- ^ the inputs+  -> target (TKS2 sh xn)+{-# INLINE sfold #-}+sfold @k = tfold (SNat @k) knownSTK knownSTK+sscan+  :: ( KnownNat k, KnownShS sh, KnownShS shm, KnownSTK xn, KnownSTK xm+     , BaseTensor target, LetTensor target )+  => (forall f. ADReady f+      => f (TKS2 sh xn) -> f (TKS2 shm xm) -> f (TKS2 sh xn))+       -- ^ the function to scan with+  -> target (TKS2 sh xn)  -- ^ the initial accumulator+  -> target (TKS2 (k ': shm) xm)  -- ^ the inputs+  -> target (TKS2 (1 + k ': sh) xn)+{-# INLINE sscan #-}+sscan @k = tscan (SNat @k) knownSTK knownSTK+xfold+  :: ( KnownNat k, KnownShX sh, KnownShX shm, KnownSTK xn, KnownSTK xm+     , BaseTensor target, LetTensor target )+  => (forall f. ADReady f+      => f (TKX2 sh xn) -> f (TKX2 shm xm) -> f (TKX2 sh xn))+       -- ^ the function to scan with+  -> target (TKX2 sh xn)  -- ^ the initial accumulator+  -> target (BuildTensorKind k (TKX2 shm xm))  -- ^ the inputs+  -> target (TKX2 sh xn)+{-# INLINE xfold #-}+xfold @k = tfold (SNat @k) knownSTK knownSTK+xscan+  :: ( KnownNat k, KnownShX sh, KnownShX shm, KnownSTK xn, KnownSTK xm+     , BaseTensor target, LetTensor target )+  => (forall f. ADReady f+      => f (TKX2 sh xn) -> f (TKX2 shm xm) -> f (TKX2 sh xn))+       -- ^ the function to scan with+  -> target (TKX2 sh xn)  -- ^ the initial accumulator+  -> target (BuildTensorKind k (TKX2 shm xm))  -- ^ the inputs+  -> target (BuildTensorKind (1 + k) (TKX2 sh xn))+{-# INLINE xscan #-}+xscan @k = tscan (SNat @k) knownSTK knownSTK++-- | Reverse derivative.+--+-- The function argument needs to be quantified,+-- because otherwise in the ADVal instance one could put an illegal+-- @DeltaInput@ there, confusing different levels of contangents.+kgrad :: forall x r target. BaseTensor target+      => (forall f. ADReady f => f x -> f (TKScalar r))  -- ^ x |-> TKScalar r+      -> FullShapeTK x  -- ^ shape of x and dx+      -> target x  -- ^ input x+      -> target (ADTensorKind x)  -- ^ gradient dx+kgrad f xftk =+  \ !es -> tApply (tgrad @target xftk (HFun f)) es+rvjp :: forall n x r target. BaseTensor target+     => (forall f. ADReady f => f x -> f (TKR2 n r))  -- ^ x |-> z+     -> FullShapeTK x  -- ^ shape of x and dx+     -> target x  -- ^ input x+     -> target (ADTensorKind (TKR2 n r))  -- ^ incoming cotangent dz+     -> target (ADTensorKind x)  -- ^ gradient dx+rvjp f xftk =+  \ !es !dt -> tApply (tvjp @target xftk $ HFun f) (tpair dt es)+rjvp :: forall n x r target. BaseTensor target+     => (forall f. ADReady f => f x -> f (TKR2 n r))  -- ^ x |-> z+     -> FullShapeTK x  -- ^ shape of x and dx+     -> target x  -- ^ input x+     -> target (ADTensorKind x)  -- ^ incoming tangent dx+     -> target (ADTensorKind (TKR2 n r))  -- ^ derivative dz+rjvp f xftk =+  \ !es !ds -> tApply (tjvp @target xftk $ HFun f) (tpair ds es)+svjp :: forall sh x r target. BaseTensor target+     => (forall f. ADReady f => f x -> f (TKS2 sh r))  -- ^ x |-> z+     -> FullShapeTK x  -- ^ shape of x and dx+     -> target x  -- ^ input x+     -> target (ADTensorKind (TKS2 sh r))  -- ^ incoming cotangent dz+     -> target (ADTensorKind x)  -- ^ gradient dx+svjp f xftk =+  \ !es !dt -> tApply (tvjp @target xftk $ HFun f) (tpair dt es)+sjvp :: forall sh x r target. BaseTensor target+     => (forall f. ADReady f => f x -> f (TKS2 sh r))+     -> FullShapeTK x  -- ^ shape of x and dx+     -> target x  -- ^ input x+     -> target (ADTensorKind x)  -- ^ incoming tangent dx+     -> target (ADTensorKind (TKS2 sh r))  -- ^ derivative dz+sjvp f xftk =+  \ !es !ds -> tApply (tjvp @target xftk $ HFun f) (tpair ds es)++-- These take @target@ first, because they change the target.+rprimalPart :: BaseTensor target+            => target (TKR2 n x) -> PrimalOf target (TKR2 n x)+rprimalPart = tprimalPart+rdualPart :: (BaseTensor target, KnownNat n, KnownSTK x)+          => target (TKR2 n x) -> DualOf target (TKR2 n x)+rdualPart = tdualPart knownSTK+rfromPrimal :: (BaseTensor target, KnownNat n, KnownSTK x)+            => PrimalOf target (TKR2 n x) -> target (TKR2 n x)+rfromPrimal = tfromPrimal knownSTK+rfromDual :: BaseTensor target+          => DualOf target (TKR2 n x) -> target (TKR2 n x)+rfromDual = tfromDual+rScale :: ( BaseTensor target, KnownNat n, GoodScalar r+          , Num (target (TKR n r)), Num (PrimalOf target (TKR n r)) )+       => PrimalOf target (TKR n r) -> DualOf target (TKR n r)+       -> DualOf target (TKR n r)+rScale @target = tScale @target knownSTK++sprimalPart :: BaseTensor target+            => target (TKS2 sh x) -> PrimalOf target (TKS2 sh x)+sprimalPart = tprimalPart+sdualPart :: (BaseTensor target, KnownShS sh, KnownSTK x)+          => target (TKS2 sh x) -> DualOf target (TKS2 sh x)+sdualPart = tdualPart knownSTK+sfromPrimal :: (BaseTensor target, KnownShS sh, KnownSTK x)+            => PrimalOf target (TKS2 sh x) -> target (TKS2 sh x)+sfromPrimal = tfromPrimal knownSTK+sfromDual :: BaseTensor target+          => DualOf target (TKS2 sh x) -> target (TKS2 sh x)+sfromDual = tfromDual+sScale :: ( BaseTensor target, KnownShS sh, GoodScalar r+          , Num (target (TKS sh r)), Num (PrimalOf target (TKS sh r)) )+       => PrimalOf target (TKS sh r) -> DualOf target (TKS sh r)+       -> DualOf target (TKS sh r)+sScale @target = tScale @target knownSTK++xprimalPart :: BaseTensor target+            => target (TKX2 sh x) -> PrimalOf target (TKX2 sh x)+xprimalPart = tprimalPart+xdualPart :: (BaseTensor target, KnownShX sh, KnownSTK x)+          => target (TKX2 sh x) -> DualOf target (TKX2 sh x)+xdualPart = tdualPart knownSTK+xfromPrimal :: (BaseTensor target, KnownShX sh, KnownSTK x)+            => PrimalOf target (TKX2 sh x) -> target (TKX2 sh x)+xfromPrimal = tfromPrimal knownSTK+xfromDual :: BaseTensor target+          => DualOf target (TKX2 sh x) -> target (TKX2 sh x)+xfromDual = tfromDual+xScale :: ( BaseTensor target, KnownShX sh, GoodScalar r+          , Num (target (TKX sh r)), Num (PrimalOf target (TKX sh r)) )+       => PrimalOf target (TKX sh r) -> DualOf target (TKX sh r)+       -> DualOf target (TKX sh r)+xScale @target = tScale @target knownSTK++kprimalPart :: BaseTensor target+            => target (TKScalar r) -> PrimalOf target (TKScalar r)+kprimalPart = tprimalPart+kdualPart :: (BaseTensor target, GoodScalar r)+          => target (TKScalar r) -> DualOf target (TKScalar r)+kdualPart = tdualPart knownSTK+kfromPrimal :: (BaseTensor target, GoodScalar r)+            => PrimalOf target (TKScalar r) -> target (TKScalar r)+kfromPrimal = tfromPrimal knownSTK+kfromDual :: BaseTensor target+          => DualOf target (TKScalar r) -> target (TKScalar r)+kfromDual = tfromDual+kScale :: ( BaseTensor target, GoodScalar r+          , Num (target (TKScalar r)), Num (PrimalOf target (TKScalar r)) )+       => PrimalOf target (TKScalar r) -> DualOf target (TKScalar r)+       -> DualOf target (TKScalar r)+kScale @target = tScale @target knownSTK
+ src/HordeAd/OpsTensorMixed.hs view
@@ -0,0 +1,52 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | The mixed tensor operations intended for the casual library user.+-- All these operations together with operations of the remaining two tensor+-- variants are gathered in "HordeAd.OpsTensor".+--+-- The less user-friendly+-- prototypes of most of these operation can be found in "HordeAd.Core.Ops"+-- where some additional rarely used operations reside.+-- All these operations, together with instances of numerical classes+-- such as @Num@, @Fractional@, @IntegralH@, @RealFloatH@, @EqH@ and others+-- (see class instances of type 'HordeAd.Core.Ast.AstTensor' for the full list),+-- are a major part of the high-level API of the horde-ad library,+-- which is relatively orthogonal to the other major part,+-- the differentiation interface exposed in "HordeAd.ADEngine".+module HordeAd.OpsTensorMixed+  ( -- * Shape manipulation+    xshape, xlength, xsize, xwidth+  , tsize, tftk+    -- * Constructing arrays from concrete values, lists and vectors+  , xconcrete, xscalar, xrepl, xingestData, xfromListLinear+  , kconcrete+  , xfromList, xfromVector, xfromVector0N, xfromList0N, xunravelToList+    -- * Main array operations+  , tunit, tlet, ifH, minH, maxH, tpair, tproject1, tproject2+  , xsum, xsum0, xdot0, xdot1In, xmatvecmul, xmatmul2, xreplicate, xreplicate0N+  , xindex, xindex0, xoneHot, xscatter, xscatter1, xgather, xgather1+  , xtr, xtranspose, xflatten, xreshape+   -- * Auxiliary array operations+  , xfloor, xfromIntegral, xcast, xminIndex, xmaxIndex, xiota+  , kfloor, kfromIntegral, kcast+  , xappend, xappend0, xconcat, xslice, xuncons, xreverse+    -- * Array operations derived from @build@+  , xbuild, xbuild1+    -- * Array operations derived from @mapAccum@+  , xfold, xscan, tmapAccumR, tmapAccumL+    -- * Array operations producing derivatives+  , kgrad+    -- * Operations dealing with dual numbers+  , xprimalPart, xdualPart, xfromPrimal, xfromDual, xScale+  , kprimalPart, kdualPart, kfromPrimal, kfromDual, kScale+    -- * Array operations that utilize unwinding of nested arrays+  , treplTarget, tdefTarget, taddTarget, tmultTarget, tsum0Target, tdot0Target+    -- * Minimal re-exports to make this module a higher level replacement for the mixed part of "HordeAd.Core.Ops"+  , ADReady+  , LetTensor, BaseTensor+  ) where++import Prelude ()++import HordeAd.OpsTensor
+ src/HordeAd/OpsTensorRanked.hs view
@@ -0,0 +1,53 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | The ranked tensor operations intended for the casual library user.+-- All these operations together with operations of the remaining two tensor+-- variants are gathered in "HordeAd.OpsTensor".+--+-- The less user-friendly+-- prototypes of most of these operation can be found in "HordeAd.Core.Ops"+-- where some additional rarely used operations reside.+-- All these operations, together with instances of numerical classes+-- such as @Num@, @Fractional@, @IntegralH@, @RealFloatH@, @EqH@ and others+-- (see class instances of type 'HordeAd.Core.Ast.AstTensor' for the full list),+-- are a major part of the high-level API of the horde-ad library,+-- which is relatively orthogonal to the other major part,+-- the differentiation interface exposed in "HordeAd.ADEngine".+module HordeAd.OpsTensorRanked+  ( -- * Shape manipulation+    rshape, rlength, rsize, rwidth+  , tsize, tftk+    -- * Constructing arrays from concrete values, lists and vectors+  , rconcrete, rscalar, rrepl, ringestData, rfromListLinear+  , kconcrete+  , rfromList, rfromVector, rfromVector0N, rfromList0N, runravelToList+    -- * Main array operations+  , tunit, tlet, ifH, minH, maxH, tpair, tproject1, tproject2+  , rsum, rsum0, rdot0, rdot1In, rmatvecmul, rmatmul2, rreplicate, rreplicate0N+  , rindex, (!), rindex0, roneHot, rscatter, rscatter1, rgather, rgather1+  , rtr, rtranspose, rflatten, rreshape+   -- * Auxiliary array operations+  , rfloor, rfromIntegral, rcast, rminIndex, rmaxIndex, riota+  , kfloor, kfromIntegral, kcast+  , rappend, rconcat, rslice, runcons, rreverse+    -- * Array operations derived from @build@+  , rbuild, rbuild1, rmap, rmap1, rmap0N, rzipWith, rzipWith1, rzipWith0N+  , rzipWith3, rzipWith31, rzipWith30N, rzipWith4, rzipWith41, rzipWith40N+    -- * Array operations derived from @mapAccum@+  , rfold, rscan, tmapAccumR, tmapAccumL+    -- * Array operations producing derivatives+  , kgrad, rvjp, rjvp+    -- * Operations dealing with dual numbers+  , rprimalPart, rdualPart, rfromPrimal, rfromDual, rScale+  , kprimalPart, kdualPart, kfromPrimal, kfromDual, kScale+    -- * Array operations that utilize unwinding of nested arrays+  , treplTarget, tdefTarget, taddTarget, tmultTarget, tsum0Target, tdot0Target+    -- * Minimal re-exports to make this module a higher level replacement for the ranked part of "HordeAd.Core.Ops"+  , ADReady+  , LetTensor, BaseTensor+  ) where++import Prelude ()++import HordeAd.OpsTensor
+ src/HordeAd/OpsTensorShaped.hs view
@@ -0,0 +1,53 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | The shaped tensor operations intended for the casual library user.+-- All these operations together with operations of the remaining two tensor+-- variants are gathered in "HordeAd.OpsTensor".+--+-- The less user-friendly+-- prototypes of most of these operation can be found in "HordeAd.Core.Ops"+-- where some additional rarely used operations reside.+-- All these operations, together with instances of numerical classes+-- such as @Num@, @Fractional@, @IntegralH@, @RealFloatH@, @EqH@ and others+-- (see class instances of type 'HordeAd.Core.Ast.AstTensor' for the full list),+-- are a major part of the high-level API of the horde-ad library,+-- which is relatively orthogonal to the other major part,+-- the differentiation interface exposed in "HordeAd.ADEngine".+module HordeAd.OpsTensorShaped+  ( -- * Shape manipulation+    sshape, slength, ssize, swidth+  , tsize, tftk+    -- * Constructing arrays from concrete values, lists and vectors+  , sconcrete, sscalar, srepl, singestData, sfromListLinear+  , kconcrete+  , sfromList, sfromVector, sfromVector0N, sfromList0N, sunravelToList+    -- * Main array operations+  , tunit, tlet, ifH, minH, maxH, tpair, tproject1, tproject2+  , ssum, ssum0, sdot0, sdot1In, smatvecmul, smatmul2, sreplicate, sreplicate0N+  , sindex, (!$), sindex0, soneHot, sscatter, sscatter1, sgather, sgather1+  , str, stranspose, sflatten, sreshape+   -- * Auxiliary array operations+  , sfloor, sfromIntegral, scast, sminIndex, smaxIndex, siota+  , kfloor, kfromIntegral, kcast+  , sappend, sslice, suncons, sreverse+    -- * Array operations derived from @build@+  , sbuild, sbuild1, smap, smap1, smap0N, szipWith, szipWith1, szipWith0N+  , szipWith3, szipWith31, szipWith30N, szipWith4, szipWith41, szipWith40N+    -- * Array operations derived from @mapAccum@+  , sfold, sscan, tmapAccumR, tmapAccumL+    -- * Array operations producing derivatives+  , kgrad, svjp, sjvp+    -- * Operations dealing with dual numbers+  , sprimalPart, sdualPart, sfromPrimal, sfromDual, sScale+  , kprimalPart, kdualPart, kfromPrimal, kfromDual, kScale+    -- * Array operations that utilize unwinding of nested arrays+  , treplTarget, tdefTarget, taddTarget, tmultTarget, tsum0Target, tdot0Target+    -- * Minimal re-exports to make this module a higher level replacement for the shaped part of "HordeAd.Core.Ops"+  , ADReady+  , LetTensor, BaseTensor+  ) where++import Prelude ()++import HordeAd.OpsTensor
+ test/CAFlessTest.hs view
@@ -0,0 +1,40 @@+module Main (main) where++import Prelude++import Data.Proxy+import System.IO qualified as SIO+import Test.Tasty+import Test.Tasty.Options+import Test.Tasty.Runners++import EqEpsilon+import TestAdaptorSimplified qualified+import TestConvSimplified qualified+import TestGatherSimplified qualified+import TestHighRankSimplified qualified+import TestMnistPP qualified+import TestRevFwdFold qualified++main :: IO ()+main = do+  -- Limit interleaving of characters in parallel tests.+  SIO.hSetBuffering SIO.stdout SIO.LineBuffering+  SIO.hSetBuffering SIO.stderr SIO.LineBuffering+  opts <- parseOptions (ingredients : defaultIngredients) tests+  setEpsilonEq (lookupOption opts :: EqEpsilon)+  defaultMainWithIngredients (ingredients : defaultIngredients) tests+ where+  ingredients = includingOptions [Option (Proxy :: Proxy EqEpsilon)]++tests :: TestTree+tests =+  testGroup "Tests for simplified horde-ad that don't create big CAFs"+    [ testGroup "Short_tests"+        (TestAdaptorSimplified.testTrees+         ++ TestConvSimplified.testTrees+         ++ TestGatherSimplified.testTrees+         ++ TestHighRankSimplified.testTrees+         ++ TestRevFwdFold.testTrees+         ++ TestMnistPP.testTrees)+    ]
+ test/FullTest.hs view
@@ -0,0 +1,51 @@+module Main (main) where++import Prelude++import Data.Proxy+import System.IO qualified as SIO+import Test.Tasty+import Test.Tasty.Options+import Test.Tasty.Runners++import EqEpsilon+import TestAdaptorSimplified qualified+import TestConvSimplified qualified+import TestGatherSimplified qualified+import TestHighRankSimplified qualified+import TestMnistCNNR qualified+import TestMnistCNNS qualified+import TestMnistFCNNR qualified+import TestMnistPP qualified+import TestMnistRNNR qualified+import TestMnistRNNS qualified+import TestRevFwdFold qualified++main :: IO ()+main = do+  -- Limit interleaving of characters in parallel tests.+  SIO.hSetBuffering SIO.stdout SIO.LineBuffering+  SIO.hSetBuffering SIO.stderr SIO.LineBuffering+  opts <- parseOptions (ingredients : defaultIngredients) tests+  setEpsilonEq (lookupOption opts :: EqEpsilon)+  defaultMainWithIngredients (ingredients : defaultIngredients) tests+ where+  ingredients = includingOptions [Option (Proxy :: Proxy EqEpsilon)]++tests :: TestTree+tests =+  testGroup "Full set of tests for simplified horde-ad"+    [ testGroup "Short_tests"+        (TestAdaptorSimplified.testTrees+         ++ TestConvSimplified.testTrees+         ++ TestGatherSimplified.testTrees+         ++ TestHighRankSimplified.testTrees+         ++ TestRevFwdFold.testTrees+         ++ TestMnistPP.testTrees)+    , testGroup "Long_tests"+        (TestMnistCNNR.testTrees+         ++ TestMnistCNNS.testTrees+         ++ TestMnistFCNNR.testTrees+         ++ TestMnistRNNR.testTrees+         ++ TestMnistRNNS.testTrees)+    ]
+ test/MinimalTest.hs view
@@ -0,0 +1,28 @@+module Main (main) where++import Prelude++import Data.Proxy+import System.IO qualified as SIO+import Test.Tasty+import Test.Tasty.Options+import Test.Tasty.Runners++import EqEpsilon+import TestGatherSimplified qualified++main :: IO ()+main = do+  -- Limit interleaving of characters in parallel tests.+  SIO.hSetBuffering SIO.stdout SIO.LineBuffering+  SIO.hSetBuffering SIO.stderr SIO.LineBuffering+  opts <- parseOptions (ingredients : defaultIngredients) tests+  setEpsilonEq (lookupOption opts :: EqEpsilon)+  defaultMainWithIngredients (ingredients : defaultIngredients) tests+ where+  ingredients = includingOptions [Option (Proxy :: Proxy EqEpsilon)]++tests :: TestTree+tests =+  testGroup "Minimal test that doesn't require any dataset"+    TestGatherSimplified.testTrees
+ test/ParallelTest.hs view
@@ -0,0 +1,38 @@+module Main (main) where++import Prelude++import Data.Proxy+import System.IO qualified as SIO+import Test.Tasty+import Test.Tasty.Options+import Test.Tasty.Runners++import EqEpsilon+import TestMnistCNNR qualified+import TestMnistCNNS qualified+import TestMnistFCNNR qualified+import TestMnistRNNR qualified+import TestMnistRNNS qualified++main :: IO ()+main = do+  -- Limit interleaving of characters in parallel tests.+  SIO.hSetBuffering SIO.stdout SIO.LineBuffering+  SIO.hSetBuffering SIO.stderr SIO.LineBuffering+  opts <- parseOptions (ingredients : defaultIngredients) tests+  setEpsilonEq (lookupOption opts :: EqEpsilon)+  defaultMainWithIngredients (ingredients : defaultIngredients) tests+ where+  ingredients = includingOptions [Option (Proxy :: Proxy EqEpsilon)]++tests :: TestTree+tests =+  testGroup "The set of tests for simplified horde-ad that can run in parallel"+    [ testGroup "Long_tests"+        (TestMnistCNNR.testTrees+         ++ TestMnistCNNS.testTrees+         ++ TestMnistFCNNR.testTrees+         ++ TestMnistRNNR.testTrees+         ++ TestMnistRNNS.testTrees)+    ]
+ test/simplified/TestAdaptorSimplified.hs view
@@ -0,0 +1,2458 @@+{-# LANGUAGE OverloadedLists #-}+{-# OPTIONS_GHC -fno-cse #-}+-- | Assorted rather low rank tensor tests.+module TestAdaptorSimplified+  ( testTrees+  ) where++import Prelude++import Data.Int (Int64)+import Data.Kind (Type)+import Data.List.NonEmpty qualified as NonEmpty+import Foreign.C (CInt)+import GHC.Exts (IsList (..))+import GHC.TypeLits (KnownNat)+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId (funToAst, resetVarCounter)+import HordeAd.Core.AstInterpret+import HordeAd.Core.DeltaFreshId (resetIdCounter)++import CrossTesting+import EqEpsilon++testTrees :: [TestTree]+testTrees =+  [ testCase "2zeroZ" testZeroZ+  , testCase "2zeroS" testZeroS+  , testCase "2CFwdZeroS" testCFwdZeroS+  , testCase "2FwdZeroS" testFwdZeroS+  , testCase "2zero2S" testZero2S+  , testCase "2CFwdZero2S" testCFwdZero2S+  , testCase "2FwdZero2S" testFwdZero2S+  , testCase "2zero3S" testZero3S+  , testCase "2CFwdZero3S" testCFwdZero3S+  , testCase "2FwdZero3S" testFwdZero3S+  , testCase "2zero4s" testZero4s+  , testCase "2zero5S" testZero5S+  , testCase "2zero6S" testZero6S+  , testCase "2zero7S" testZero7S+  , testCase "2zero8" testZero8+  , testCase "2zero9S" testZero9S+  , testCase "2CFwdZero9S" testCFwdZero9S+  , testCase "2FwdZero9S" testFwdZero9S+  , testCase "2zero10S" testZero10S+  , testCase "2CFwdZero10S" testCFwdZero10S+  , testCase "2FwdZero10S" testFwdZero10S+  , testCase "2zero11S" testZero11S+  , testCase "2CFwdZero11S" testCFwdZero11S+  , testCase "2FwdZero11S" testFwdZero11S+  , testCase "2piecewiseLinear" testPiecewiseLinear+  , testCase "2piecewiseLinearPP" testPiecewiseLinearPP+  , testCase "2piecewiseLinear2" testPiecewiseLinear2+  , testCase "2piecewiseLinear2PP" testPiecewiseLinear2PP+  , testCase "2overleaf" testOverleaf+  , testCase "2overleafInt64n" testOverleafInt64n+  , testCase "2overleafCIntn" testOverleafCIntn+  , testCase "2overleafCIntToFloatn" testOverleafCIntToFloatn+  , testCase "2overleafInt64p" testOverleafInt64p+  , testCase "2overleafCIntp" testOverleafCIntp+  , testCase "2overleafCIntToFloatp" testOverleafCIntToFloatp+  , testCase "2overleafPP" testOverleafPP+  , testCase "2foo" testFoo+  , testCase "2fooGradDouble" testGradFooDouble+  , testCase "2fooMatrix" testFooMatrix+  , testCase "2fooGradMatrix" testGradFooMatrix+  , testCase "2fooLetGradMatrixPP" testGradFooLetMatrixPP+  , testCase "2fooGradMatrixVjp" testGradFooMatrixVjp+  , testCase "2fooGradMatrixRev" testGradFooMatrixRev+  , testCase "2fooLetGradMatrixSimpPP" testGradFooLetMatrixSimpPP+  , testCase "2fooLetGradMatrixSimpRPP" testGradFooLetMatrixSimpRPP+  , testCase "2fooSumMatrix" testfooSumMatrix+  , testCase "2fooGradMatrix2" testGradFooMatrix2+  , testCase "2fooGradMatrixPP" testGradFooMatrixPP+  , testCase "2fooGradMatrixSimpPP" testGradFooMatrixSimpPP+  , testCase "2fooGradScalar" testGradFooScalar+  , testCase "2fooGradCScalar" testGradCFooScalar+  , testCase "2fooS" testFooS+  , testCase "2fooSToFloat" testFooSToFloat+  , testCase "2fooSBoth" testFooSBoth+  , testCase "2fooBoth" testFooBoth+  , testCase "2trustVstackConcatRepl10" testTrustVstackConcatRepl10+  , testCase "2trustVstackConcatIota10" testTrustVstackConcatIota10+  , testCase "2trustVstackConcatReplIota10" testTrustVstackConcatReplIota10+  , testCase "2vstackWarmup" testVstackWarmup+  , testCase "2vstackConcatConcrete" testVstackConcatConcrete+  , testCase "2vstackBuildConcrete" testVstackBuildConcrete+  , testCase "2vstackConcatAst" testVstackConcatAst+  , testCase "2vstackBuildAst" testVstackBuildAst+  , testCase "2vstackBuildAstSimp" testVstackBuildAstSimp+  , testCase "2vstackBuildAstPP" testVstackBuildAstPP+  , testCase "2vstackConcatConcrete2" testVstackConcatConcrete2+  , testCase "2vstackBuildConcrete2" testVstackBuildConcrete2+  , testCase "2vstackConcatAst2" testVstackConcatAst2+  , testCase "2vstackBuildAst2" testVstackBuildAst2+  , testCase "2vstackBuildAstSimp2" testVstackBuildAstSimp2+  , testCase "2vstackBuildAstPP2" testVstackBuildAstPP2+  , testCase "2fooPP" testFooPP+  , testCase "2fooLet" testFooLet+  , testCase "2fooLetPP" testFooLetPP+  , testCase "2listProdPP" testListProdPP+  , testCase "2listProdrPP" testListProdrPP+  , testCase "2listProdrLongPP" testListProdrLongPP+  , testCase "2listProd" testListProd+  , testCase "2listProdr" testListProdr+  , testCase "2listSumrPP" testListSumrPP+  , testCase "2listSum2rPP" testListSum2rPP+  , testCase "2listSum22rPP" testListSum22rPP+  , testCase "2listSumk22rPP" testListSumk22rPP+  , testCase "2listSum2xpyrPP" testListSum2xpyrPP+  , testCase "2listSum2xyrPP" testListSum2xyrPP+  , testCase "2list2xyPP" test2xyPP+  , testCase "2listSum23rPP" testListSum23rPP+  , testCase "2list23PP" test23PP+  , testCase "2reluPP" testReluPP+  , testCase "2reluPP2" testReluPP2+  , testCase "2reluSimpler" testReluSimpler+  , testCase "2reluSimplerPP" testReluSimplerPP+  , testCase "2reluSimplerPP2" testReluSimplerPP2+  , testCase "2reluSimplerPP3" testReluSimplerPP3+  , testCase "2reluSimpler3" testReluSimpler3+  , testCase "2reluSimplerPP4" testReluSimplerPP4+  , testCase "2reluSimpler4" testReluSimpler4+  , testCase "2reluSimplerPP4s" testReluSimplerPP4s+  , testCase "2reluSimplerPP4s2" testReluSimplerPP4s2+  , testCase "2reluSimpler4s" testReluSimpler4s+  , testCase "2reluSimplerPP7s2" testReluSimplerPP7s2+  , testCase "2reluSimpler7s" testReluSimpler7s+  , testCase "2reluMax" testReluMax+  , testCase "2reluMaxPP" testReluMaxPP+  , testCase "2reluMaxPP2" testReluMaxPP2+  , testCase "2reluMax3" testReluMax3+  , testCase "2dot1PP" testDot1PP+  , testCase "2dot2PP" testDot2PP+  , testCase "2matvecmulPP" testMatvecmulPP+  , testCase "2matmul2PP" testMatmul2PP+  , testCase "2matmul2FromMatvecmulPP" testMatmul2FromMatvecmulPP+  , testCase "2matmul2PaperPP" testMatmul2PaperPP+  , testCase "2matmul2PPS" testMatmul2PPS+  , testCase "2addSpeedBig" testAddSpeedBig+  , testCase "2matmul2SpeedBig" testMatmul2SpeedBig+  , testCase "2bar" testBar+  , testCase "2barS" testBarS+  , testCase "2bar2S" testBar2S+  , testCase "2barCFwd" testBarCFwd+  , testCase "2barFwd" testBarFwd+  , testCase "2baz old to force fooFromPrimal" testBaz+  , testCase "2baz if repetition breaks things" testBaz+  , testCase "2baz again with renumbered terms" testBazRenumbered+  , testCase "2fooD T Double [1.1, 2.2, 3.3]" testFooD+  , testCase "2fooBuildDt" testFooBuildDt+  , testCase "2fooBuildCFwd" testFooBuildCFwd+  , testCase "2fooBuildFwd" testFooBuildFwd+  , testCase "2fooBuild" testFooBuild+  , testCase "2fooNoGo0" testFooNoGo0+  , testCase "2nestedBuildMap1" testNestedBuildMap1+  , testCase "2nestedSumBuild" testNestedSumBuild+  , testCase "2nestedBuildIndex" testNestedBuildIndex+  , testCase "2barReluDt" testBarReluDt+  , testCase "2barRelu" testBarRelu+  , testCase "2barRelu3" testBarRelu3+  , testCase "2barReluMaxDt" testBarReluMaxDt+  , testCase "2barReluMax" testBarReluMax+  , testCase "2barReluMax30" testBarReluMax30+  , testCase "2barReluMax31" testBarReluMax31+  , testCase "2barReluMax3CFwd" testBarReluMax3CFwd+  , testCase "2barReluMax3FwdS" testBarReluMax3FwdS+  , testCase "2barReluMax3FwdFrom" testBarReluMax3FwdFrom+  , testCase "2barReluMax3FwdR" testBarReluMax3FwdR+  , testCase "2F1" testF1+  , testCase "2F11" testF11+  , testCase "2F2" testF2+  , testCase "2F21" testF21+  , testCase "2F2CFwd" testF2CFwd+  , testCase "2F2Fwd" testF2Fwd+  , testCase "2braidedBuilds1" testBraidedBuilds1+  , testCase "2recycled1" testRecycled1+  , testCase "2concatBuild1" testConcatBuild1+  , testCase "2concatBuild2" testConcatBuild2+  , testCase "2concatBuild3" testConcatBuild3+  , testCase "2concatBuild4" testConcatBuild4+  , testCase "2concatBuild5" testConcatBuild5+  , testCase "2concatBuild6" testConcatBuild6+  , testCase "2concatBuild7" testConcatBuild7+--  , testCase "2concatBuild8" testConcatBuild8+--  , testCase "2concatBuild9" testConcatBuild9+  , testCase "2concatBuild10" testConcatBuild10+  , testCase "2concatBuild11" testConcatBuild11+  , testCase "2concatBuild12" testConcatBuild12+  , testCase "2emptyArgs0" testEmptyArgs0+  , testCase "2emptyArgs1" testEmptyArgs1+  , testCase "2emptyArgs4" testEmptyArgs4+  , testCase "2filterPositiveFail" testFilterPositiveFail+  , testCase "2blowupPP" fblowupPP+  , testCase "2blowup2LetPP" fblowupLetPP+  , testCase "2blowup2LetPP23" fblowupLetPP23+  , testCase "2blowup2LetPP10" fblowupLetPP10+  , blowupTests+  , testCase "22concatBuild3PP" testConcatBuild3PP+  , testCase "22concatBuild3PP2" testConcatBuild3PP2+  ]++testZeroZ :: Assertion+testZeroZ =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @0 (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+                          f = const (rscalar 3)+                      in f) (rscalar 42))++testZeroS :: Assertion+testZeroS =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])+    (cgrad (kfromS . ssum0 .+            let f :: ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+                  -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+                f = const (srepl 3)+            in f) (srepl 42))++testCFwdZeroS :: Assertion+testCFwdZeroS =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])+    (cjvp (let f :: ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+                 -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+               f = const (srepl 3)+           in f) 42 41)++testFwdZeroS :: Assertion+testFwdZeroS =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])+    (jvp (let f :: AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double)+                -> AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double)+              f = const (srepl 3)+          in f) (srepl 42) (srepl 41))++testZero2S :: Assertion+testZero2S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [1])+    (cgrad (kfromS @_ @Double .+           let f :: a -> a+               f = id+           in f) (srepl 42))++testCFwdZero2S :: Assertion+testCFwdZero2S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [41])+    (cjvp @_ @(TKS '[] Double)+          (let f :: a -> a+               f = id+           in f) (srepl 42) (srepl 41))++testFwdZero2S :: Assertion+testFwdZero2S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [41])+    (jvp @_ @(TKS '[] Double)+          (let f :: a -> a+               f = id+           in f) (srepl 42) (srepl 41))++testZero3S :: Assertion+testZero3S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[33, 2] knownShS (replicate 66 3.6174114266850617))+    (cgrad (kfromS . ssum0 . (\x -> barF @(ADVal Concrete (TKS '[33, 2] Double)) (x, x))) (srepl 1))++testCFwdZero3S :: Assertion+testCFwdZero3S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[33, 2] knownShS (replicate 66 3.9791525693535674))+    (cjvp (\x -> barF @(ADVal Concrete (TKS '[33, 2] Double)) (x, x)) (srepl 1) (srepl 1.1))++testFwdZero3S :: Assertion+testFwdZero3S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[33, 2] knownShS (replicate 66 3.9791525693535674))+    (jvp (\x -> barF @(AstTensor AstMethodLet FullSpan (TKS '[33, 2] Double)) (x, x)) (srepl 1) (srepl 1.1))++testZero4s :: Assertion+testZero4s =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [0])+    (grad @(AstTensor AstMethodLet FullSpan (TKS '[] Double))+          (kfromS @_ @Double .+          let f = const (srepl 3)+          in f) (srepl 42))++testZero5S :: Assertion+testZero5S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[44] knownShS (replicate 44 1))+    (grad (kfromS . ssum0 .+           let f :: a -> a+               f = id+           in f @(AstTensor AstMethodLet FullSpan (TKS '[44] Double))) (srepl 42))++testZero6S :: Assertion+testZero6S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,11,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,111,1,1,1,1, 2, 2, 2, 2] knownShS (replicate (product ([2, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,11,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,111,1,1,1,1, 2, 2, 2, 2] :: [Int])) 3.6174114266850617))+    (grad (kfromS . ssum0 @'[2, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 2, 2, 2, 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,11,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,111,1,1,1,1, 2, 2, 2, 2] @(TKScalar Double) . (\x -> barF (x, x))) (srepl 1))++testZero7S :: Assertion+testZero7S =+  assertEqualUpToEpsilon 1e-10+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[] knownShS [0])+    (grad (kfromR . (const (rscalar 3) :: AstTensor AstMethodLet FullSpan (TKS '[] Double) -> AstTensor AstMethodLet FullSpan (TKR 0 Double))) (srepl 42))++testZero8 :: Assertion+testZero8 =+  assertEqualUpToEpsilon 1e-10+    (rfromList0N [] [rscalar 0])+    (grad (kfromS . (const (sscalar 3) :: AstTensor AstMethodLet FullSpan (TKR 0 Double) -> AstTensor AstMethodLet FullSpan (TKS '[] Double))) (rscalar 42))++testZero9S :: Assertion+testZero9S =+  assertEqualUpToEpsilon 1e-9+    (rfromList0N [0, 2, 4, 0, 1] [])+    (cgrad (kfromS . ssum0 .+            let f :: ADVal Concrete (TKR 5 Double)+                  -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+                f = const (srepl 3)+            in f)+           (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42)))++testCFwdZero9S :: Assertion+testCFwdZero9S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])+    (cjvp (let f :: ADVal Concrete (TKR 5 Double)+                 -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double)+               f = const (srepl 3)+           in f)+          (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42)) (rreplicate0N [0, 2, 4, 0, 1] (rscalar 41)))++testFwdZero9S :: Assertion+testFwdZero9S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])+    (jvp (let f :: AstTensor AstMethodLet FullSpan (TKR 5 Double)+                -> AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double)+              f = const (srepl 3)+          in f)+         (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42)) (rreplicate0N [0, 2, 4, 0, 1] (rscalar 41)))++testZero10S :: Assertion+testZero10S =+  assertEqualUpToEpsilon 1e-9+    ( rfromList0N [0, 2, 4, 0, 1] []+    , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )+    (cgrad (kfromS . ssum0 .+            let f = const (srepl 3) . snd+            in f :: ( ADVal Concrete (TKR 5 Double)+                    , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )+                 -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double))+          (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42), srepl 21))++testCFwdZero10S :: Assertion+testCFwdZero10S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])+    (cjvp (let f = const (srepl 3) . snd+           in f :: ( ADVal Concrete (TKR 5 Double)+                   , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )+                   -> ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double))+          ( rfromList0N [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )+          ( rfromList0N [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] ))++testFwdZero10S :: Assertion+testFwdZero10S =+  assertEqualUpToEpsilon 1e-9+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [])+    (jvp  (let f = const (srepl 3) . snd+           in f :: ( AstTensor AstMethodLet FullSpan (TKR 5 Double)+                   , AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double) )+                   -> AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double))+          ( rfromList0N [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )+          ( rfromList0N [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] ))++testZero11S :: Assertion+testZero11S =+  assertEqualUpToEpsilon 1e-9+    ( rfromList0N [0, 2, 4, 0, 1] []+    , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )+    (cgrad (kfromR . rsum0 .+            let f = const (rreplicate0N [0, 2, 4, 0, 1] (rscalar 3)) . snd+            in f :: ( ADVal Concrete (TKR 5 Double)+                    , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )+                 -> ADVal Concrete (TKR 5 Double))+          (rreplicate0N [0, 2, 4, 0, 1] (rscalar 42), srepl 21))++testCFwdZero11S :: Assertion+testCFwdZero11S =+  assertEqualUpToEpsilon 1e-9+    (rfromList0N [0, 2, 4, 0, 1] [])+    (cjvp (let f = const (rreplicate0N [0, 2, 4, 0, 1] (rscalar 3)) . snd+           in f :: ( ADVal Concrete (TKR 5 Double)+                   , ADVal Concrete (TKS '[0, 2, 4, 0, 1] Double) )+                   -> ADVal Concrete (TKR 5 Double))+          ( rfromList0N [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )+          ( rfromList0N [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] ))++testFwdZero11S :: Assertion+testFwdZero11S =+  assertEqualUpToEpsilon 1e-9+    (rfromList0N [0, 2, 4, 0, 1] [])+    (jvp  (let f = const (rreplicate0N [0, 2, 4, 0, 1] (rscalar 3)) . snd+           in f :: ( AstTensor AstMethodLet FullSpan (TKR 5 Double)+                   , AstTensor AstMethodLet FullSpan (TKS '[0, 2, 4, 0, 1] Double) )+                   -> AstTensor AstMethodLet FullSpan (TKR 5 Double))+          ( rfromList0N [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] )+          ( rfromList0N [0, 2, 4, 0, 1] []+          , sconcrete $ Nested.sfromListPrimLinear @_ @'[0, 2, 4, 0, 1] knownShS [] ))++testPiecewiseLinear :: Assertion+testPiecewiseLinear =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2)+    (let fT :: ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+         fT x = ifH (x >. rscalar 0) (rscalar 2 * x) (rscalar 5 * x)+     in rev' @Double @0 fT (rscalar 42))++testPiecewiseLinearPP :: Assertion+testPiecewiseLinearPP = do+  resetVarCounter+  let fT :: AstTensor AstMethodLet FullSpan (TKR 0 Double)+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT x = ifH (x >. rscalar 0) (rscalar 2 * x) (rscalar 5 * x)+      (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKR ZSR FTKScalar)+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> rfromS (let v4 = soneHot (sfromR dret) [ifH (sscalar -0.0 <=. negate (sfromR x1)) 0 1] in sscalar 2.0 * v4 !$ [1] + sscalar 5.0 * v4 !$ [0])"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (ifH (sscalar -0.0 <=. negate (sfromR x1)) (sscalar 5.0 * sfromR x1) (sscalar 2.0 * sfromR x1))"++testPiecewiseLinear2 :: Assertion+testPiecewiseLinear2 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 5)+    (let fT :: ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+         fT x = ifH (x >. rscalar 0) (rscalar 2) (rscalar 5) * x+     in rev' @Double @0 fT (rscalar (-42)))++testPiecewiseLinear2PP :: Assertion+testPiecewiseLinear2PP = do+  resetVarCounter+  let fT :: AstTensor AstMethodLet FullSpan (TKR 0 Double)+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT x = ifH (x >. rscalar 0) (rscalar 2) (rscalar 5) * x+      (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKR ZSR FTKScalar)+  printArtifactPretty artifactRev+    @?= "\\dret x1 -> let x2 = ifH (sscalar -0.0 <=. negate (sfromR x1)) (sscalar 5.0) (sscalar 2.0) in rfromS (x2 * sfromR dret)"+  printArtifactPrimalPretty artifactRev+    @?= "\\x1 -> let x2 = ifH (sscalar -0.0 <=. negate (sfromR x1)) (sscalar 5.0) (sscalar 2.0) in rfromS (x2 * sfromR x1)"++overleaf :: forall r target. (BaseTensor target, GoodScalar r)+         => target (TKR 1 r) -> target (TKR 0 r)+overleaf v = let wrap i = i `remH` fromIntegral (rwidth v)+             in rsum (rbuild @1 [50] (\[i] -> rindex v [wrap i]))++testOverleaf :: Assertion+testOverleaf =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @_ @Double [28] [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @0 overleaf (ringestData [28] [0 .. 27]))++testOverleafInt64n :: Assertion+testOverleafInt64n =+  assertEqualUpToEpsilon 1e-10+    (ringestData [28] (map round [0 :: Double,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]))+    (cgrad (kfromR @_ @Int64 . overleaf) (ringestData [28] [0 .. 27]))++testOverleafCIntn :: Assertion+testOverleafCIntn =+  assertEqualUpToEpsilon 1e-10+    (ringestData [28] (map round [0 :: Double,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]))+    (grad (kfromR @_ @CInt . overleaf) (ringestData [28] [0 .. 27]))++testOverleafCIntToFloatn :: Assertion+testOverleafCIntToFloatn =+  assertEqualUpToEpsilon 1e-10+    (rfromList0N [28] (replicate 28 (rscalar 0.0)))+    (grad (kfromR @_ @CInt . rfromIntegral . overleaf @CInt . rfloor) (ringestData @_ @Float [28] [0 .. 27]))++testOverleafInt64p :: Assertion+testOverleafInt64p =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @_ @Int64 [28] (map round [2.0 :: Double,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0]))+    (rev' @Int64 @0 overleaf (ringestData [28] [0 .. 27]))++testOverleafCIntp :: Assertion+testOverleafCIntp =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @_ @CInt [28] (map round [2.0 :: Double,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,1.0,1.0,1.0,1.0,1.0,1.0]))+    (rev' @CInt @0 overleaf (ringestData [28] [0 .. 27]))++testOverleafCIntToFloatp :: Assertion+testOverleafCIntToFloatp =+  assertEqualUpToEpsilon' 1e-10+    (ringestData @1 @Float [28] (replicate 28 0.0))+    (let f :: forall f. ADReady f => f (TKR 1 Float) -> f (TKR 0 Float)+         f = rfromIntegral . overleaf @CInt . rfloor+     in rev' @Float @0 f (ringestData [28] [0 .. 27]))++testOverleafPP :: Assertion+testOverleafPP = do+  resetVarCounter >> resetIdCounter+  let fT :: AstTensor AstMethodLet FullSpan (TKR 1 Double)+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = overleaf+      (var3, ast3) = funToAst (FTKR [28] FTKScalar) Nothing fT+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> rfromS (ssum @50 (sgather [] (sfromR v1) (\\[i2] -> [remH i2 28])))"+  resetVarCounter+  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent fT (FTKR [28] FTKScalar)+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret v1 -> rfromS (sscatter (sreplicate @50 (sfromR dret)) (\\[i5] -> [remH i5 28]))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\v1 -> rfromS (ssum0 (sgather (sfromR v1) (\\[i3] -> [remH i3 28])))"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000002 (DeltaSum (SNat @50) (STKS [] STKScalar) (DeltaShare 100000001 (DeltaGatherS [50] [] [28] (DeltaConvert (ConvCmp (ConvXS' (FTKS [28] FTKScalar)) ConvRX) (DeltaInput (InputId 0))) <function>))))"++foo :: RealFloat a => (a, a, a) -> a+foo (x, y, z) =+  let w = x * sin y+  in atan2 z w + z * w  -- note that w appears twice++testFoo :: Assertion+testFoo = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627)+    (cgrad (kfromR @_ @Double . foo) (rscalar 1.1, rscalar 2.2, rscalar 3.3))++gradFooDouble :: (Double, Double, Double) -> (Double, Double, Double)+gradFooDouble = fromDValue . cgrad foo . fromValue++testGradFooDouble :: Assertion+testGradFooDouble =+  assertEqualUpToEpsilon 1e-10+    (2.4396285219055063, -1.953374825727421, 0.9654825811012627)+    (gradFooDouble (1.1, 2.2, 3.3))++type Matrix2x2 :: Target -> Type -> Type+type Matrix2x2 f r = f (TKS '[2, 2] r)+type ThreeMatrices r = (Matrix2x2 Concrete r, Matrix2x2 Concrete r, Matrix2x2 Concrete r)+threeSimpleMatrices :: ThreeMatrices Double+threeSimpleMatrices = (srepl 1.1, srepl 2.2, srepl 3.3)+fooMatrixValue :: Matrix2x2 Concrete Double+fooMatrixValue = foo threeSimpleMatrices+gradSumFooMatrix :: ThreeMatrices Double -> ThreeMatrices Double+gradSumFooMatrix = cgrad (kfromS . ssum0 . foo)++testFooMatrix :: Assertion+testFooMatrix =+  assertEqualUpToEpsilon 1e-10+    (sfromListLinear [2,2] [4.242393641025528,4.242393641025528,4.242393641025528,4.242393641025528])+    fooMatrixValue++testGradFooMatrix :: Assertion+testGradFooMatrix =+  assertEqualUpToEpsilon 1e-10+    (sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])+    (gradSumFooMatrix threeSimpleMatrices)++fooLet :: (RealFloatH (f r), ADReady f)+       => (f r, f r, f r) -> f r+fooLet (x, y, z) =+  tlet (x * sin y) $ \w ->+    atan2H z w + z * w  -- note that w still appears twice++artifact :: AstArtifactRev (X (ThreeMatrices Double)) (TKS '[2, 2] Double)+artifact = vjpArtifact fooLet threeSimpleMatrices++testGradFooLetMatrixPP :: Assertion+testGradFooLetMatrixPP = do+  resetVarCounter+  printArtifactPretty artifact+    @?= "\\dret m1 -> let m3 = sin (tproject2 (tproject1 m1)) ; m4 = tproject1 (tproject1 m1) * m3 ; m5 = recip (tproject2 m1 * tproject2 m1 + m4 * m4) ; m7 = (negate (tproject2 m1) * m5) * dret + tproject2 m1 * dret in tpair (tpair (m3 * m7) (cos (tproject2 (tproject1 m1)) * (tproject1 (tproject1 m1) * m7))) ((m4 * m5) * dret + m4 * dret)"++testGradFooMatrixVjp :: Assertion+testGradFooMatrixVjp =+  assertEqualUpToEpsilon 1e-10+    ((sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627]) :: ThreeMatrices Double)+    (vjpInterpretArtifact artifact (toTarget threeSimpleMatrices) (srepl 1))++testGradFooMatrixRev :: Assertion+testGradFooMatrixRev =+  assertEqualUpToEpsilon 1e-10+    (sfromListLinear [2,2] [2.4396285219055063,2.4396285219055063,2.4396285219055063,2.4396285219055063],sfromListLinear [2,2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421],sfromListLinear [2,2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])+    (grad (kfromS . ssum0 . fooLet) threeSimpleMatrices)++testGradFooLetMatrixSimpPP :: Assertion+testGradFooLetMatrixSimpPP = do+  resetVarCounter+  (let ftk = FTKS @'[2, 2] [2, 2] (FTKScalar @Double)+   in printArtifactPretty+        (simplifyArtifact $ revArtifactAdapt UseIncomingCotangent fooLet (FTKProduct (FTKProduct ftk ftk) ftk)))+      @?= "\\dret m1 -> let m3 = sin (tproject2 (tproject1 m1)) ; m4 = tproject1 (tproject1 m1) * m3 ; m5 = recip (tproject2 m1 * tproject2 m1 + m4 * m4) ; m7 = (negate (tproject2 m1) * m5) * dret + tproject2 m1 * dret in tpair (tpair (m3 * m7) (cos (tproject2 (tproject1 m1)) * (tproject1 (tproject1 m1) * m7))) ((m4 * m5) * dret + m4 * dret)"++testGradFooLetMatrixSimpRPP :: Assertion+testGradFooLetMatrixSimpRPP = do+  resetVarCounter+  (let ftk = FTKR (2 :$: 2 :$: ZSR) (FTKScalar @Double)+    in printArtifactPretty (simplifyArtifact $ revArtifactAdapt UseIncomingCotangent fooLet (FTKProduct (FTKProduct ftk ftk) ftk)))+       @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2] FTKScalar)) ConvSX))) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2,2] STKScalar)) (let m3 = sin (sfromR (tproject2 (tproject1 m1))) ; m4 = sfromR (tproject1 (tproject1 m1)) * m3 ; m5 = recip (sfromR (tproject2 m1) * sfromR (tproject2 m1) + m4 * m4) ; m7 = (negate (sfromR (tproject2 m1)) * m5) * sfromR dret + sfromR (tproject2 m1) * sfromR dret in tpair (tpair (m3 * m7) (cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m7))) ((m4 * m5) * sfromR dret + m4 * sfromR dret))"++sumFooMatrix :: (ADReady f, RealFloat (Matrix2x2 f r), GoodScalar r)+             => (Matrix2x2 f r, Matrix2x2 f r, Matrix2x2 f r) -> f (TKScalar r)+sumFooMatrix = kfromS . ssum0 . foo++testfooSumMatrix :: Assertion+testfooSumMatrix =+  assertEqualUpToEpsilon 1e-10+    16.96957456410211+    (sumFooMatrix threeSimpleMatrices)++foo2 :: RealFloatH a => (a, a, a) -> a+foo2 (x, y, z) =+  let w = x * sin y+  in atan2H z w + z * w++gradFooMatrix2 :: (Differentiable r, GoodScalar r)+               => (Concrete (TKR 2 r), Concrete (TKR 2 r), Concrete (TKR 2 r))+               -> (Concrete (TKR 2 r), Concrete (TKR 2 r), Concrete (TKR 2 r))+gradFooMatrix2 = grad (kfromR . rsum0 . foo2)++testGradFooMatrix2 :: Assertion+testGradFooMatrix2 =+  assertEqualUpToEpsilon 1e-10+    (ringestData [2, 2] [2.4396285219055063 :: Double,2.4396285219055063,2.4396285219055063,2.4396285219055063], ringestData [2, 2] [-1.953374825727421,-1.953374825727421,-1.953374825727421,-1.953374825727421], ringestData [2, 2] [0.9654825811012627,0.9654825811012627,0.9654825811012627,0.9654825811012627])+    (gradFooMatrix2 (rreplicate0N [2, 2] (rscalar 1.1), rreplicate0N [2, 2] (rscalar 2.2), rreplicate0N [2, 2] (rscalar 3.3)))++testGradFooMatrixPP :: Assertion+testGradFooMatrixPP = do+  resetVarCounter+  (let ftk = FTKR (2 :$: 2 :$: ZSR) (FTKScalar @Double)+    in printArtifactPretty (revArtifactAdapt UseIncomingCotangent foo2 (FTKProduct (FTKProduct ftk ftk) ftk)))+      @?= "\\dret m1 -> let m2 = sin (sfromR (tproject2 (tproject1 m1))) ; m3 = sfromR (tproject1 (tproject1 m1)) * m2 ; m4 = recip (sfromR (tproject2 m1) * sfromR (tproject2 m1) + m3 * m3) ; m5 = sin (sfromR (tproject2 (tproject1 m1))) ; m6 = sfromR (tproject1 (tproject1 m1)) * m5 ; m8 = sfromR (tproject2 m1) * sfromR dret ; m9 = (negate (sfromR (tproject2 m1)) * m4) * sfromR dret in tpair (tpair (rfromS (m2 * m9 + m5 * m8)) (rfromS (cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m9) + cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m8)))) (rfromS ((m3 * m4) * sfromR dret + m6 * sfromR dret))"++testGradFooMatrixSimpPP :: Assertion+testGradFooMatrixSimpPP = do+  resetVarCounter+  (let ftk = FTKR (2 :$: 2 :$: ZSR) (FTKScalar @Double)+    in printArtifactPretty (simplifyArtifact $ revArtifactAdapt UseIncomingCotangent foo2 (FTKProduct (FTKProduct ftk ftk) ftk)))+      @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2,2] STKScalar)) (let m2 = sin (sfromR (tproject2 (tproject1 m1))) ; m3 = sfromR (tproject1 (tproject1 m1)) * m2 ; m4 = recip (sfromR (tproject2 m1) * sfromR (tproject2 m1) + m3 * m3) ; m5 = sin (sfromR (tproject2 (tproject1 m1))) ; m8 = sfromR (tproject2 m1) * sfromR dret ; m9 = (negate (sfromR (tproject2 m1)) * m4) * sfromR dret in tpair (tpair (m2 * m9 + m5 * m8) (cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m9) + cos (sfromR (tproject2 (tproject1 m1))) * (sfromR (tproject1 (tproject1 m1)) * m8))) ((m3 * m4) * sfromR dret + (sfromR (tproject1 (tproject1 m1)) * m5) * sfromR dret))"++gradFooScalar :: forall r. r ~ Double+              => (r, r, r) -> (r, r, r)+gradFooScalar = fromDValue . grad foo2 . fromValue++testGradFooScalar :: Assertion+testGradFooScalar =+  assertEqualUpToEpsilon 1e-10+    (2.4396285219055063, -1.953374825727421, 0.9654825811012627)+    (gradFooScalar (1.1, 2.2, 3.3))++gradCFooScalar :: forall r. r ~ Float+               => (r, r, r) -> (r, r, r)+gradCFooScalar = fromDValue . cgrad foo2 . fromValue++testGradCFooScalar :: Assertion+testGradCFooScalar =+  assertEqualUpToEpsilon 1e-10+    (2.4396284,-1.9533751,0.96548253)+    (gradCFooScalar (1.1, 2.2, 3.3))++testFooS :: Assertion+testFooS = do+  assertEqualUpToEpsilon 1e-10+    (srepl 2.4396285219055063, srepl (-1.953374825727421), srepl 0.9654825811012627)+    (grad (kfromS @_ @Double . ssum0 @'[3, 534, 3] @(TKScalar Double) . foo2) (srepl 1.1, srepl 2.2, srepl 3.3))++testFooSToFloat :: Assertion+testFooSToFloat = do+  assertEqualUpToEpsilon 1e-5+    (srepl 2.4396285219055063, srepl (-1.953374825727421), srepl 0.9654825811012627)+    (grad (kfromS @_ @Float . ssum0 . scast . foo2)+         (srepl 1.1 :: Concrete (TKS '[3, 534, 3] Double), srepl 2.2, srepl 3.3))++testFooSBoth :: Assertion+testFooSBoth = do+  assertEqualUpToEpsilon 1e-5+    (srepl 2.439628436155373, srepl (-1.9533749), srepl 0.9654825479484146)+    (grad (kfromS @_ @Float . ssum0 . scast . foo2 . (\(d, f, d2) -> (d, scast f, d2)))+         ( srepl 1.1 :: Concrete (TKS '[3, 534, 3] Double)+         , srepl 2.2 :: Concrete (TKS '[3, 534, 3] Float)+         , srepl 3.3 ))++testFooBoth :: Assertion+testFooBoth = do+  assertEqualUpToEpsilon 1e-5+    (rscalar 2.439628436155373, rscalar (-1.9533749), rscalar 0.9654825479484146)+    (grad (kfromR @_ @Float . rcast . foo2 . (\(d, f, d2) -> (d, rcast f, d2)))+         ( rscalar 1.1 :: Concrete (TKR 0 Double)+         , rscalar 2.2 :: Concrete (TKR 0 Float)+         , rscalar 3.3 ))++-- Add arrays a,b,c, but shifting b and c one to left/right+-- and then remove the first and last element.+--+-- In PyTorch+-- vstack( a[0] + b[1]+--       , a[1:N-1] + b[2:N] + c[:N-2]+--       , a[N-1] + c[N-2] )+vstackABC :: (ADReady target, GoodScalar r)+          => (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))+          -> target (TKR 1 r)+vstackABC (a, b, c) =+  let n = rwidth a+  in rconcat [ rreplicate 1 (a ! [0] + b ! [1])+             , rslice 1 (n - 2) a + rslice 2 (n - 2) b + rslice 0 (n - 2) c+             , rreplicate 1 (a ! [fromIntegral n - 1]+                             + c ! [fromIntegral n - 2]) ]++vstackBuild :: (ADReady target, GoodScalar r)+            => (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))+            -> target (TKR 1 r)+vstackBuild (a, b, c) =+  let n = rwidth a+  in rbuild1 n (\i ->+       ifH (i ==. 0)+           (a ! [0] + b ! [1])+           (ifH (i ==. fromIntegral n - 1)+                (a ! [fromIntegral n - 1] + c ! [fromIntegral n - 2])+                (a ! [i] + b ! [i + 1] + c ! [i - 1])))++testTrustVstackConcatRepl10 :: Assertion+testTrustVstackConcatRepl10 = do+  vstackABC @Concrete @Double (rrepl [10] 1, rrepl [10] 2, rrepl [10] 3)+  @?= rfromListLinear [10] [3.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,4.0]++testTrustVstackConcatIota10 :: Assertion+testTrustVstackConcatIota10 = do+  vstackABC @Concrete @Double (riota 10, riota 10, riota 10)+  @?= rfromListLinear [10] [1.0,3.0,6.0,9.0,12.0,15.0,18.0,21.0,24.0,17.0]++replIota :: (ADReady target, GoodScalar r)+         => Int -> (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))+replIota n =+  ( rconcrete (unConcrete $ rrepl [n] 1 * riota n)+  , rconcrete (unConcrete $ rrepl [n] 2 * riota n)+  , rconcrete (unConcrete $ rrepl [n] 3 * riota n) )++testTrustVstackConcatReplIota10 :: Assertion+testTrustVstackConcatReplIota10 = do+  vstackABC @Concrete @Double (replIota 10)+  @?= rfromListLinear [10] [2.0,5.0,11.0,17.0,23.0,29.0,35.0,41.0,47.0,33.0]++nN :: Int+nN = (round :: Double -> Int) 1e6  -- 1e6++trustedResult :: Concrete (TKR 1 Double)+trustedResult =+  rcast (vstackABC @Concrete @Float (replIota nN))+    -- the cast prevents computation sharing with the first test below++testVstackWarmup :: Assertion+testVstackWarmup = do+  trustedResult+  @?= trustedResult++testVstackConcatConcrete :: Assertion+testVstackConcatConcrete = do+  vstackABC @Concrete @Double (replIota nN)+  @?= trustedResult++testVstackBuildConcrete :: Assertion+testVstackBuildConcrete = do+  vstackBuild @Concrete @Double (replIota nN)+  @?= trustedResult++testVstackConcatAst :: Assertion+testVstackConcatAst = do+  interpretAstFull @Concrete+    emptyEnv+    (vstackABC @(AstTensor AstMethodLet FullSpan) @Double+               (replIota nN))+  @?= trustedResult++testVstackBuildAst :: Assertion+testVstackBuildAst = do+  interpretAstFull @Concrete+    emptyEnv+    (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                 (replIota nN))+  @?= trustedResult++testVstackBuildAstSimp :: Assertion+testVstackBuildAstSimp = do+  interpretAstFull @Concrete+    emptyEnv+      (simplifyInlineContract+         (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                      (replIota nN)))+  @?= trustedResult++testVstackBuildAstPP :: Assertion+testVstackBuildAstPP = do+  resetVarCounter+  let (var3, ast3) =+        funToAst (FTKProduct (FTKProduct (FTKR [10] FTKScalar) (FTKR [10] FTKScalar)) (FTKR [10] FTKScalar))+          Nothing (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double . fromTarget)+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstPretty ast3+    @?= "\\v1 -> rfromS (let x9 = sfromR (tproject1 (tproject1 v1)) !$ [0] + sfromR (tproject2 (tproject1 v1)) !$ [1] ; v10 = let x6 = sfromR (tproject1 (tproject1 v1)) !$ [9] + sfromR (tproject2 v1) !$ [8] ; v7 = (sfromR (tproject1 (tproject1 v1)) + sappend (sslice (SNat @1) (SNat @9) (sfromR (tproject2 (tproject1 v1)))) (sconcrete (sfromListLinear [1] [0.0]))) + sappend (sconcrete (sfromListLinear [1] [0.0])) (sslice (SNat @0) (SNat @9) (sfromR (tproject2 v1))) in sappend (sslice (SNat @0) (SNat @9) v7) (sreplicate @1 x6) in sappend (sreplicate @1 x9) (sslice (SNat @1) (SNat @9) v10))"+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstPretty (simplifyInlineContract ast3)+    @?= "\\v1 -> rfromS (sappend (sreplicate @1 (sfromR (tproject1 (tproject1 v1)) !$ [0] + sfromR (tproject2 (tproject1 v1)) !$ [1])) (sappend ((sslice (SNat @1) (SNat @8) (sfromR (tproject1 (tproject1 v1))) + sslice (SNat @2) (SNat @8) (sfromR (tproject2 (tproject1 v1)))) + sslice (SNat @0) (SNat @8) (sfromR (tproject2 v1))) (sreplicate @1 (sfromR (tproject1 (tproject1 v1)) !$ [9] + sfromR (tproject2 v1) !$ [8]))))"++{- The above is:+\v1 ->+  rfromS+    (sappend+       (sreplicate @1+          (sfromR (tproject1 (tproject1 v1)) !$ [0] ++           sfromR (tproject2 (tproject1 v1)) !$ [1]))+       (sappend+          ((sslice (SNat @1) (SNat @8) (sfromR (tproject1 (tproject1 v1))) ++            sslice (SNat @2) (SNat @8) (sfromR (tproject2 (tproject1 v1)))) ++           sslice (SNat @0) (SNat @8) (sfromR (tproject2 v1)))+          (sreplicate @1+             (sfromR (tproject1 (tproject1 v1)) !$ [9] ++              sfromR (tproject2 v1) !$ [8]))))+-}++replIota2 :: (ADReady target, GoodScalar r)+          => Int -> (target (TKR 1 r), target (TKR 1 r), target (TKR 1 r))+replIota2 n =+  (rrepl [n] 1 * riota n, rrepl [n] 2 * riota n, rrepl [n] 3 * riota n)++testVstackConcatConcrete2 :: Assertion+testVstackConcatConcrete2 = do+  vstackABC @Concrete @Double (replIota nN)+  @?= trustedResult++testVstackBuildConcrete2 :: Assertion+testVstackBuildConcrete2 = do+  vstackBuild @Concrete @Double (replIota2 nN)+  @?= trustedResult++testVstackConcatAst2 :: Assertion+testVstackConcatAst2 = do+  interpretAstFull @Concrete+    emptyEnv+    (vstackABC @(AstTensor AstMethodLet FullSpan) @Double+               (replIota2 nN))+  @?= trustedResult++testVstackBuildAst2 :: Assertion+testVstackBuildAst2 = do+  interpretAstFull @Concrete+    emptyEnv+    (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                 (replIota2 nN))+  @?= trustedResult++testVstackBuildAstSimp2 :: Assertion+testVstackBuildAstSimp2 = do+  interpretAstFull @Concrete+    emptyEnv+      (simplifyInlineContract+         (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                      (replIota2 nN)))+  @?= trustedResult++testVstackBuildAstPP2 :: Assertion+testVstackBuildAstPP2 = do+  resetVarCounter+  (printAstPretty+     (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                  (replIota2 10)))+    @?= "rfromS (let x14 = sfromIntegral (sscalar 0) + sscalar 2.0 * sfromIntegral (sscalar 1) ; v15 = let x11 = sfromIntegral (sscalar 9) + sscalar 3.0 * sfromIntegral (sscalar 8) ; v12 = (siota (SNat @10) + (let v2 = sconcrete (sreplicate [10] 1) + siota (SNat @10) in sgather (sconcrete (sreplicate [10] 2.0)) (\\[i8] -> [kfromS (v2 !$ [i8])]) * sgather (sfromVector (fromList [sfromIntegral v2, sconcrete (sreplicate [10] 0.0)])) (\\[i9] -> [ifH (sscalar (-9) <=. negate (v2 !$ [i9])) 0 1, i9]))) + (let v3 = sconcrete (sreplicate [10] (-1)) + siota (SNat @10) in sgather (sconcrete (sreplicate [10] 3.0)) (\\[i6] -> [kfromS (v3 !$ [i6])]) * sgather (sfromVector (fromList [sfromIntegral v3, sconcrete (sreplicate [10] 0.0)])) (\\[i7] -> [ifH (sscalar 0 <=. v3 !$ [i7]) 0 1, i7])) in sappend (sslice (SNat @0) (SNat @9) v12) (sreplicate @1 x11) in sappend (sreplicate @1 x14) (sslice (SNat @1) (SNat @9) v15))"+  (printAstPretty+     (simplifyInlineContract+        (vstackBuild @(AstTensor AstMethodLet FullSpan) @Double+                     (replIota2 10))))+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [2.0])) (sappend (sconcrete (sfromListLinear [8] [1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0]) + (sgather (sconcrete (sreplicate [10] 2.0)) (\\[i23] -> [kfromS (sconcrete (sfromListLinear [10] [1,2,3,4,5,6,7,8,9,10]) !$ [1 + i23])]) * sgather (sfromVector (fromList [sfromIntegral (sconcrete (sfromListLinear [8] [2,3,4,5,6,7,8,9])), sconcrete (sreplicate [8] 0.0)])) (\\[i24] -> [ifH (sscalar (-9) <=. negate (sconcrete (sfromListLinear [10] [1,2,3,4,5,6,7,8,9,10]) !$ [1 + i24])) 0 1, i24]) + sgather (sconcrete (sreplicate [10] 3.0)) (\\[i21] -> [kfromS (sconcrete (sfromListLinear [10] [-1,0,1,2,3,4,5,6,7,8]) !$ [1 + i21])]) * sgather (sfromVector (fromList [sfromIntegral (sconcrete (sfromListLinear [8] [0,1,2,3,4,5,6,7])), sconcrete (sreplicate [8] 0.0)])) (\\[i22] -> [ifH (sscalar 0 <=. sconcrete (sfromListLinear [10] [-1,0,1,2,3,4,5,6,7,8]) !$ [1 + i22]) 0 1, i22]))) (sconcrete (sfromListLinear [1] [33.0]))))"+      -- TODO: this reduced to "rfromS (sconcrete (sfromListLinear [10] [2.0,5.0,11.0,17.0,23.0,29.0,35.0,41.0,47.0,33.0]))" when tshare/tlet were inserted less often++testFooPP :: Assertion+testFooPP = do+  resetVarCounter+  let fooT = foo2 @(AstTensor AstMethodLet FullSpan (TKR 0 Double))+      foo3 x = fooT (x, x, x)+      (var3, ast3) = funToAst (FTKR ZSR FTKScalar) Nothing foo3+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\x1 -> rfromS (atan2H (sfromR x1) (sfromR x1 * sin (sfromR x1)) + sfromR x1 * (sfromR x1 * sin (sfromR x1)))"+  resetVarCounter+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fooT (FTKProduct (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar)) (FTKR ZSR FTKScalar))+  printArtifactSimple artifactRev+    @?= "\\dret x1 -> tlet (sin (sfromR (tproject2 (tproject1 x1)))) (\\x2 -> tlet (sfromR (tproject1 (tproject1 x1)) * x2) (\\x3 -> tlet (recip (sfromR (tproject2 x1) * sfromR (tproject2 x1) + x3 * x3)) (\\x4 -> tlet (sin (sfromR (tproject2 (tproject1 x1)))) (\\x5 -> tlet (sfromR (tproject1 (tproject1 x1)) * x5) (\\x6 -> tlet (sfromR (tproject2 x1) * sfromR dret) (\\x8 -> tlet ((negate (sfromR (tproject2 x1)) * x4) * sfromR dret) (\\x9 -> tpair (tpair (rfromS (x2 * x9 + x5 * x8)) (rfromS (cos (sfromR (tproject2 (tproject1 x1))) * (sfromR (tproject1 (tproject1 x1)) * x9) + cos (sfromR (tproject2 (tproject1 x1))) * (sfromR (tproject1 (tproject1 x1)) * x8)))) (rfromS ((x3 * x4) * sfromR dret + x6 * sfromR dret)))))))))"+  printArtifactPrimalSimple artifactRev+    @?= "\\x1 -> tlet (sin (sfromR (tproject2 (tproject1 x1)))) (\\x2 -> tlet (sfromR (tproject1 (tproject1 x1)) * x2) (\\x3 -> tlet (sin (sfromR (tproject2 (tproject1 x1)))) (\\x5 -> tlet (sfromR (tproject1 (tproject1 x1)) * x5) (\\x6 -> rfromS (atan2H (sfromR (tproject2 x1)) x3 + sfromR (tproject2 x1) * x6)))))"++fooLetOld :: forall target r n.+          (RealFloatH (target (TKR n r)), LetTensor target)+       => (target (TKR n r), target (TKR n r), target (TKR n r)) -> target (TKR n r)+fooLetOld (x, y, z) =+  let w0 = x * sin y+  in tlet w0 $ \w ->+     atan2H z w + z * w++testFooLet :: Assertion+testFooLet = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627)+    (grad (kfromR @_ @Double . fooLetOld) (rscalar 1.1, rscalar 2.2, rscalar 3.3))++testFooLetPP :: Assertion+testFooLetPP = do+  resetVarCounter+  let fooLetT = fooLetOld @(AstTensor AstMethodLet FullSpan) @Double+      fooLet3 x = fooLetT (x, x, x)+      (var3, ast3) = funToAst (FTKR ZSR FTKScalar) Nothing fooLet3+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\x1 -> rfromS (tlet (sfromR x1 * sin (sfromR x1)) (\\x2 -> atan2H (sfromR x1) x2 + sfromR x1 * x2))"+  resetVarCounter+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fooLetT (FTKProduct (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar)) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX))) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (STKS [] STKScalar)) (let x3 = sin (sfromR (tproject2 (tproject1 x1))) ; x4 = sfromR (tproject1 (tproject1 x1)) * x3 ; x5 = recip (sfromR (tproject2 x1) * sfromR (tproject2 x1) + x4 * x4) ; x7 = (negate (sfromR (tproject2 x1)) * x5) * sfromR dret + sfromR (tproject2 x1) * sfromR dret in tpair (tpair (x3 * x7) (cos (sfromR (tproject2 (tproject1 x1))) * (sfromR (tproject1 (tproject1 x1)) * x7))) ((x4 * x5) * sfromR dret + x4 * sfromR dret))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (let x4 = sfromR (tproject1 (tproject1 x1)) * sin (sfromR (tproject2 (tproject1 x1))) in atan2H (sfromR (tproject2 x1)) x4 + sfromR (tproject2 x1) * x4)"++shapedListProd :: forall k target r. (BaseTensor target, GoodScalar r)+               => ListR k (target (TKS '[] r)) -> target (TKS '[] r)+shapedListProd = foldr1 (*)++testListProdPP :: Assertion+testListProdPP = do+  resetVarCounter+  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKS '[] Double)) -> AstTensor AstMethodLet FullSpan (TKS '[] Double)+      fT = shapedListProd+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKS ZSS FTKScalar) (FTKProduct (FTKS ZSS FTKScalar) (FTKProduct (FTKS ZSS FTKScalar) (FTKProduct (FTKS ZSS FTKScalar) ftkUnit))))+  printArtifactSimple artifactRev+    @?= "\\dret x1 -> tlet (tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1)))) (\\x2 -> tlet (tproject1 (tproject2 x1) * x2) (\\x3 -> tlet (tproject1 x1 * dret) (\\x5 -> tlet (tproject1 (tproject2 x1) * x5) (\\x6 -> tpair (x3 * dret) (tpair (x2 * x5) (tpair (tproject1 (tproject2 (tproject2 (tproject2 x1))) * x6) (tpair (tproject1 (tproject2 (tproject2 x1)) * x6) Z1)))))))"+  printArtifactPrimalSimple artifactRev+    @?= "\\x1 -> tlet (tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1)))) (\\x2 -> tlet (tproject1 (tproject2 x1) * x2) (\\x3 -> tproject1 x1 * x3))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> let x2 = tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1))) ; x5 = tproject1 x1 * dret ; x6 = tproject1 (tproject2 x1) * x5 in tpair ((tproject1 (tproject2 x1) * x2) * dret) (tpair (x2 * x5) (tpair (tproject1 (tproject2 (tproject2 (tproject2 x1))) * x6) (tpair (tproject1 (tproject2 (tproject2 x1)) * x6) Z1)))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> tproject1 x1 * (tproject1 (tproject2 x1) * (tproject1 (tproject2 (tproject2 x1)) * tproject1 (tproject2 (tproject2 (tproject2 x1)))))"++rankedListProdr :: forall k target r. (BaseTensor target, GoodScalar r)+                => ListR k (target (TKR 0 r)) -> target (TKR 0 r)+rankedListProdr = foldr1 (*)++testListProdrPP :: Assertion+testListProdrPP = do+  resetVarCounter+  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListProdr+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x2 = sfromR (tproject1 (tproject2 (tproject2 x1))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) ; x5 = sfromR (tproject1 x1) * sfromR dret ; x6 = sfromR (tproject1 (tproject2 x1)) * x5 in tpair ((sfromR (tproject1 (tproject2 x1)) * x2) * sfromR dret) (tpair (x2 * x5) (tpair (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x6) (tpair (sfromR (tproject1 (tproject2 (tproject2 x1))) * x6) Z1))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sfromR (tproject1 x1) * (sfromR (tproject1 (tproject2 x1)) * (sfromR (tproject1 (tproject2 (tproject2 x1))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++testListProdrLongPP :: Assertion+testListProdrLongPP = do+  resetVarCounter+  let fT :: ListR 13 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListProdr+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit)))))))))))))+  printArtifactSimple artifactRev+    @?= "\\dret x1 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))))) (\\x2 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x2) (\\x3 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x3) (\\x4 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x4) (\\x5 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x5) (\\x6 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x6) (\\x7 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x7) (\\x8 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x8) (\\x9 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x9) (\\x10 -> tlet (sfromR (tproject1 (tproject2 (tproject2 x1))) * x10) (\\x11 -> tlet (sfromR (tproject1 (tproject2 x1)) * x11) (\\x12 -> tlet (sfromR (tproject1 x1) * sfromR dret) (\\x14 -> tlet (sfromR (tproject1 (tproject2 x1)) * x14) (\\x15 -> tlet (sfromR (tproject1 (tproject2 (tproject2 x1))) * x15) (\\x16 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x16) (\\x17 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x17) (\\x18 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x18) (\\x19 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x19) (\\x20 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x20) (\\x21 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x21) (\\x22 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x22) (\\x23 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x23) (\\x24 -> tpair (rfromS (x12 * sfromR dret)) (tpair (rfromS (x11 * x14)) (tpair (rfromS (x10 * x15)) (tpair (rfromS (x9 * x16)) (tpair (rfromS (x8 * x17)) (tpair (rfromS (x7 * x18)) (tpair (rfromS (x6 * x19)) (tpair (rfromS (x5 * x20)) (tpair (rfromS (x4 * x21)) (tpair (rfromS (x3 * x22)) (tpair (rfromS (x2 * x23)) (tpair (rfromS (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))))) * x24)) (tpair (rfromS (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * x24)) Z1))))))))))))))))))))))))))))))))))"+  printArtifactPrimalSimple artifactRev+    @?= "\\x1 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))))) (\\x2 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x2) (\\x3 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x3) (\\x4 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x4) (\\x5 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x5) (\\x6 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x6) (\\x7 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x7) (\\x8 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x8) (\\x9 -> tlet (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x9) (\\x10 -> tlet (sfromR (tproject1 (tproject2 (tproject2 x1))) * x10) (\\x11 -> tlet (sfromR (tproject1 (tproject2 x1)) * x11) (\\x12 -> rfromS (sfromR (tproject1 x1) * x12))))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId))))))))))))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar))))))))))))) (let x2 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))))) ; x3 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x2 ; x4 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x3 ; x5 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x4 ; x6 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x5 ; x7 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x6 ; x8 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x7 ; x9 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x8 ; x10 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x9 ; x11 = sfromR (tproject1 (tproject2 (tproject2 x1))) * x10 ; x14 = sfromR (tproject1 x1) * sfromR dret ; x15 = sfromR (tproject1 (tproject2 x1)) * x14 ; x16 = sfromR (tproject1 (tproject2 (tproject2 x1))) * x15 ; x17 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x16 ; x18 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * x17 ; x19 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * x18 ; x20 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * x19 ; x21 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * x20 ; x22 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * x21 ; x23 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * x22 ; x24 = sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * x23 in tpair ((sfromR (tproject1 (tproject2 x1)) * x11) * sfromR dret) (tpair (x11 * x14) (tpair (x10 * x15) (tpair (x9 * x16) (tpair (x8 * x17) (tpair (x7 * x18) (tpair (x6 * x19) (tpair (x5 * x20) (tpair (x4 * x21) (tpair (x3 * x22) (tpair (x2 * x23) (tpair (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))))) * x24) (tpair (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * x24) Z1)))))))))))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sfromR (tproject1 x1) * (sfromR (tproject1 (tproject2 x1)) * (sfromR (tproject1 (tproject2 (tproject2 x1))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 x1))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1))))))))))) * (sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 x1)))))))))))))))))))))))))"++testListProd :: Assertion+testListProd = do+  assertEqualUpToEpsilon 1e-10+    [srepl 24, srepl 12, srepl 8, srepl 6]+    (grad (kfromS @_ @Double . shapedListProd @4) [srepl 1, srepl 2, srepl 3, srepl 4])++testListProdr :: Assertion+testListProdr = do+  assertEqualUpToEpsilon 1e-10+    [rscalar 24, rscalar 12, rscalar 8, rscalar 6]+    (grad (kfromR @_ @Double . rankedListProdr @4) [rscalar 1, rscalar 2, rscalar 3, rscalar 4])++rankedListSumr :: (BaseTensor target, GoodScalar r)+               => ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSumr = foldr1 (+)++testListSumrPP :: Assertion+testListSumrPP = do+  resetVarCounter >> resetIdCounter+  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSumr+  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tpair dret (tpair dret (tpair dret (tpair dret Z1)))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sfromR (tproject1 x1) + (sfromR (tproject1 (tproject2 x1)) + (sfromR (tproject1 (tproject2 (tproject2 x1))) + sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000003 (DeltaAdd (DeltaConvert (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (DeltaInput (InputId 0))) (DeltaShare 100000002 (DeltaAdd (DeltaConvert (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (DeltaInput (InputId 1))) (DeltaShare 100000001 (DeltaAdd (DeltaConvert (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (DeltaInput (InputId 2))) (DeltaConvert (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (DeltaInput (InputId 3)))))))))"++-- Note that the function is not associative, so foldr vs foldl matters.+rankedListSum2r :: (BaseTensor target, GoodScalar r)+                => ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum2r = foldr1 (\x y -> x + rscalar 2 * y)++testListSum2rPP :: Assertion+testListSum2rPP = do+  resetVarCounter+  let fT ::  ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum2r+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKR (SNat @0) STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x5 = sscalar 2.0 * sfromR dret ; x6 = sscalar 2.0 * x5 in tpair dret (tpair x5 (tpair x6 (tpair (sscalar 2.0 * x6) Z1))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sfromR (tproject1 x1) + (sscalar 2.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 4.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++rankedListSum22r :: (BaseTensor target, GoodScalar r)+                 => ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum22r = foldr1 (\x y -> rscalar 2 * x + rscalar 2 * y)++testListSum22rPP :: Assertion+testListSum22rPP = do+  resetVarCounter+  let fT ::  ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum22r+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x5 = sscalar 2.0 * sfromR dret ; x6 = sscalar 2.0 * x5 in tpair (sscalar 2.0 * sfromR dret) (tpair (sscalar 2.0 * x5) (tpair (sscalar 2.0 * x6) (tpair (sscalar 2.0 * x6) Z1))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + (sscalar 4.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++-- Note how this tlet did not change anything, in particular the sharing.+rankedListSumk22r :: ( BaseTensor target, LetTensor target+                     , GoodScalar r )+                 =>  ListR k (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSumk22r = foldr1 (\x y -> tlet (rscalar 2) (\k -> k * x + k * y))++testListSumk22rPP :: Assertion+testListSumk22rPP = do+  resetVarCounter+  let fT ::  ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSumk22r+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x5 = sscalar 2.0 * sfromR dret ; x6 = sscalar 2.0 * x5 in tpair (sscalar 2.0 * sfromR dret) (tpair (sscalar 2.0 * x5) (tpair (sscalar 2.0 * x6) (tpair (sscalar 2.0 * x6) Z1))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + (sscalar 4.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++rankedListSum2xpyr :: (BaseTensor target, GoodScalar r)+                   =>  ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum2xpyr = foldr1 (\x y -> rscalar 2 * (x + y))++testListSum2xpyrPP :: Assertion+testListSum2xpyrPP = do+  resetVarCounter+  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum2xpyr+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x6 = sscalar 2.0 * sfromR dret ; x7 = sscalar 2.0 * x6 ; x8 = sscalar 2.0 * x7 in tpair x6 (tpair x7 (tpair x8 (tpair x8 Z1))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + (sscalar 4.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 8.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++rankedListSum2xyr :: (BaseTensor target, GoodScalar r)+                  => ListR 4 (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum2xyr = foldr1 (\x y -> rscalar 2 * (x * y))++testListSum2xyrPP :: Assertion+testListSum2xyrPP = do+  resetVarCounter+  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum2xyr+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x3 = sscalar 2.0 * (sfromR (tproject1 (tproject2 (tproject2 x1))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1))))) ; x8 = sscalar 2.0 * sfromR dret ; x9 = sscalar 2.0 * (sfromR (tproject1 x1) * x8) ; x10 = sscalar 2.0 * (sfromR (tproject1 (tproject2 x1)) * x9) in tpair (sscalar 2.0 * ((sfromR (tproject1 (tproject2 x1)) * x3) * x8)) (tpair (x3 * x9) (tpair (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))) * x10) (tpair (sfromR (tproject1 (tproject2 (tproject2 x1))) * x10) Z1))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sscalar 8.0 * (sfromR (tproject1 x1) * (sfromR (tproject1 (tproject2 x1)) * (sfromR (tproject1 (tproject2 (tproject2 x1))) * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1))))))))"++ranked2xy :: (BaseTensor target, GoodScalar r)+          => (target (TKR 0 r), target (TKR 0 r)) -> target (TKR 0 r)+ranked2xy = \(x, y) -> rscalar 2 * x * y++test2xyPP :: Assertion+test2xyPP = do+  resetVarCounter+  let fT :: (AstTensor AstMethodLet FullSpan (TKR 0 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = ranked2xy+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair (sscalar 2.0 * (sfromR (tproject2 x1) * sfromR dret)) (sscalar 2.0 * (sfromR (tproject1 x1) * sfromR dret)))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sscalar 2.0 * (sfromR (tproject1 x1) * sfromR (tproject2 x1)))"++-- Note that the function is not associative, so foldr vs foldl matters.+rankedListSum23r :: forall k target r. (BaseTensor target, GoodScalar r)+                 => ListR k (target (TKR 0 r)) -> target (TKR 0 r)+rankedListSum23r = foldr1 (\x y -> rscalar 2 * x + rscalar 3 * y)++testListSum23rPP :: Assertion+testListSum23rPP = do+  resetVarCounter+  let fT :: ListR 4 (AstTensor AstMethodLet FullSpan (TKR 0 Double)) -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = rankedListSum23r+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) (FTKProduct (FTKR ZSR FTKScalar) ftkUnit))))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) STKScalar)))) (let x5 = sscalar 3.0 * sfromR dret ; x6 = sscalar 3.0 * x5 in tpair (sscalar 2.0 * sfromR dret) (tpair (sscalar 2.0 * x5) (tpair (sscalar 2.0 * x6) (tpair (sscalar 3.0 * x6) Z1))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + (sscalar 6.0 * sfromR (tproject1 (tproject2 x1)) + (sscalar 18.0 * sfromR (tproject1 (tproject2 (tproject2 x1))) + sscalar 27.0 * sfromR (tproject1 (tproject2 (tproject2 (tproject2 x1)))))))"++ranked23 :: (BaseTensor target, GoodScalar r)+         => (target (TKR 0 r), target (TKR 0 r)) -> target (TKR 0 r)+ranked23 = \(x, y) -> rscalar 2 * x + rscalar 3 * y++test23PP :: Assertion+test23PP = do+  resetVarCounter+  let fT :: (AstTensor AstMethodLet FullSpan (TKR 0 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+         -> AstTensor AstMethodLet FullSpan (TKR 0 Double)+      fT = ranked23+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent fT (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret x1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair (sscalar 2.0 * sfromR dret) (sscalar 3.0 * sfromR dret))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\x1 -> rfromS (sscalar 2.0 * sfromR (tproject1 x1) + sscalar 3.0 * sfromR (tproject2 x1))"++reluPrimal+  :: forall target n r.+     (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+  => target (TKR n r) -> target (TKR n r)+reluPrimal v =+  let oneIfGtZero = rmap0N (\x -> ifH (x <=. rscalar 0) (rscalar 0.0) (rscalar 1.0))+                           (rprimalPart v)+  in scale2 oneIfGtZero v++scale2 :: forall target r n.+          (ADReady target, GoodScalar r, KnownNat n)+       => PrimalOf target (TKR n r) -> target (TKR n r) -> target (TKR n r)+scale2 a d = rfromPrimal @target a * d++testReluPP :: Assertion+testReluPP = do+  resetVarCounter >> resetIdCounter+  let reluT :: AstTensor AstMethodLet FullSpan (TKR 2 Double) -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT = reluPrimal+      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing reluT+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> rfromS (tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i4, i3] -> [ifH (sscalar -0.0 <=. negate (tprimalPart (sfromR m1) !$ [i4, i3])) 0 1])) * sfromR m1)"+  resetVarCounter+  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent reluT (FTKR [3, 4] FTKScalar)+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i5, i6])) 0 1]) * sfromR dret)"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i5, i6])) 0 1]) * sfromR m1)"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000004 (DeltaScale <primal> (DeltaConvert (ConvCmp (ConvXS' (FTKS [3,4] FTKScalar)) ConvRX) (DeltaInput (InputId 0)))))"++testReluPP2 :: Assertion+testReluPP2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 1 Double)+      reluT2 (t, r) = reluPrimal (t * rreplicate 5 r)+      (var3, ast3) = funToAst (FTKR [5] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> rfromS (tfromPrimal (STKS [5] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i2] -> [ifH (sscalar -0.0 <=. sscalar (-7.0) * tprimalPart (sfromR v1) !$ [i2]) 0 1])) * (sfromR v1 * tfromPrimal (STKS [5] STKScalar) (sconcrete (sreplicate [5] 7.0))))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [5] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret v1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [5] STKScalar) (STKS [] STKScalar)) (let v9 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 v1) !$ [i3]) * sfromR (tproject2 v1)) 0 1]) * sfromR dret in tpair (sreplicate @5 (sfromR (tproject2 v1)) * v9) (sdot0 (sfromR (tproject1 v1)) v9))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\v1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 v1) !$ [i3]) * sfromR (tproject2 v1)) 0 1]) * (sfromR (tproject1 v1) * sreplicate @5 (sfromR (tproject2 v1))))"++testReluSimpler :: Assertion+testReluSimpler = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3, 4] [1.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,0.0,0.0,1.0,1.0])+    (rev' @Double @2 relu (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12]))++testReluSimplerPP :: Assertion+testReluSimplerPP = do+  resetVarCounter >> resetIdCounter+  let reluT :: AstTensor AstMethodLet FullSpan (TKR 2 Double) -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT = relu+      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing reluT+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> rfromS (tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i4, i3] -> [ifH (sscalar -0.0 <=. negate (tprimalPart (sfromR m1) !$ [i4, i3])) 0 1])) * sfromR m1)"+  resetVarCounter+  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent reluT (FTKR [3, 4] FTKScalar)+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i5, i6])) 0 1]) * sfromR dret)"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i6] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i5, i6])) 0 1]) * sfromR m1)"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000004 (DeltaScale <primal> (DeltaConvert (ConvCmp (ConvXS' (FTKS [3,4] FTKScalar)) ConvRX) (DeltaInput (InputId 0)))))"++testReluSimplerPP2 :: Assertion+testReluSimplerPP2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 1 Double)+      reluT2 (t, r) = relu (t * rreplicate 5 r)+      (var3, ast3) = funToAst (FTKR [5] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> rfromS (tlet (sfromR v1 * tfromPrimal (STKS [5] STKScalar) (sconcrete (sreplicate [5] 7.0))) (\\v2 -> tfromPrimal (STKS [5] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3] -> [ifH (sscalar -0.0 <=. negate (sfromR (tprimalPart (rfromS (v2 !$ [i3]))))) 0 1])) * v2))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [5] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret v1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [5] STKScalar) (STKS [] STKScalar)) (let v8 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 v1) !$ [i5]) * sfromR (tproject2 v1)) 0 1]) * sfromR dret in tpair (sreplicate @5 (sfromR (tproject2 v1)) * v8) (sdot0 (sfromR (tproject1 v1)) v8))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\v1 -> rfromS (let v4 = sfromR (tproject1 v1) * sreplicate @5 (sfromR (tproject2 v1)) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5] -> [ifH (sscalar -0.0 <=. negate (v4 !$ [i5])) 0 1]) * v4)"++testReluSimplerPP3 :: Assertion+testReluSimplerPP3 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = relu (t * rreplicate 3 (rreplicate 4 r))+      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> rfromS (tlet (sfromR m1 * tfromPrimal (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 7.0))) (\\m2 -> tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i4] -> [ifH (sscalar -0.0 <=. negate (tprimalPart m2 !$ [i5, i4])) 0 1])) * m2))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [3, 4] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [3,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [3,4] STKScalar) (STKS [] STKScalar)) (let m11 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 m1) !$ [i7, i8]) * sfromR (tproject2 m1)) 0 1]) * sfromR dret in tpair (sreplicate @3 (sreplicate @4 (sfromR (tproject2 m1))) * m11) (sdot0 (sfromR (tproject1 m1)) m11))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (let m6 = sfromR (tproject1 m1) * sreplicate @3 (sreplicate @4 (sfromR (tproject2 m1))) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) * m6)"++testReluSimpler3 :: Assertion+testReluSimpler3 = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = relu (t * rreplicate 3 (rreplicate 4 r))+  assertEqualUpToEpsilon 1e-10+    ( ringestData [3, 4] [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , rscalar 57.1 )+    (grad (kfromR . rsum0 . reluT2) (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], rscalar 7))++testReluSimplerPP4 :: Assertion+testReluSimplerPP4 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = relu (t * rreplicate0N [3, 4] r)+      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> rfromS (tlet (sfromR m1 * tfromPrimal (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 7.0))) (\\m2 -> tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i4] -> [ifH (sscalar -0.0 <=. negate (tprimalPart m2 !$ [i5, i4])) 0 1])) * m2))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [3, 4] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [3,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX))) (STKProduct (STKS [3,4] STKScalar) (STKS [] STKScalar)) (let m11 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 m1) !$ [i7, i8]) * sfromR (tproject2 m1)) 0 1]) * sfromR dret in tpair (sreplicate @3 (sreplicate @4 (sfromR (tproject2 m1))) * m11) (sdot0 (sfromR (tproject1 m1)) m11))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (let m6 = sfromR (tproject1 m1) * sreplicate @3 (sreplicate @4 (sfromR (tproject2 m1))) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) * m6)"++testReluSimpler4 :: Assertion+testReluSimpler4 = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = relu (t * rreplicate0N [3, 4] r)+  assertEqualUpToEpsilon 1e-10+    ( ringestData [3, 4] [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , rscalar 57.1 )+    (grad (kfromR . rsum0 . reluT2) (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], rscalar 7))++testReluSimplerPP4s :: Assertion+testReluSimplerPP4s = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Float), AstTensor AstMethodLet FullSpan (TKS '[] Float))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Float)+      reluT2 (t, r) = reluS (t * sreplicate0N r)+      (var3, ast3) = funToAst (FTKS knownShS FTKScalar) Nothing (\t -> reluT2 (t, srepl 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> tlet (m1 * tfromPrimal (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 7.0))) (\\m2 -> tfromPrimal (STKS [3,4] STKScalar) (sgather [] (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i5, i4] -> [ifH (sscalar -0.0 <=. negate (tprimalPart m2 !$ [i5, i4])) 0 1])) * m2)"++testReluSimplerPP4s2 :: Assertion+testReluSimplerPP4s2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)+      -- This is tweaked compared to above to avoid test artifacts coming+      -- from counter resets, which are inherently unsafe (cse, etc.).+      reluT2 (t, r) = reluS (t * sreplicate0N r)+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKS [3, 4] FTKScalar) (FTKS ZSS FTKScalar))+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let m6 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) ; m9 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) ; m11 = m9 * dret in tpair (sreplicate @3 (sreplicate @4 (tproject2 m1)) * m11) (ssum @4 (ssum @3 (tproject1 m1 * m11)))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> let m6 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) ; m9 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) in m9 * m6"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> let m11 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (tproject1 m1 !$ [i7, i8]) * tproject2 m1) 0 1]) * dret in tpair (sreplicate @3 (sreplicate @4 (tproject2 m1)) * m11) (sdot0 (tproject1 m1) m11)"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> let m6 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i7, i8] -> [ifH (sscalar -0.0 <=. negate (m6 !$ [i7, i8])) 0 1]) * m6"++testReluSimpler4s :: Assertion+testReluSimpler4s = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)+      reluT2 (t, r) = reluS (t * sreplicate0N r)+  assertEqualUpToEpsilon 1e-10+    ( sconcrete+      $ Nested.sfromListPrimLinear @_ @'[3, 4] knownShS [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , srepl 57.1 )+    (grad (kfromS . ssum0 . reluT2) (sconcrete $ Nested.sfromListPrimLinear @_ @'[3, 4] knownShS [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], srepl 7))++testReluSimplerPP7s2 :: Assertion+testReluSimplerPP7s2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)+      reluT2 (t, r) = smap0N (flip tlet $ \x -> ifH (x <=. sscalar 0) (sscalar 0.0) x) (t * sreplicate0N r)+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKS [3, 4] FTKScalar) (FTKS ZSS FTKScalar))+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let m9 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) ; t15 = sscatter dret (\\[i13, i14] -> [ifH (sscalar -0.0 <=. negate (m9 !$ [i13, i14])) 0 1, i13, i14]) ; m16 = t15 !$ [1] in tpair (sreplicate @3 (sreplicate @4 (tproject2 m1)) * m16) (ssum @4 (ssum @3 (tproject1 m1 * m16)))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> let m9 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) in sgather (sfromVector (fromList [sconcrete (sreplicate [3,4] 0.0), m9])) (\\[i10, i11] -> [ifH (sscalar -0.0 <=. negate (m9 !$ [i10, i11])) 0 1, i10, i11])"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> let m16 = sscatter dret (\\[i13, i14] -> [ifH (sscalar -0.0 <=. negate (tproject1 m1 !$ [i13, i14]) * tproject2 m1) 0 1, i13, i14]) !$ [1] in tpair (sreplicate @3 (sreplicate @4 (tproject2 m1)) * m16) (sdot0 (tproject1 m1) m16)"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> let m9 = tproject1 m1 * sreplicate @3 (sreplicate @4 (tproject2 m1)) in sgather (sfromVector (fromList [sconcrete (sreplicate [3,4] 0.0), m9])) (\\[i10, i11] -> [ifH (sscalar -0.0 <=. negate (m9 !$ [i10, i11])) 0 1, i10, i11])"++testReluSimpler7s :: Assertion+testReluSimpler7s = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double), AstTensor AstMethodLet FullSpan (TKS '[] Double))+             -> AstTensor AstMethodLet FullSpan (TKS '[3, 4] Double)+      reluT2 (t, r) = smap0N (flip tlet $ \x -> ifH (x <=. sscalar 0) (sscalar 0.0) x) (t * sreplicate0N r)+  assertEqualUpToEpsilon 1e-10+    ( sconcrete+      $ Nested.sfromListPrimLinear @_ @'[3, 4] knownShS [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , srepl 57.1 )+    (grad (kfromS . ssum0 . reluT2) (sconcrete $ Nested.sfromListPrimLinear @_ @'[3, 4] knownShS [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], srepl 7))++reluMax :: forall target n r. (ADReady target, GoodScalar r, KnownNat n)+        => target (TKR n r) -> target (TKR n r)+reluMax = rmap0N (maxH (rscalar 0))++testReluMax :: Assertion+testReluMax = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3, 4] [1.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,0.0,0.0,1.0,1.0])+    (rev' @Double @2 reluMax (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12]))++testReluMaxPP :: Assertion+testReluMaxPP = do+  resetVarCounter >> resetIdCounter+  let reluT :: AstTensor AstMethodLet FullSpan (TKR 2 Double) -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT = reluMax+      (var3, ast3) = funToAst (FTKR [3, 4] FTKScalar) Nothing reluT+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\m1 -> rfromS (sgather [] (tfromVector (SNat @2) (STKS [3,4] STKScalar) (fromList [tfromPrimal (STKS [3,4] STKScalar) (sconcrete (sreplicate [3,4] 0.0)), sfromR m1])) (\\[i5, i4] -> [ifH (sscalar -0.0 <=. negate (tprimalPart (sfromR m1) !$ [i5, i4])) 0 1, i5, i4]))"+  resetVarCounter+  let (artifactRev, deltas) = revArtifactDelta UseIncomingCotangent reluT (FTKR [3, 4] FTKScalar)+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> rfromS (sscatter (sfromR dret) (\\[i9, i10] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i9, i10])) 0 1, i9, i10]) !$ [1])"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (sgather (sfromVector (fromList [sconcrete (sreplicate [3,4] 0.0), sfromR m1])) (\\[i6, i7] -> [ifH (sscalar -0.0 <=. negate (sfromR m1 !$ [i6, i7])) 0 1, i6, i7]))"+  show deltas+    @?= "DeltaConvert (ConvCmp (ConvXR STKScalar) ConvSX) (DeltaShare 100000004 (DeltaGatherS [3,4] [] [2,3,4] (DeltaShare 100000001 (DeltaFromVector (SNat @2) (STKS [3,4] STKScalar) [DeltaZero (FTKS [3,4] FTKScalar),DeltaConvert (ConvCmp (ConvXS' (FTKS [3,4] FTKScalar)) ConvRX) (DeltaInput (InputId 0))])) <function>))"++testReluMaxPP2 :: Assertion+testReluMaxPP2 = do+  resetVarCounter+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 1 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 1 Double)+      reluT2 (t, r) = reluMax (t * rreplicate 5 r)+      (var3, ast3) = funToAst (FTKR [5] FTKScalar) Nothing (\t -> reluT2 (t, rscalar 7))+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> rfromS (sgather [] (tfromVector (SNat @2) (STKS [5] STKScalar) (fromList [tfromPrimal (STKS [5] STKScalar) (sconcrete (sreplicate [5] 0.0)), sfromR v1 * tfromPrimal (STKS [5] STKScalar) (sconcrete (sreplicate [5] 7.0))])) (\\[i3] -> [ifH (sscalar -0.0 <=. sscalar (-7.0) * tprimalPart (sfromR v1) !$ [i3]) 0 1, i3]))"+  resetVarCounter+  let (artifactRev, _deltas) = revArtifactDelta UseIncomingCotangent reluT2 (FTKProduct (FTKR [5] FTKScalar) (FTKR ZSR FTKScalar))+  printArtifactPretty artifactRev+    @?= "\\dret v1 -> let m11 = sscatter (sfromR dret) (\\[i8] -> [let x9 = negate (sfromR (tproject1 v1) !$ [i8]) ; x10 = sfromR (rreplicate 5 (tproject2 v1)) !$ [i8] in ifH (sscalar -0.0 <=. x9 * x10) 0 1, i8]) ; v12 = m11 !$ [1] in tpair (rfromS (sfromR (rreplicate 5 (tproject2 v1)) * v12)) (rsum (rfromS (sfromR (tproject1 v1) * v12)))"+  printArtifactPrimalPretty artifactRev+    @?= "\\v1 -> rfromS (sgather (sfromVector (fromList [sconcrete (sreplicate [5] 0.0), sfromR (tproject1 v1) * sfromR (rreplicate 5 (tproject2 v1))])) (\\[i4] -> [let x5 = negate (sfromR (tproject1 v1) !$ [i4]) ; x6 = sfromR (rreplicate 5 (tproject2 v1)) !$ [i4] in ifH (sscalar -0.0 <=. x5 * x6) 0 1, i4]))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret v1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [5] STKScalar) (STKS [] STKScalar)) (let v12 = sscatter (sfromR dret) (\\[i8] -> [ifH (sscalar -0.0 <=. negate (sfromR (tproject1 v1) !$ [i8]) * sfromR (tproject2 v1)) 0 1, i8]) !$ [1] in tpair (sreplicate @5 (sfromR (tproject2 v1)) * v12) (sdot0 (sfromR (tproject1 v1)) v12))"++testReluMax3 :: Assertion+testReluMax3 = do+  let reluT2 :: (AstTensor AstMethodLet FullSpan (TKR 2 Double), AstTensor AstMethodLet FullSpan (TKR 0 Double))+             -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      reluT2 (t, r) = reluMax (t * rreplicate 3 (rreplicate 4 r))+  assertEqualUpToEpsilon 1e-10+    ( ringestData [3, 4] [7.0,0.0,0.0,7.0,7.0,7.0,7.0,7.0,0.0,0.0,7.0,7.0]+    , rscalar 57.1 )+    (grad (kfromR . rsum0 . reluT2) (ringestData [3, 4] [1.1, -2.2, 0, 4.4, 5.5, 6.6, 7.7, 8.8, -9.9, -10, 11, 12], rscalar 7))++testDot1PP :: Assertion+testDot1PP = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent (uncurry (rdot0 @1 @Double))+                 (FTKProduct (FTKR [3] FTKScalar) (FTKR [3] FTKScalar))+  printArtifactPretty artifactRev+    @?= "\\dret v1 -> tpair (rfromS (sfromR (tproject2 v1) * sreplicate @3 (sfromR dret))) (rfromS (sfromR (tproject1 v1) * sreplicate @3 (sfromR dret)))"+  printArtifactPrimalPretty artifactRev+    @?= "\\v1 -> rfromS (ssum @3 (sfromR (tproject1 v1) * sfromR (tproject2 v1)))"++testDot2PP :: Assertion+testDot2PP = do+  resetVarCounter+  let (artifactRev, _deltas) =+        revArtifactDelta UseIncomingCotangent (uncurry (rdot0 @2 @Double))+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [2, 3] FTKScalar))+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let m3 = sreshape @[2,3] (sreplicate @6 (sfromR dret)) in tpair (rfromS (sfromR (tproject2 m1) * m3)) (rfromS (sfromR (tproject1 m1) * m3))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromS (ssum @6 (sreshape @[6] (sfromR (tproject1 m1) * sfromR (tproject2 m1))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,3] FTKScalar)) ConvSX))) (STKProduct (STKS [2,3] STKScalar) (STKS [2,3] STKScalar)) (tpair (sfromR (tproject2 m1) * sreplicate @2 (sreplicate @3 (sfromR dret))) (sfromR (tproject1 m1) * sreplicate @2 (sreplicate @3 (sfromR dret))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (sdot0 (sfromR (tproject1 m1)) (sfromR (tproject2 m1)))"++testMatvecmulPP :: Assertion+testMatvecmulPP = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry rmatvecmul)+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3] FTKScalar))+  printArtifactPretty @_ @(TKR 1 Double) artifactRev+    @?= "\\dret m1 -> tpair (rfromS (str (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret)))) (rfromS (ssum @2 (str (str (sfromR (tproject1 m1)) * sreplicate @3 (sfromR dret)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromS (ssum @3 (str (sreplicate @2 (sfromR (tproject2 m1))) * str (sfromR (tproject1 m1))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [3] FTKScalar)) ConvSX))) (STKProduct (STKS [2,3] STKScalar) (STKS [3] STKScalar)) (tpair (sreplicate @2 (sfromR (tproject2 m1)) * str (sreplicate @3 (sfromR dret))) (sdot1In (str (sfromR (tproject1 m1))) (sreplicate @3 (sfromR dret))))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (sdot1In (sreplicate @2 (sfromR (tproject2 m1))) (sfromR (tproject1 m1)))"++testMatmul2PP :: Assertion+testMatmul2PP = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry rmatmul2)+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3, 4] FTKScalar))+  printArtifactPretty @_ @(TKR 2 Double) artifactRev+    @?= "\\dret m1 -> tpair (rfromS (ssum @4 (stranspose @[2,1,0] (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret))))) (rfromS (ssum @2 (str (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * sreplicate @3 (sfromR dret)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromS (ssum @3 (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * str (sreplicate @2 (sfromR (tproject2 m1)))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [3,4] FTKScalar)) ConvSX))) (STKProduct (STKS [2,3] STKScalar) (STKS [3,4] STKScalar)) (tpair (smatmul2 (sfromR dret) (str (sfromR (tproject2 m1)))) (smatmul2 (str (sfromR (tproject1 m1))) (sfromR dret)))"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (smatmul2 (sfromR (tproject1 m1)) (sfromR (tproject2 m1)))"++testMatmul2FromMatvecmulPP :: Assertion+testMatmul2FromMatvecmulPP = do+  resetVarCounter+  let rmatmul2F :: (BaseTensor target, GoodScalar r)+                => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)+      rmatmul2F m1 m2 =+        rbuild1 (rwidth m1) (\i -> rmatvecmul (rtr m2) (m1 ! [i]))+      (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry rmatmul2F)+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3, 4] FTKScalar))+  printArtifactPretty @_ @(TKR 2 Double) artifactRev+    @?= "\\dret m1 -> tpair (rfromS (ssum @4 (stranspose @[2,1,0] (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret))))) (rfromS (ssum @2 (str (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * sreplicate @3 (sfromR dret)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromS (ssum @3 (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * str (sreplicate @2 (sfromR (tproject2 m1)))))"++testMatmul2PaperPP :: Assertion+testMatmul2PaperPP = do+  resetVarCounter+  let rmatmul2P :: (BaseTensor target, GoodScalar r)+                => target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)+      rmatmul2P a b =+        let k :$: m :$: _ = rshape a+            _ :$: n :$: _ = rshape b+        in rbuild1 k (\i ->+             rbuild1 n (\j ->+               rsum (rbuild1 m (\p -> a ! [i, p] * b ! [p, j]))))+      (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry rmatmul2P)+                 (FTKProduct (FTKR [2, 3] FTKScalar) (FTKR [3, 4] FTKScalar))+  printArtifactPretty @_ @(TKR 2 Double) artifactRev+    @?= "\\dret m1 -> tpair (rfromS (ssum @4 (stranspose @[2,1,0] (str (sreplicate @2 (sfromR (tproject2 m1))) * sreplicate @3 (sfromR dret))))) (rfromS (ssum @2 (str (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * sreplicate @3 (sfromR dret)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> rfromS (ssum @3 (stranspose @[2,1,0] (sreplicate @4 (sfromR (tproject1 m1))) * str (sreplicate @2 (sfromR (tproject2 m1)))))"++testMatmul2PPS :: Assertion+testMatmul2PPS = do+  resetVarCounter+  let (artifactRev, _) =+        revArtifactDelta+                 UseIncomingCotangent (uncurry smatmul2)+                 (FTKProduct (FTKS (SNat @2 :$$ SNat @3 :$$ ZSS) (FTKScalar @Float)) (FTKS (SNat @3 :$$ SNat @4 :$$ ZSS) (FTKScalar @Float)))+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> tpair (ssum @4 (stranspose @[2,1,0] (str (sreplicate @2 (tproject2 m1)) * sreplicate @3 dret))) (ssum @2 (str (stranspose @[2,1,0] (sreplicate @4 (tproject1 m1)) * sreplicate @3 dret)))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> ssum @3 (stranspose @[2,1,0] (sreplicate @4 (tproject1 m1)) * str (sreplicate @2 (tproject2 m1)))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> tpair (smatmul2 dret (str (tproject2 m1))) (smatmul2 (str (tproject1 m1)) dret)"+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> smatmul2 (tproject1 m1) (tproject2 m1)"++testAddSpeedBig :: Assertion+testAddSpeedBig =+  let m1 :: Concrete (TKR 2 Double)+      m1 = ringestData [1000,1000] [1 .. 1000000]+      m2 = m1 + m1+  in assertEqualUpToEpsilon 1e-10 m2 m2++testMatmul2SpeedBig :: Assertion+testMatmul2SpeedBig =+  let m1 :: Concrete (TKR 2 Double)+      m1 = ringestData [1000,1000] [1 .. 1000000]+      m2 = rmatmul2 m1 m1+  in assertEqualUpToEpsilon 1e-10 m2 m2++bar :: forall a. RealFloatH a => (a, a) -> a+bar (x, y) =+  let w = foo2 (x, y, x) * sin y+  in atan2H x w + y * w++barF :: forall a. RealFloatH a => (a, a) -> a+barF (x, y) =+  let w = foo2 (x, y, x) * sin y+  in atan2H x w + y * w++testBar :: Assertion+testBar =+  assertEqualUpToEpsilon 1e-9+    (rscalar 3.1435239435581166,rscalar (-1.1053869545195814))+    (cgrad (kfromR . bar @(ADVal Concrete (TKR 0 Double))) (rscalar 1.1, rscalar 2.2))++testBarS :: Assertion+testBarS =+  assertEqualUpToEpsilon 1e-9+    (srepl 3.1435239435581166, srepl (-1.1053869545195814))+    (cgrad (kfromS . barF @(ADVal Concrete (TKS '[] Double))) (srepl 1.1, srepl 2.2))++testBar2S :: Assertion+testBar2S =+  assertEqualUpToEpsilon 1e-9+    (srepl 3.1435239435581166, srepl (-1.1053869545195814))+    (grad (kfromS . ssum0 . barF @(AstTensor AstMethodLet FullSpan (TKS '[52, 2, 2, 1, 1, 3] Double))) (srepl 1.1, srepl 2.2))++testBarCFwd :: Assertion+testBarCFwd =+  assertEqualUpToEpsilon 1e-9+    (rscalar 9.327500345189534e-2)+    (cjvp (bar @(ADVal Concrete (TKR 0 Double))) (rscalar 1.1, rscalar 2.2) (rscalar 0.1, rscalar 0.2))++testBarFwd :: Assertion+testBarFwd =+  assertEqualUpToEpsilon 1e-9+    (rscalar 9.327500345189534e-2)+    (jvp @_ @(TKR 0 Double)+         bar (rscalar 1.1, rscalar 2.2) (rscalar 0.1, rscalar 0.2))++barADVal2 :: forall a. RealFloatH a+          => (a, a, a) -> a+barADVal2 (x,y,z) =+  let w = foo2 (x,y,z) * sin y+  in atan2H z w + z * w++-- A check if gradient computation is re-entrant.+-- This test fails (not on first run, but on repetition) if old terms,+-- from before current iteration of gradient descent, are not soundly+-- handled in new computations (due to naive impurity, TH, plugins,+-- or just following the papers that assume a single isolated run).+-- This example requires monomorphic types and is contrived,+-- but GHC does optimize and factor out some accidentally constant+-- subterms in real examples (presumably after it monomorphizes them)+-- causing exactly the same danger.+-- This example also tests unused parameters (x), another common cause+-- of crashes in naive gradient computing code.+baz :: ( ADVal Concrete (TKR 0 Double)+       , ADVal Concrete (TKR 0 Double)+       , ADVal Concrete (TKR 0 Double) )+    -> ADVal Concrete (TKR 0 Double)+baz (_x,y,z) =+  let w = fooFromPrimal * barADVal2 (y,y,z) * sin y+  in atan2H z w + z * w++-- An "old term", computed once, stored at top level.+fooFromPrimal :: ADVal Concrete (TKR 0 Double)+fooFromPrimal = foo2 (rscalar 7, rscalar 8, rscalar 9)++testBaz :: Assertion+testBaz =+  assertEqualUpToEpsilon 1e-9+    (rscalar 0, rscalar (-5219.20995030263), rscalar 2782.276274462047)+    (cgrad (kfromR . baz) (rscalar 1.1, rscalar 2.2, rscalar 3.3))++-- If terms are numbered and @z@ is, wrongly, decorated with number 0,+-- here @fooFromPrimal@ is likely to clash with @z@, since it was numbered+-- starting from 0, too.+-- The test fails if term numbering is reset to 0 between gradient computation+-- runs (verified) as well as when delta-expression evaluation assumes+-- it can only encounter terms numbered in the current run and,+-- e.g., allocates an array with only that much space (not verified).+-- Actually, with term counter zeroed just before @cgrad@ application,+-- even a single @testBaz@ execution produces wrong results, but this test+-- is likely to fail in less naive implementations, as well.+testBazRenumbered :: Assertion+testBazRenumbered =+  assertEqualUpToEpsilon 1e-9+    (rscalar 0, rscalar (-5219.20995030263), rscalar 2783.276274462047)+    (cgrad (kfromR . (\(x,y,z) -> z + baz (x,y,z))) (rscalar 1.1, rscalar 2.2, rscalar 3.3))++-- A dual-number and list-based version of a function that goes+-- from `R^3` to `R`.+fooD :: forall r. r ~ Double+     => ListR 3 (ADVal Concrete (TKR 0 r)) -> ADVal Concrete (TKR 0 r)+fooD (x ::: y ::: z ::: ZR) =+  let w = x * sin y+  in atan2H z w + z * w++testFooD :: Assertion+testFooD =+  assertEqualUpToEpsilon 1e-10+    (fromList [rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627])+    (cgrad (kfromR . fooD) (fromList [rscalar 1.1, rscalar 2.2, rscalar 3.3]))++fooBuild1 :: (ADReady target, GoodScalar r, Differentiable r)+          => target (TKR 1 r) -> target (TKR 1 r)+fooBuild1 v =+  let r = rsum0 v+      v' = rminimum v+  in rbuild1 3 $ \ix ->+       r * foo2 ( rscalar 3+               , rscalar 5 * r+               , r * rminimum v * v')+       + bar (r, rindex v [ix + 1])++testFooBuildDt :: Assertion+testFooBuildDt =+  assertEqualUpToEpsilon 1e-5+    (rconcrete $ Nested.rfromListPrimLinear [4] [-189890.46351219364,-233886.08744601303,-222532.22669716467,-206108.68889329425])+    (vjp @_ @(TKR 1 Double)+           fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]) (rreplicate0N [3] (rscalar 42)))++testFooBuildCFwd :: Assertion+testFooBuildCFwd =+  assertEqualUpToEpsilon 1e-5+    (rconcrete $ Nested.rfromListPrimLinear [3] [-296584.8166864211,-290062.472288043,-265770.1775742018])+    (cjvp @_ @(TKR 1 Double)+          fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]) (rreplicate0N [4] (rscalar 42)))++testFooBuildFwd :: Assertion+testFooBuildFwd =+  assertEqualUpToEpsilon 1e-5+    (rconcrete $ Nested.rfromListPrimLinear [3] [-296584.8166864211,-290062.472288043,-265770.1775742018])+    (jvp @_ @(TKR 1 Double)+         fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]) (rreplicate0N [4] (rscalar 42)))++testFooBuild :: Assertion+testFooBuild =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [4] [-4521.201512195087,-5568.7163677622175,-5298.386349932494,-4907.349735554627])+    (rev' @Double @1 fooBuild1 (ringestData [4] [1.1, 2.2, 3.3, 4]))++fooNoGo :: forall target r. (ADReady target, GoodScalar r, Differentiable r)+        => target (TKR 1 r) -> target (TKR 1 r)+fooNoGo v =+  let r = rsum0 v+  in rbuild1 3 (\ix' -> let ix :: PrimalOf target (TKS '[] Int64)+                            ix = sfromR $ rfromK ix' in+       bar (rscalar 3.14, bar (rscalar 3.14, rindex v [kfromS $ (ix + (sprimalPart . sfloor . sfromR) r) `minH` sscalar 2 `maxH` sscalar 0]))+       + ifH ((&&*) (rindex @1 v [ix' * 2] <=. rscalar 0)+                    (sscalar 6 >. abs ix))+               r (rscalar 5 * r))+     / rslice 1 3 (rmap0N (\x -> ifH (x >. r) r x) v)+     * rbuild1 3 (const (rscalar 1))++testFooNoGo0 :: Assertion+testFooNoGo0 =+  assertEqualUpToEpsilon' 1e-6+   (ringestData [5] [5.037878787878788,-14.394255484765257,43.23648655081373,-0.8403418295960368,5.037878787878788])+   (rev' @Double @1 fooNoGo+         (ringestData [5] [1.1 :: Double, 2.2, 3.3, 4, 5]))++nestedBuildMap :: forall target r.+                  (ADReady target, GoodScalar r, Differentiable r)+               => target (TKR 0 r) -> target (TKR 1 r)+nestedBuildMap r =+  let w = rreplicate0N [4]+      v0' = rreplicate0N [177] r :: target (TKR 1 r)+  in tlet v0' $ \v' ->+    let nestedMap x0 = tlet x0 $ \x -> rmap0N (x /) (w x)+        variableLengthBuild iy = rbuild1 7 (\ix -> rindex v' (ix + iy :.: ZIR))+        doublyBuild = rbuild1 5 (rminimum . variableLengthBuild)+        fooMap1 :: target (TKR 0 r) -> target (TKR 1 r)+        fooMap1 r2 =+          let v = fooBuild1 $ rreplicate0N [130] r2+          in rmap0N (\x -> x * r2 + rscalar 5) v+    in rmap0N (\x0 -> tlet x0 $ \x -> x * rsum0+                           (rbuild1 3 (\ix -> bar (x, rindex v' [ix]))+                            + (tlet (nestedMap x) $ \x3 -> fooBuild1 x3)+                            / (tlet x $ \x4 -> fooMap1 x4))+              ) doublyBuild++testNestedBuildMap1 :: Assertion+testNestedBuildMap1 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 107.25984443006627)+    (rev' @Double @1 nestedBuildMap (rscalar 1.1))++nestedSumBuild :: (ADReady target, GoodScalar r, Differentiable r)+               => target (TKR 1 r) -> target (TKR 1 r)+nestedSumBuild v0 = tlet v0 $ \v ->+ tlet (rsum (rbuild1 9 rfromIndex0)) (\tbtf ->+  (rbuild1 13 (\ix ->+    rsum (rbuild1 4 (\ix2 ->+      flip rindex [ix2]+        (tlet (rsum v) $ \tsumv -> rbuild1 5 (const tsumv)+         * rfromList+             [ rfromIndex0 ix+             , tbtf+             , rsum (rbuild1 6 (\_ -> rsum v))+             , rindex v [ix2]+             , rsum (rbuild1 3 (\ix7 ->+                 rsum (rreplicate 5 (rfromIndex0 ix7))))+             ]))))))+ + tlet (nestedBuildMap (rsum0 v)) (\nbmt -> (rbuild1 13 (\ix ->+     nbmt `rindex` [minH ix 4])))++testNestedSumBuild :: Assertion+testNestedSumBuild =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [5] [-14084.715065313612,-14084.715065313612,-14084.715065313612,-14014.775065313623,-14084.715065313612])+    (rev' @Double @1 nestedSumBuild (ringestData [5] [1.1, 2.2, 3.3, 4, -5.22]))++nestedBuildIndex :: forall target r. (ADReady target, GoodScalar r)+                 => target (TKR 1 r) -> target (TKR 1 r)+nestedBuildIndex v =+  rbuild1 2 $ \ix2 -> rindex (rbuild1 3 $ \ix3 -> rindex (rbuild1 4 $ \ix4 -> rindex v (ix4 :.: ZIR)) [ix3]) (ix2 :.: ZIR)++testNestedBuildIndex :: Assertion+testNestedBuildIndex =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [5] [1,1,0,0,0])+    (rev' @Double @1 nestedBuildIndex (ringestData [5] [1.1, 2.2, 3.3, 4, -5.22]))++barRelu+  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )+  => target (TKR n r) -> target (TKR n r)+barRelu x = relu $ bar (x, relu x)++testBarReluDt :: Assertion+testBarReluDt =+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [] [191.20462646925841])+    (vjp @_ @(TKR 0 Double)+           barRelu (rscalar 1.1) (rscalar 42.2))++testBarRelu :: Assertion+testBarRelu =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [] [4.5309153191767395])+    (rev' @Double @0 barRelu (ringestData [] [1.1]))++testBarRelu3 :: Assertion+testBarRelu3 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 1, 2] [4.5309153191767395,4.5302138998556,-9.39547533946234,95.29759282497125])+    (rev' @Double @3 barRelu (ringestData [2, 1, 2] [1.1, 2, 3, 4.2]))++barReluMax0+  :: ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+  => target (TKR n r) -> target (TKR n r)+barReluMax0 x = reluMax $ bar (x, x)++barReluMax+  :: ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+  => target (TKR n r) -> target (TKR n r)+barReluMax x = reluMax $ bar (x, reluMax x)++testBarReluMaxDt :: Assertion+testBarReluMaxDt =+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [] [191.20462646925841])+    (vjp @_ @(TKR 0 Double)+           barReluMax (rfromList0N [] [rscalar 1.1]) (rscalar 42.2))++testBarReluMax :: Assertion+testBarReluMax =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [] [4.5309153191767395])+    (rev' @Double @0 barReluMax (ringestData [] [1.1]))++testBarReluMax30 :: Assertion+testBarReluMax30 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [4.5309153191767395])+    (rev' @Double @1 barReluMax0 (ringestData [1] [1.1]))++testBarReluMax31 :: Assertion+testBarReluMax31 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 1, 2] [4.5309153191767395,4.5302138998556,-9.39547533946234,95.29759282497125])+    (rev' @Double @3 barReluMax (ringestData [2, 1, 2] [1.1, 2, 3, 4.2]))++testBarReluMax3CFwd :: Assertion+testBarReluMax3CFwd =+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2, 1, 2] [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])+    (cjvp @_ @(TKR 3 Double)+          barReluMax+                     (rconcrete $ Nested.rfromListPrimLinear (fromList [2, 1, 2]) [1.1, 2, 3, 4.2])+                     (ringestData [2, 1, 2] [0.1, 0.2, 0.3, 0.42]))++reluMaxS :: forall target sh r. (ADReady target, GoodScalar r, KnownShS sh)+         => target (TKS sh r) -> target (TKS sh r)+reluMaxS = smap0N (maxH (srepl 0))++barReluMaxS+  :: ( ADReady target, GoodScalar r, KnownShS sh+     , RealFloatH (target (TKS sh r)) )+  => target (TKS sh r) -> target (TKS sh r)+barReluMaxS x = reluMaxS $ barF (x, reluMaxS x)++-- Previously the shape of FromListR[DeltaZero] couldn't be determined+-- in buildDerivative, so this was needed. See below that it now works fine.+testBarReluMax3FwdS :: Assertion+testBarReluMax3FwdS =+  assertEqualUpToEpsilon 1e-10+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])+    (jvp @_ @(TKS '[2, 1, 2] Double)+         barReluMaxS+         (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [1.1, 2, 3, 4.2])+         (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [0.1, 0.2, 0.3, 0.42]))++testBarReluMax3FwdFrom :: Assertion+testBarReluMax3FwdFrom =+  assertEqualUpToEpsilon 1e-10+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])+    (jvp @_ @(TKS '[2, 1, 2] Double)+         (sfromR . barReluMax . rfromS)+         (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [1.1, 2, 3, 4.2])+         (sconcrete $ Nested.sfromListPrimLinear @_ @'[2, 1, 2] knownShS [0.1, 0.2, 0.3, 0.42]))++testBarReluMax3FwdR :: Assertion+testBarReluMax3FwdR =+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2, 1, 2] [0.45309153191767404,0.9060427799711201,-2.8186426018387007,40.02498898648793])+    (jvp @_ @(TKR 3 Double)+         barReluMax+         (ringestData [2, 1, 2] [1.1, 2, 3, 4.2])+         (ringestData [2, 1, 2] [0.1, 0.2, 0.3, 0.42]))++f1 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 0 r)+f1 = \arg -> rsum0 (rbuild1 10 (\i -> arg * rfromIndex0 i))++testF1 :: Assertion+testF1 =+  assertEqualUpToEpsilon 1e-10+    (rscalar 45.0)+    (grad (kfromR @_ @Double . f1) (rscalar 1.1))++testF11 :: Assertion+testF11 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 45.0)+    (rev' @Double @0 f1 (rscalar 1.1))++f2 :: forall target r. (ADReady target, GoodScalar r)+   => target (TKR 0 r) -> target (TKR 0 r)+f2 = \arg ->+  let fun1 i = arg * rfromIndex0 i+      v1a = rsum0 (rbuild1 10 fun1)+      v1b = rsum0 (rbuild1 20 fun1)+      fun2 y i = y * rfromIndex0 @r i+      v2a = rsum0 (rbuild1 10 (fun2 arg))+      v2b = rsum0 (rbuild1 20 (fun2 (arg + rscalar 1)))+  in v1a + v1b + v2a + v2b++testF2 :: Assertion+testF2 =+  assertEqualUpToEpsilon 1e-10+    (rscalar 470)+    (grad (kfromR @_ @Double . f2) (rscalar 1.1))++testF21 :: Assertion+testF21 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 470)+    (rev' @Double @0 f2 (rscalar 1.1))++testF2CFwd :: Assertion+testF2CFwd =+  assertEqualUpToEpsilon 1e-10+    (rscalar 47)+    (cjvp @_ @(TKR 0 Double)+          f2 (rscalar 1.1) (rscalar 0.1))++testF2Fwd :: Assertion+testF2Fwd =+  assertEqualUpToEpsilon 1e-10+    (rscalar 47)+    (jvp @_ @(TKR 0 Double)+         f2 (rscalar 1.1) (rscalar 0.1))++braidedBuilds :: forall target r.+                 (ADReady target, GoodScalar r, Differentiable r)+              => target (TKR 0 r) -> target (TKR 2 r)+braidedBuilds r =+  rbuild1 3 (\ix1 ->+    rbuild1 4 (\ix2 -> rindex (rfromList0N [4]+      [rfromIndex0 ix2, rscalar 7, r, rscalar (-0.2)]) (ix1 :.: ZIR)))++testBraidedBuilds1 :: Assertion+testBraidedBuilds1 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 4.0)+    (rev' @Double @2 braidedBuilds (rscalar 3.4))++recycled :: (ADReady target, GoodScalar r, Differentiable r)+         => target (TKR 0 r) -> target (TKR 5 r)+recycled r =+  tlet (nestedSumBuild (rreplicate0N [7] r)) $ \nsb ->+    rbuild1 2 $ \_ -> rbuild1 4 $ \_ -> rbuild1 2 $ \_ -> rbuild1 3 $ const nsb++testRecycled1 :: Assertion+testRecycled1 =+  assertEqualUpToEpsilon 1e-6+    (rscalar 348356.9278600814)+    (grad (kfromR @_ @Double . rsum0 . recycled) (rscalar 0.0000001))++concatBuild :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 2 r)+concatBuild r =+  rbuild1 7 (\i ->+    rappend (rappend (rbuild1 5 (const r))+                     (rreplicate 1 (rfromIndex0 i)))+            (rbuild1 13 (const r)))++testConcatBuild1 :: Assertion+testConcatBuild1 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 126.0)+    (rev' @Double @2 concatBuild (rscalar 3.4))++concatBuild2 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild2 r =+  tlet (rfromList [r, rscalar 1, rscalar 2, rscalar 3, rscalar 4]) $ \a ->+    rbuild1 10 (\i -> ifH (i <. 5) (rindex a [i]) (rindex a [i - 5]))++testConcatBuild2 :: Assertion+testConcatBuild2 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2.0)+    (rev' @Double @1 concatBuild2 (rscalar 3.4))++concatBuild3 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild3 r =+  tlet (rfromList [r, rscalar 1, rscalar 2, rscalar 3, rscalar 4]) $ \a ->+    rbuild1 10 (\i -> ifH (i <. 5) (rindex a [i]) (rindex a [i - 5 + (1 `quotH` maxH 1 (i - 5))]))++testConcatBuild3 :: Assertion+testConcatBuild3 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1)+    (rev' @Double @1 concatBuild3 (rscalar 3.4))++concatBuild4 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild4 r =+  tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (4 + i)) :.: ZIR)) $ \a ->+    rappend a a++testConcatBuild4 :: Assertion+testConcatBuild4 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 10)+    (rev' @Double @1 concatBuild4 (rscalar 3.4))++concatBuild5 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild5 r =+  tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (5 + i)) :.: ZIR)) $ \a ->+    (rappend a (rgather1 5 (rreplicate 1 r)+                           (\i -> (1 `quotH` (5 + i)) :.: ZIR)))++testConcatBuild5 :: Assertion+testConcatBuild5 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 10)+    (rev' @Double @1 concatBuild5 (rscalar 3.4))++concatBuild6 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 2 r)+concatBuild6 r =+  rbuild1 7 (\j ->+    rappend (rappend+             (tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (4 + i)) :.: ZIR)) $ \a ->+    (rappend (rgather1 5 (rreplicate 1 r)+                         (\i -> (1 `quotH` (100 * maxH 1 (i - j))) :.: ZIR)) a))+                     (rreplicate 1 (rfromIndex0 j)))+            (rbuild1 13 (const r)))++testConcatBuild6 :: Assertion+testConcatBuild6 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 161)+    (rev' @Double @2 concatBuild6 (rscalar 3.4))++concatBuild7 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+concatBuild7 r =+  rbuild1 10 $ \j ->+    (rappend (rreplicate 5 r) (rgather1 5 (rreplicate 1 r)+                                 (\i -> (1 `quotH` maxH 1 (j - i)) :.: ZIR)))+     ! (j :.: ZIR)++testConcatBuild7 :: Assertion+testConcatBuild7 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 10)+    (rev' @Double @1 concatBuild7 (rscalar 3.4))++-- These tests show that term rewriting changes the value+-- of out-of-bounds indexing, e.g., of gather(replicate)+-- that reduces to a non-gather. Vectorization is not needed for that.+_concatBuild8 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 1 r)+_concatBuild8 r =+  tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (5 - i)) :.: ZIR)) $ \a ->+    (rappend a (rgather1 5 (rreplicate 1 r)+                           (\i -> (1 `quotH` (5 - i)) :.: ZIR)))++_testConcatBuild8 :: Assertion+_testConcatBuild8 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 10)+    (rev' @Double @1 _concatBuild8 (rscalar 3.4))++_concatBuild9 :: (ADReady target, GoodScalar r) => target (TKR 0 r) -> target (TKR 2 r)+_concatBuild9 r =+  rbuild1 7 (\j ->+    rappend (rappend+             (tlet (rgather1 5 (rreplicate 1 r)+                   (\i -> (1 `quotH` (4 - i)) :.: ZIR)) $ \a ->+    (rappend (rgather1 5 (rreplicate 1 r)+                         (\i -> (1 `quotH` (100 * maxH 0 (i - j))) :.: ZIR)) a))+                     (rreplicate 1 (rfromIndex0 j)))+            (rbuild1 13 (const r)))++_testConcatBuild9 :: Assertion+_testConcatBuild9 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 161)+    (rev' @Double @2 _concatBuild9 (rscalar 3.4))++concatBuild10 :: (ADReady target, GoodScalar r)+              => target (TKR 0 r) -> target (TKR 2 r)+concatBuild10 r =+  rbuild1 7 (\j ->+    rappend (rappend+               (tlet (rgather1 5 (rreplicate 1 r)+                       (\_i -> 10000 :.: ZIR)) $ \a ->+      (rappend (rgather1 5 (rreplicate 1 r)+                 (\_i -> (-1) :.: ZIR)) a))+                   (rreplicate 1 (rfromIndex0 j)))+            (rbuild1 13 (const r)))++testConcatBuild10 :: Assertion+testConcatBuild10 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 91)+    (rev' @Double @2 concatBuild10 (rscalar 3.4))++concatBuild11 :: (ADReady target, GoodScalar r)+              => target (TKR 0 r) -> target (TKR 1 r)+concatBuild11 r =+  rgather1 5 (rreplicate 1 r) (\_i -> (-1) :.: ZIR)++testConcatBuild11 :: Assertion+testConcatBuild11 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @1 concatBuild11 (rscalar 3.4))++concatBuild12 :: (ADReady target, GoodScalar r)+              => target (TKR 0 r) -> target (TKR 0 r)+concatBuild12 r =+  rindex (rreplicate 1 r) ((-1) :.: ZIR)++testConcatBuild12 :: Assertion+testConcatBuild12 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @0 concatBuild12 (rscalar 3.4))++-- TODO: copy-paste a variant of emptyArgs with r not GoodScalar+-- or maybe generalize emptyArgs and then in half of the tests do TKScalar+emptyArgs :: forall target r. (ADReady target, GoodScalar r, Differentiable r)+          => target (TKR 1 r) -> target (TKR 1 r)+emptyArgs t =+  emptyTensor+  - rfromList0N (rshape @target @_ @(TKScalar r) emptyTensor) []+  + rconcrete Nested.remptyArray+  - rsum (rfromList0N (0 :$: rshape @target @_ @(TKScalar r) emptyTensor) [])+  * rsum (rconcrete $ Nested.rreshape (0 :$: 0 :$: ZSR) Nested.remptyArray)+  * rconcrete (Nested.rsumOuter1 $ Nested.rfromListOuter+               $ NonEmpty.fromList [Nested.remptyArray])+  * rsum (rconcrete $ Nested.rfromListOuter+          $ NonEmpty.fromList [Nested.remptyArray])+  * rsum (rfromList [rconcrete Nested.remptyArray])+  * rsum (rfromList [emptyTensor])+  * rsum (rfromList [rsum (rfromList0N+                             (0 :$: rshape @target @_ @(TKScalar r)+                                      emptyTensor) [])])+  - rindex (rfromList0N (0 :$: 0 :$: ZSR) []) (42 :.: ZIR)+  - rindex (rfromList0N (0 :$: rshape @target @_ @(TKScalar r)+                                     emptyTensor) []) (42 :.: ZIR)+  - rreshape @1 [0] emptyTensor+  - rsum (rreshape @1 [0, 0] emptyTensor)+-- TODO:+--  - rgather1 0 emptyTensor (:.: ZIR)+--  - rsum (rgather1 0 emptyTensor (const ZIR))+--  - rsum (rgather @target @(TKScalar r) @2 (0 :$: 0 :$: ZSR) emptyTensor (const (0 :.: ZIR)))+--  - rsum (rgather @target @(TKScalar r) @2 @0 @1 [0, 0] emptyTensor (const [0]))+--  * rflatten (rtr (rgather1 0 t (const ZIR)))+  + rbuild1 0 (\i -> t ! (i :.: ZIR))+  + rbuild1 0 (\i -> t ! [fromIntegral (rlength t) `quotH` i] / rfromIndex0 i)+  + rbuild @1 (0 :$: ZSR) (const $ rscalar 73)+  - rsum (rbuild @0 (0 :$: 0 :$: ZSR)+                 (const (rreplicate 1 emptyTensor)))+  + rfromS+      (sfromList0N []+       + sconcrete (Nested.semptyArray ZSS)+       - ssum @0 (sfromList0N [])+       * ssum @0 (sconcrete $ Nested.semptyArray (SNat @0 :$$ ZSS))+       * sconcrete (Nested.ssumOuter1 $ Nested.sfromListOuter (SNat @1)+                    $ NonEmpty.fromList [Nested.semptyArray ZSS])+       * ssum @1 (sconcrete $ Nested.sfromListOuter SNat+                        $ NonEmpty.fromList [Nested.semptyArray ZSS])+       * ssum @1 (sfromList [sconcrete $ Nested.semptyArray ZSS])+       * ssum @1 (sfromList [sfromR @_ @'[0] emptyTensor])+       * ssum @1 (sfromList [ssum @0 (sfromList0N [])])+       - sindex @'[0] (sfromList0N []) (42 :.$ ZIS)+       - sindex @'[0] (sfromList0N []) (42 :.$ ZIS)+       - sreshape @_ @'[0] (sfromR @_ @'[0] emptyTensor)+       - ssum (sreshape @_ @'[0, 0] (sfromR @_ @'[0] emptyTensor))+       * sbuild1 @0 (\i -> sfromR @_ @'[0] (rslice 0 0 t) !$ (i :.$ ZIS))+       + sbuild1 @0 (\i -> sfromR @_ @'[0] (rslice 0 0 t)+                              !$ (fromIntegral (rlength t) `quotH` i :.$ ZIS)+                              / sfromIndex0 i)+       + sbuild @1 (const $ sscalar 73)+       - ssum (sbuild @0+                      (const (sreplicate @1 (sfromR emptyTensor)))))+  + rfromX+      (xfromList0N (SKnown (SNat @0) :$% ZSX) []+       + xconcrete (Nested.memptyArray ZSX)+       - xsum @0 (xfromList0N+                          (SKnown (SNat @0) :$% SKnown (SNat @0) :$% ZSX) [])+       * xsum @0 (xconcrete+                        $ Nested.memptyArray (SKnown (SNat @0) :$% ZSX))+       * xconcrete (Nested.msumOuter1 $ Nested.mfromListOuter+                    $ NonEmpty.fromList [Nested.memptyArray ZSX])+       * xsum @1 (xconcrete+                        $ Nested.mcast+                            (SKnown (SNat @1) :!% SKnown (SNat @0) :!% ZKX)+                        $ Nested.mfromListOuter+                        $ NonEmpty.fromList [Nested.memptyArray ZSX])+       * xsum @1 (xfromList [xconcrete $ Nested.memptyArray ZSX])+       * xsum @1 (xfromList [xfromR @_ @'[Just 0] emptyTensor])+       * xsum @1 (xfromList [xsum @0+                                     (xfromList0N+                                        (SKnown (SNat @0)+                                         :$% xshape @target @_ @(TKScalar r)+                                               (xfromR @_ @'[Just 0]+                                                  emptyTensor)) [])])+       - xindex @_ @'[Just 0] (xfromList0N+                                    (SKnown (SNat @0)+                                     :$% SKnown (SNat @0)+                                     :$% ZSX) []) (42 :.% ZIX)+       - xindex @_ @'[Just 0] (xfromList0N+                                    (SKnown (SNat @0)+                                     :$% xshape @target @_ @(TKScalar r)+                                             (xfromR @_ @'[Just 0]+                                                emptyTensor)) []) (42 :.% ZIX)+       - xreshape (SKnown (SNat @0) :$% ZSX) (xfromR @_ @'[Just 0] emptyTensor)+       - xsum (xreshape (SKnown (SNat @0) :$% SKnown (SNat @0) :$% ZSX) (xfromR @_ @'[Just 0] emptyTensor))+       * xbuild1 @0 (\i -> xfromR @_ @'[Nothing] (rslice 0 0 t)+                            `xindex` (i :.% ZIX))+       + xbuild1 @0 (\i -> xfromR @_ @'[Nothing] (rslice 0 0 t)+                              `xindex`+                              (fromIntegral (rlength t) `quotH` i :.% ZIX)+                              / xfromIndex0 i)+       + xbuild @1 (SKnown (SNat @0) :$% ZSX)+                (const $ xscalar 73)+       - xsum (xbuild @0 (SKnown (SNat @0)+                                                :$% SKnown (SNat @0)+                                                :$% ZSX)+                      (const (xreplicate (xfromR emptyTensor)))))+  -- - rsum (rbuild @2 (0 :$: 0 :$: ZSR) (const 73))+  -- - rsum (rbuild @1 (0 :$: 0 :$: ZSR) (const emptyTensor))+       -- these two fail and rightly so; TODO: make them fail earlier+ where+  emptyTensor :: target (TKR 1 r)+  emptyTensor = rconcrete $ Nested.rfromListPrimLinear (fromList [0]) []++testEmptyArgs0 :: Assertion+testEmptyArgs0 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [0] [])+    (rev' @Double @1 emptyArgs (ringestData [0] []))++testEmptyArgs1 :: Assertion+testEmptyArgs1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [0])+    (rev' @Float @1 emptyArgs (ringestData [1] [0.24]))++testEmptyArgs4 :: Assertion+testEmptyArgs4 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1] [0])+    (rev' @Double @1+          (\t -> rbuild [2, 5, 11, 0] (const $ emptyArgs t))+          (ringestData [1] [0.24]))++filterPositiveFail :: ADReady target+                   => target (TKR 1 Double) -> target (TKR 1 Double)+filterPositiveFail v =+  let l = runravelToList v+      -- l2 = filter (\x -> x >= 0) l+      -- Could not deduce ‘Ord (target Double 0)’+      -- l2 = filter (\x -> x >=. 0) l+      -- Could not deduce ‘BoolOf target ~ Bool’+      l2 = take 3 l  -- the most I can do+  in rfromList $ NonEmpty.fromList l2++testFilterPositiveFail :: Assertion+testFilterPositiveFail =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [5] [1.0,1.0,1.0,0.0,0.0])+    (rev' @Double @1+          filterPositiveFail+          (ringestData [5] [0.24, 52, -0.5, 0.33, 0.1]))++-- Catastrophic loss of sharing prevented.+fblowup :: forall target r. (ADReady target, GoodScalar r, Differentiable r)+        => Int -> target (TKR 1 r) -> target (TKR 0 r)+fblowup k inputs =+  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)+      blowup 0 y = y - rfromIndex0 0+      blowup n y =+        let ysum = y + y - rfromIndex0 0+            yscaled = rscalar 0.499999985 * ysum+              -- without the scaling we'd get NaN at once+        in blowup (pred n) yscaled+      y0 = (inputs ! [0]) / (inputs ! [1])+  in blowup k y0++fblowupLet :: forall target r. (ADReady target, GoodScalar r, Differentiable r)+           => IntOf target -> Int -> target (TKR 1 r) -> target (TKR 0 r)+fblowupLet i k inputs =+  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)+      blowup 0 y = y - rfromIndex0 i+      blowup n y1 = tlet y1 $ \y ->+        let ysum = y + y - rfromIndex0 i+            yscaled = rscalar 0.499999985 * ysum+              -- without the scaling we'd get NaN at once+        in blowup (pred n) yscaled+      y0 = (inputs ! [0]) / (inputs ! [1])+  in blowup k y0++-- Catastrophic loss of sharing prevented also with non-trivial multiplication.+fblowupMult :: forall target r. (ADReady target, GoodScalar r, Differentiable r)+            => Int -> target (TKR 1 r) -> target (TKR 0 r)+fblowupMult k inputs =+  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)+      blowup 0 y = y+      blowup n y =+        let ysum = y + y * y / (y - rscalar 0.000000001)+            yscaled = sqrt $ rscalar 0.499999985 * rscalar 0.499999985 * ysum * ysum+              -- without the scaling we'd get NaN at once+        in blowup (pred n) yscaled - rfromIndex0 0+      y0 = (inputs ! [0 `remH` 2]) * (inputs ! [1])+  in blowup k y0++fblowupMultLet :: forall target r.+                  (ADReady target, GoodScalar r, Differentiable r)+               => IntOf target -> Int -> target (TKR 1 r) -> target (TKR 0 r)+fblowupMultLet i k inputs =+  let blowup :: Int -> target (TKR 0 r) -> target (TKR 0 r)+      blowup 0 y = y+      blowup n y1 = tlet y1 $ \y ->+        let ysum0 = y + y * y / (y - rscalar 0.000001)+            yscaled = tlet ysum0 $ \ysum ->+                        sqrt $ rscalar 0.499999985 * rscalar 0.499999985 * ysum * ysum+              -- without the scaling we'd get NaN at once+        in blowup (pred n) yscaled - rfromIndex0 i+      y0 = (inputs ! [i `remH` 2]) * (inputs ! [1])+  in blowup k y0++fblowupPP :: Assertion+fblowupPP = do+  resetVarCounter+  let fblowupT = fblowup @(AstTensor AstMethodLet FullSpan) @Double 1+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupT (FTKR [4] FTKScalar)+  printArtifactSimple artifactRev+    @?= "\\dret v1 -> tlet (sfromR v1 !$ [0]) (\\x2 -> tlet (sfromR v1 !$ [1]) (\\x3 -> tlet (sfromR v1 !$ [0]) (\\x4 -> tlet (sfromR v1 !$ [1]) (\\x5 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x8 -> rfromS (soneHot (recip x3 * x8) [0] + (soneHot ((negate x2 / (x3 * x3)) * x8) [1] + (soneHot (recip x5 * x8) [0] + soneHot ((negate x4 / (x5 * x5)) * x8) [1]))))))))"+  printArtifactPrimalSimple artifactRev+    @?= "\\v1 -> tlet (sfromR v1 !$ [0]) (\\x2 -> tlet (sfromR v1 !$ [1]) (\\x3 -> tlet (sfromR v1 !$ [0]) (\\x4 -> tlet (sfromR v1 !$ [1]) (\\x5 -> tlet ((x2 / x3 + x4 / x5) + negate (sfromIntegral (sscalar 0))) (\\x6 -> rfromS (sscalar 0.499999985 * x6 + negate (sfromIntegral (sscalar 0))))))))"++fblowupLetPP :: Assertion+fblowupLetPP = do+  resetVarCounter+  let fblowupLetT = fblowupLet @(AstTensor AstMethodLet FullSpan) @Double 0 1+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupLetT (FTKR [4] FTKScalar)+  printArtifactSimple artifactRev+    @?= "\\dret v1 -> tlet (sfromR v1 !$ [0]) (\\x3 -> tlet (sfromR v1 !$ [1]) (\\x4 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x8 -> tlet (x8 + x8) (\\x9 -> rfromS (soneHot (recip x4 * x9) [0] + soneHot ((negate x3 / (x4 * x4)) * x9) [1])))))"+  printArtifactPrimalSimple artifactRev+    @?= "\\v1 -> tlet (sfromR v1 !$ [0]) (\\x3 -> tlet (sfromR v1 !$ [1]) (\\x4 -> tlet (x3 / x4) (\\x5 -> tlet ((x5 + x5) + negate (sfromIntegral (sscalar 0))) (\\x6 -> rfromS (sscalar 0.499999985 * x6 + negate (sfromIntegral (sscalar 0)))))))"++fblowupLetPP23 :: Assertion+fblowupLetPP23 = do+  resetVarCounter+  let fblowupLetT = fblowupLet @(AstTensor AstMethodLet FullSpan) @Double 2 3+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupLetT (FTKR [4] FTKScalar)+  printArtifactSimple artifactRev+    @?= "\\dret v1 -> tlet (sfromR v1 !$ [0]) (\\x5 -> tlet (sfromR v1 !$ [1]) (\\x6 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x14 -> tlet (sscalar 0.499999985 * x14 + sscalar 0.499999985 * x14) (\\x15 -> tlet (sscalar 0.499999985 * x15 + sscalar 0.499999985 * x15) (\\x16 -> tlet (x16 + x16) (\\x17 -> rfromS (soneHot (recip x6 * x17) [0] + soneHot ((negate x5 / (x6 * x6)) * x17) [1])))))))"+  printArtifactPrimalSimple artifactRev+    @?= "\\v1 -> tlet (sfromR v1 !$ [0]) (\\x5 -> tlet (sfromR v1 !$ [1]) (\\x6 -> tlet (x5 / x6) (\\x7 -> tlet ((x7 + x7) + negate (sfromIntegral (sscalar 2))) (\\x8 -> tlet (sscalar 0.499999985 * x8) (\\x9 -> tlet ((x9 + x9) + negate (sfromIntegral (sscalar 2))) (\\x10 -> tlet (sscalar 0.499999985 * x10) (\\x11 -> tlet ((x11 + x11) + negate (sfromIntegral (sscalar 2))) (\\x12 -> rfromS (sscalar 0.499999985 * x12 + negate (sfromIntegral (sscalar 2)))))))))))"+  printArtifactSimple (simplifyArtifact artifactRev)+    @?= "\\dret v1 -> rfromS (tlet (sfromR v1 !$ [1]) (\\x6 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x14 -> tlet (sscalar 0.499999985 * x14 + sscalar 0.499999985 * x14) (\\x15 -> tlet (sscalar 0.499999985 * x15 + sscalar 0.499999985 * x15) (\\x16 -> tlet (x16 + x16) (\\x17 -> soneHot (recip x6 * x17) [0] + soneHot ((negate (sfromR v1 !$ [0]) / (x6 * x6)) * x17) [1]))))))"++fblowupLetPP10 :: Assertion+fblowupLetPP10 = do+  resetVarCounter+  let fblowupLetT = fblowupLet @(AstTensor AstMethodLet FullSpan) @Double 0 6+  let (artifactRev, _) = revArtifactDelta UseIncomingCotangent fblowupLetT (FTKR [2] FTKScalar)+  printArtifactSimple artifactRev+    @?= "\\dret v1 -> tlet (sfromR v1 !$ [0]) (\\x8 -> tlet (sfromR v1 !$ [1]) (\\x9 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x23 -> tlet (sscalar 0.499999985 * x23 + sscalar 0.499999985 * x23) (\\x24 -> tlet (sscalar 0.499999985 * x24 + sscalar 0.499999985 * x24) (\\x25 -> tlet (sscalar 0.499999985 * x25 + sscalar 0.499999985 * x25) (\\x26 -> tlet (sscalar 0.499999985 * x26 + sscalar 0.499999985 * x26) (\\x27 -> tlet (sscalar 0.499999985 * x27 + sscalar 0.499999985 * x27) (\\x28 -> tlet (x28 + x28) (\\x29 -> rfromS (soneHot (recip x9 * x29) [0] + soneHot ((negate x8 / (x9 * x9)) * x29) [1]))))))))))"+  printArtifactPrimalSimple artifactRev+    @?= "\\v1 -> tlet (sfromR v1 !$ [0]) (\\x8 -> tlet (sfromR v1 !$ [1]) (\\x9 -> tlet (x8 / x9) (\\x10 -> tlet ((x10 + x10) + negate (sfromIntegral (sscalar 0))) (\\x11 -> tlet (sscalar 0.499999985 * x11) (\\x12 -> tlet ((x12 + x12) + negate (sfromIntegral (sscalar 0))) (\\x13 -> tlet (sscalar 0.499999985 * x13) (\\x14 -> tlet ((x14 + x14) + negate (sfromIntegral (sscalar 0))) (\\x15 -> tlet (sscalar 0.499999985 * x15) (\\x16 -> tlet ((x16 + x16) + negate (sfromIntegral (sscalar 0))) (\\x17 -> tlet (sscalar 0.499999985 * x17) (\\x18 -> tlet ((x18 + x18) + negate (sfromIntegral (sscalar 0))) (\\x19 -> tlet (sscalar 0.499999985 * x19) (\\x20 -> tlet ((x20 + x20) + negate (sfromIntegral (sscalar 0))) (\\x21 -> rfromS (sscalar 0.499999985 * x21 + negate (sfromIntegral (sscalar 0)))))))))))))))))"+  printArtifactSimple (simplifyArtifact artifactRev)+    @?= "\\dret v1 -> rfromS (tlet (sfromR v1 !$ [1]) (\\x9 -> tlet (sscalar 0.499999985 * sfromR dret) (\\x23 -> tlet (sscalar 0.499999985 * x23 + sscalar 0.499999985 * x23) (\\x24 -> tlet (sscalar 0.499999985 * x24 + sscalar 0.499999985 * x24) (\\x25 -> tlet (sscalar 0.499999985 * x25 + sscalar 0.499999985 * x25) (\\x26 -> tlet (sscalar 0.499999985 * x26 + sscalar 0.499999985 * x26) (\\x27 -> tlet (sscalar 0.499999985 * x27 + sscalar 0.499999985 * x27) (\\x28 -> tlet (x28 + x28) (\\x29 -> soneHot (recip x9 * x29) [0] + soneHot ((negate (sfromR v1 !$ [0]) / (x9 * x9)) * x29) [1])))))))))"++-- TODO: should do 1000000 in a few seconds+blowupTests :: TestTree+blowupTests = testGroup "Catastrophic blowup avoidance tests"+  [ testCase "blowup prim 7" $ do+      assertEqualUpToEpsilon' 1e-5+        (ringestData [2] [0.3333332333333467,-0.22222215555556446])+        (rev' @Double @0 (fblowup 7) (ringestData [2] [2, 3]))+  , testCase "blowupLet prim 2000" $ do+      assertEqualUpToEpsilon' 1e-10+        (ringestData [2] [0.3333133339329949,-0.22220888928866325])+        (rev' @Double @0 (fblowupLet 1 2000) (ringestData [2] [2, 3]))+  , testCase "blowupLet 7000" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [0.3332633406816766,-0.22217556045445108])+        (grad (kfromR @_ @Double . rsum0 . fblowupLet 0 7000) (ringestData [2] [2, 3]))+  , testCase "blowupLet tbuild0" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [333.2633406816765,-222.175560454451])+        (grad (kfromR @_ @Double . rsum0 . (\intputs -> rbuild1 1000 (\_ -> fblowupLet 0 7000 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupLet tbuild2" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [333.2633406816765,-222.175560454451])+        (grad (kfromR @_ @Double . rsum0 . (\intputs -> rbuild1 1000 (\_ -> fblowupLet 2 7000 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupLet tbuildi" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [333.2983351701977,-222.19889011346513])+        (grad (kfromR @_ @Double . rsum0+               . (\intputs -> rbuild1 1000 (\i -> fblowupLet i 3500 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupLet tbuildc" $ do+      assertEqualUpToEpsilon 1e-7+        (ringestData [2] [333.326333406717,-222.21755560448116])+        (cgrad (kfromR @_ @Double . rsum0+                . (\intputs -> rbuild1 1000 (\i -> fblowupLet i 700 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupLet prim tbuild" $ do+      assertEqualUpToEpsilon 1e-7+        (ringestData [2] [33.33263334067178,-22.221755560447928])+        (grad (kfromR @_ @Double . rsum0+               . (\intputs -> rbuild1 100 (\i -> fblowupLet i 700 intputs)))+              (ringestData [2] [2, 3]))+  , testCase "blowupMult 3" $ do+      assertEqualUpToEpsilon' 1e-5+        (ringestData [2] [2.999999730000007,1.9999998200000046])+        (rev' @Double @0 (fblowupMult 3) (ringestData [2] [2, 3]))+  , testCase "blowupMultLet 5" $ do+      assertEqualUpToEpsilon' 1e-10+        (ringestData [2] [2.9999995500000267,1.9999997000000178])+        (rev' @Double @0 (fblowupMultLet 0 5)+                                   (ringestData [2] [2, 3]))+  , testCase "blowupMultLet 50" $ do+      assertEqualUpToEpsilon' 1e-10+        (ringestData [2] [2.999995500001215,1.99999700000081])+        (rev' @Double @0 (fblowupMultLet 0 50)+                                   (ringestData [2] [2, 3]))+  , testCase "blowupMultLet tbuild1" $ do+      assertEqualUpToEpsilon 1e-10+        (ringestData [2] [14.9999773958889,39.9999398380561])+        (grad (kfromR @_ @Double . rsum0+               . (\intputs -> rbuild1 100 (\i -> fblowupMultLet i 50 intputs)))+              (ringestData [2] [0.2, 0.3]))+  ]++concatBuild33 :: (ADReady target, GoodScalar r)+             => target (TKR 1 r) -> target (TKR 2 r)+concatBuild33 _r =+  rbuild1 5 (\i ->+    rbuild1 2 (\j -> rfromIndex0 (maxH j (i `quotH` (j + 1)))))++testConcatBuild3PP :: Assertion+testConcatBuild3PP = do+  resetVarCounter+  let t = concatBuild33 @(AstTensor AstMethodLet FullSpan) @Float+      (var3, ast3) = funToAst (FTKR [3] FTKScalar) Nothing t+  "\\" ++ printAstVarName var3+       ++ " -> " ++ printAstSimple ast3+    @?= "\\v1 -> tfromPrimal (STKR (SNat @2) STKScalar) (rfromS (sgather [] (sfromIntegral (tfromVector (SNat @2) (STKS [5,2] STKScalar) (fromList [sreplicate @5 (siota (SNat @2)), quotH (str (sreplicate @2 (siota (SNat @5)))) (sreplicate @5 (sconcrete (sreplicate [2] 1) + siota (SNat @2)))]))) (\\[i5, i4] -> [ifH (0 <=. i4 + quotH (negate i5) (1 + i4)) 0 1, i5, i4])))"++testConcatBuild3PP2 :: Assertion+testConcatBuild3PP2 = do+  resetVarCounter+  let t = concatBuild33 @(AstTensor AstMethodLet FullSpan) @Double+  let (artifactRev, _) =+        revArtifactDelta UseIncomingCotangent t (FTKR [3] FTKScalar)+  printArtifactSimple artifactRev+    @?= "\\dret v1 -> rfromS (sconcrete (sreplicate [3] 0.0))"+  printArtifactPrimalSimple artifactRev+    @?= "\\v1 -> rfromS (sgather [] (sfromIntegral (tfromVector (SNat @2) (STKS [5,2] STKScalar) (fromList [sreplicate @5 (siota (SNat @2)), quotH (str (sreplicate @2 (siota (SNat @5)))) (sreplicate @5 (sconcrete (sreplicate [2] 1) + siota (SNat @2)))]))) (\\[i6, i7] -> [ifH (0 <=. i7 + quotH (negate i6) (1 + i7)) 0 1, i6, i7]))"+  printArtifactPrimalSimple (simplifyArtifact artifactRev)+    @?= "\\v1 -> rfromS (sgather [] (sconcrete (sfromListLinear [2,5,2] [0.0,1.0,0.0,1.0,0.0,1.0,0.0,1.0,0.0,1.0,0.0,0.0,1.0,0.0,2.0,1.0,3.0,1.0,4.0,2.0])) (\\[i6, i7] -> [ifH (0 <=. i7 + quotH (negate i6) (1 + i7)) 0 1, i6, i7]))"
+ test/simplified/TestConvSimplified.hs view
@@ -0,0 +1,1993 @@+{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}+{-# OPTIONS_GHC -fno-cse #-}+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+-- | Tests of convolution and disparity cost volume defined using the build+-- operation of ranked tensors.+module TestConvSimplified (testTrees) where++import Prelude++import Control.Exception.Assert.Sugar+import GHC.Exts (IsList (..))+import GHC.TypeLits (KnownNat)+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape++import HordeAd+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId (resetVarCounter)+import HordeAd.Core.AstInterpret+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersAst+import HordeAd.Core.Delta+import HordeAd.Core.Ops+import HordeAd.Core.OpsAst++import CrossTesting+import EqEpsilon++testTrees :: [TestTree]+testTrees =+  [ testCase "KonstG0Rev" testKonstG0Rev+  , testCase "KonstG0Tiny1" testKonstG0Tiny1+  , testCase "KonstG0TinyS" testKonstG0TinyS+  , testCase "KonstG0TinyA" testKonstG0TinyA+  , testCase "KonstG0LittleA" testKonstG0LittleA+  , testCase "Replicate0Rev" testReplicate0Rev+  , testCase "Replicate0Tiny1" testReplicate0Tiny1+  , testCase "Replicate0TinyS" testReplicate0TinyS+  , testCase "Replicate0TinyA" testReplicate0TinyA+  , testCase "Replicate0LittleA" testReplicate0LittleA+  , testCase "Konst5LittleB" testKonst5LittleB+  , testCase "Konst5LittleC" testKonst5LittleC+  , testCase "Konst5BigB" testKonst5BigB+  , testCase "KonstNotBigB" testKonstNotBigB+  , testCase "Konst5BigC" testKonst5BigC+  , testCase "KonstNotBigC" testKonstNotBigC+  , testCase "Konst5LittleB128b" testKonst5LittleB128b+  , testCase "Konst5LittleC128b" testKonst5LittleC128b+  , testCase "Konst5BigB128b" testKonst5BigB128b+  , testCase "KonstNotBigB128b" testKonstNotBigB128b+  , testCase "Konst5BigC128b" testKonst5BigC128b+  , testCase "KonstNotBigC128b" testKonstNotBigC128b+  , testCase "Konst5LittleB128c" testKonst5LittleB128c+  , testCase "Konst5LittleC128c" testKonst5LittleC128c+  , testCase "Konst5BigB128c" testKonst5BigB128c+  , testCase "KonstNotBigB128c" testKonstNotBigB128c+  , testCase "Konst5BigC128c" testKonst5BigC128c+  , testCase "KonstNotBigC128c" testKonstNotBigC128c+--  , testCase "Konst5LittleB128bc" testKonst5LittleB128bc+--  , testCase "Konst5LittleC128bc" testKonst5LittleC128bc+--  , testCase "Konst5BigB128bc" testKonst5BigB128bc+--  , testCase "KonstNotBigB128cb" testKonstNotBigB128cb+--  , testCase "Konst5BigC128cb" testKonst5BigC128cb+--  , testCase "KonstNotBigC128cb" testKonstNotBigC128cb+  , testCase "Replicate0RevLaborious" testReplicate0RevLaborious+  , testCase "Replicate0Tiny1Laborious" testReplicate0Tiny1Laborious+  , testCase "Replicate0TinySLaborious" testReplicate0TinySLaborious+  , testCase "Replicate0TinyALaborious" testReplicate0TinyALaborious+  , testCase "Replicate0LittleALaborious" testReplicate0LittleALaborious+  , testCase "Konst5LittleBLaborious" testKonst5LittleBLaborious+  , testCase "Konst5LittleCLaborious" testKonst5LittleCLaborious+  , testCase "Konst5BigBLaborious" testKonst5BigBLaborious+  , testCase "KonstNotBigBLaborious" testKonstNotBigBLaborious+  , testCase "Konst5BigCLaborious" testKonst5BigCLaborious+  , testCase "KonstNotBigCLaborious" testKonstNotBigCLaborious+  , testCase "Konst5LittleBLaborious128b" testKonst5LittleBLaborious128b+  , testCase "Konst5LittleCLaborious128b" testKonst5LittleCLaborious128b+--  , testCase "Konst5BigBLaborious128b" testKonst5BigBLaborious128b+--  , testCase "KonstNotBigBLaborious128b" testKonstNotBigBLaborious128b+--  , testCase "Konst5BigCLaborious128b" testKonst5BigCLaborious128b+--  , testCase "KonstNotBigCLaborious128b" testKonstNotBigCLaborious128b+  , testCase "Konst5LittleBLaborious128c" testKonst5LittleBLaborious128c+  , testCase "Konst5LittleCLaborious128c" testKonst5LittleCLaborious128c+--  , testCase "Konst5BigBLaborious128c" testKonst5BigBLaborious128c+--  , testCase "KonstNotBigBLaborious128c" testKonstNotBigBLaborious128c+--  , testCase "Konst5BigCLaborious128c" testKonst5BigCLaborious128c+--  , testCase "KonstNotBigCLaborious128c" testKonstNotBigCLaborious128c+--  , testCase "Konst5LittleBLaborious128bc" testKonst5LittleBLaborious128bc+--  , testCase "Konst5LittleCLaborious128bc" testKonst5LittleCLaborious128bc+--  , testCase "Konst5BigBLaborious128bc" testKonst5BigBLaborious128bc+--  , testCase "KonstNotBigBLaborious128cb" testKonstNotBigBLaborious128cb+--  , testCase "Konst5BigCLaborious128cb" testKonst5BigCLaborious128cb+--  , testCase "KonstNotBigCLaborious128cb" testKonstNotBigCLaborious128cb+--  , testCase "Replicate0RevPadded" testReplicate0RevPadded+  , testCase "Replicate0Tiny1Padded" testReplicate0Tiny1Padded+  , testCase "Replicate0TinySPadded" testReplicate0TinySPadded+  , testCase "Replicate0TinyAPadded" testReplicate0TinyAPadded+  , testCase "Replicate0LittleAPadded" testReplicate0LittleAPadded+--  , testCase "Konst5LittleBPadded" testKonst5LittleBPadded+--  , testCase "Konst5LittleCPadded" testKonst5LittleCPadded+--  , testCase "Konst5BigBPadded" testKonst5BigBPadded+--  , testCase "KonstNotBigBPadded" testKonstNotBigBPadded+--  , testCase "Konst5BigCPadded" testKonst5BigCPadded+--  , testCase "KonstNotBigCPadded" testKonstNotBigCPadded+--  , testCase "Konst5LittleBPadded128b" testKonst5LittleBPadded128b+--  , testCase "Konst5LittleCPadded128b" testKonst5LittleCPadded128b+--  , testCase "Konst5BigBPadded128b" testKonst5BigBPadded128b+--  , testCase "KonstNotBigBPadded128b" testKonstNotBigBPadded128b+--  , testCase "Konst5BigCPadded128b" testKonst5BigCPadded128b+--  , testCase "KonstNotBigCPadded128b" testKonstNotBigCPadded128b+--  , testCase "Konst5LittleBPadded128c" testKonst5LittleBPadded128c+--  , testCase "Konst5LittleCPadded128c" testKonst5LittleCPadded128c+--  , testCase "Konst5BigBPadded128c" testKonst5BigBPadded128c+--  , testCase "KonstNotBigBPadded128c" testKonstNotBigBPadded128c+--  , testCase "Konst5BigCPadded128c" testKonst5BigCPadded128c+--  , testCase "KonstNotBigCPadded128c" testKonstNotBigCPadded128c+--  , testCase "Konst5LittleBPadded128bc" testKonst5LittleBPadded128bc+--  , testCase "Konst5LittleCPadded128bc" testKonst5LittleCPadded128bc+--  , testCase "Konst5BigBPadded128bc" testKonst5BigBPadded128bc+--  , testCase "KonstNotBigBPadded128cb" testKonstNotBigBPadded128cb+--  , testCase "Konst5BigCPadded128cb" testKonst5BigCPadded128cb+--  , testCase "KonstNotBigCPadded128cb" testKonstNotBigCPadded128cb+  , testCase "disparityKonst" test_disparityKonst+  , testCase "disparityKonst2" test_disparityKonst2+  , testCase "disparitySmall" test_disparitySmall+  , testCase "ConvTomsSliceRev" testTomsSliceRev+  , testCase "ConvTomsSlice" testTomsSlice+  , testCase "ConvTomsSlicePP" testTomsSlicePP+  , testCase "minimizedCNNOPP0c" testCNNOPP0c+  , testCase "minimizedCNNOPP0b" testCNNOPP0b+  , testCase "minimizedCNNOPP1e" testCNNOPP1e+  , testCase "minimizedCNNOPP2" testCNNOPP2+  , testCase "minimizedCNNOPP2b" testCNNOPP2b+--  , testCase "minimizedCNNOPP3" testCNNOPP3+  , testCase "minimizedCNNOPP3b" testCNNOPP3b+  , testCase "minimizedCNNOPP4" testCNNOPP4+  , testCase "minimizedCNNOPP4b" testCNNOPP4b+  , testCase "minimizedCNNOPP5" testCNNOPP5+  , testCase "minimizedCNNOPP5b" testCNNOPP5b+  , testCase "minimizedCNNOPP6" testCNNOPP6+  , testCase "minimizedCNNOPP6b" testCNNOPP6b+  , testCase "minimizedCNNOPP7" testCNNOPP7+  , testCase "minimizedCNNOPP7b" testCNNOPP7b+--  , testCase "minimizedPaddedCNNOPP0c" testPaddedCNNOPP0c+--  , testCase "minimizedPaddedCNNOPP0b" testPaddedCNNOPP0b+--  , testCase "minimizedPaddedCNNOPP1e" testPaddedCNNOPP1e+  , testCase "minimizedPaddedCNNOPP1b" testPaddedCNNOPP1b+  , testCase "minimizedPaddedCNNOPPLet" testPaddedCNNOPPLet+  , testCase "minimizedPaddedCNNOPPLet2" testPaddedCNNOPPLet2+--  , testCase "minimizedPaddedCNNOPP2" testPaddedCNNOPP2+  , testCase "minimizedCNNOPP0cW" testCNNOPP0cW+  , testCase "minimizedCNNOPP0bW" testCNNOPP0bW+  , testCase "minimizedCNNOPP1bW" testCNNOPP1bW+  , testCase "minimizedCNNOPP4bW" testCNNOPP4bW+  , testCase "minimizedCNNOPP4bD" testCNNOPP4bD+  , testCase "minimizedCNNOPP5aW" testCNNOPP5aW+  , testCase "minimizedCNNOPP5bW" testCNNOPP5bW+  , testCase "minimizedCNNOPP5cW" testCNNOPP5cW+  , testCase "minimizedCNNOPP5dW" testCNNOPP5dW+  ]++-- The examples reproduced and transformed in this file are borrowed+-- from https://github.com/benl23x5/adops.+-- Here they are defined using ranked tensors.++-- * A non-laborious version (depends on indexing OOB giving 0 consistently)++conv2d1+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2d1 = conv2dUnpadded $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 1, 1, 1]) [-0.2]++conv2dA+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dA = conv2dUnpadded $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 2, 1, 1]) [-0.2, 25.0003]++conv2dB+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dB = conv2dUnpadded (rconcrete $ unConcrete t16b)++testKonstG0Rev :: Assertion+testKonstG0Rev =+  assertEqualUpToEpsilon 1e-4+    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dB) (rrepl [2, 2, 2, 2] 0))++testKonstG0Tiny1 :: Assertion+testKonstG0Tiny1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [-0.2])+    (rev' @Double @4 conv2d1 (rrepl [1, 1, 1, 1] 0))++testKonstG0TinyS :: Assertion+testKonstG0TinyS =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [582665.99432])+    (rev' @Double @4+          (conv2dUnpadded $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b)))+          (ringestData [1, 1, 1, 1] [0]))++testKonstG0TinyA :: Assertion+testKonstG0TinyA =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 2, 1, 1] [-0.2,25.0003])+    (rev' @Double @4 conv2dA (rrepl [1, 2, 1, 1] 0))++testKonstG0LittleA :: Assertion+testKonstG0LittleA =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 2, 2, 2] [-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003])+    (rev' @Double @4 conv2dA (rrepl [2, 2, 2, 2] 0))++conv2dC+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dC = flip conv2dUnpadded (rconcrete $ unConcrete t16b)++conv2dB128b+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dB128b = conv2dUnpadded (rconcrete $ unConcrete t128b)++conv2dC128b+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dC128b = flip conv2dUnpadded (rconcrete $ unConcrete t128b)++conv2dB128c+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dB128c = conv2dUnpadded (rconcrete $ unConcrete t128c)++conv2dC128c+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dC128c = flip conv2dUnpadded (rconcrete $ unConcrete t128c)++testReplicate0Rev :: Assertion+testReplicate0Rev =+  assertEqualUpToEpsilon 1e-4+    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dB) (rrepl [2, 2, 2, 2] 0))++testReplicate0Tiny1 :: Assertion+testReplicate0Tiny1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [-0.2])+    (rev' @Double @4 conv2d1 (rrepl [1, 1, 1, 1] 0))++testReplicate0TinyS :: Assertion+testReplicate0TinyS =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [582665.99432])+    (rev' @Double @4+          (conv2dUnpadded $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b)))+          (ringestData [1, 1, 1, 1] [0]))++testReplicate0TinyA :: Assertion+testReplicate0TinyA =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 2, 1, 1] [-0.2,25.0003])+    (rev' @Double @4 conv2dA (rrepl [1, 2, 1, 1] 0))++testReplicate0LittleA :: Assertion+testReplicate0LittleA =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 2, 2, 2] [-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003])+    (rev' @Double @4 conv2dA (rrepl [2, 2, 2, 2] 0))++-- with data t16++testKonst5LittleB :: Assertion+testKonst5LittleB =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dB (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5LittleC :: Assertion+testKonst5LittleC =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8])+    (rev' @Double @4 conv2dC (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5BigB :: Assertion+testKonst5BigB =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dB (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigB :: Assertion+testKonstNotBigB =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dB+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++testKonst5BigC :: Assertion+testKonst5BigC =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dC (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigC :: Assertion+testKonstNotBigC =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dC+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++-- with data t128b++testKonst5LittleB128b :: Assertion+testKonst5LittleB128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004])+    (rev' @Double @4 conv2dB128b (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5LittleC128b :: Assertion+testKonst5LittleC128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987])+    (rev' @Double @4 conv2dC128b (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5BigB128b :: Assertion+testKonst5BigB128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dB128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigB128b :: Assertion+testKonstNotBigB128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dB128b+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++testKonst5BigC128b :: Assertion+testKonst5BigC128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dC128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigC128b :: Assertion+testKonstNotBigC128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dC128b+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++-- with data t128c++testKonst5LittleB128c :: Assertion+testKonst5LittleB128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4])+    (rev' @Double @4 conv2dB128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5LittleC128c :: Assertion+testKonst5LittleC128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992])+    (rev' @Double @4 conv2dC128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5BigB128c :: Assertion+testKonst5BigB128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dB128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigB128c :: Assertion+testKonstNotBigB128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dB128c+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++testKonst5BigC128c :: Assertion+testKonst5BigC128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dC128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigC128c :: Assertion+testKonstNotBigC128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dC128c+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++-- with data t128b and t128c+{-+testKonst5LittleB128bc :: Assertion+testKonst5LittleB128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])+    (rev' @Double @4 conv2dB128b t128c)++testKonst5LittleC128bc :: Assertion+testKonst5LittleC128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dC128b t128c)++testKonst5BigB128bc :: Assertion+testKonst5BigB128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])+    (rev' @Double @4 conv2dB128b t128c)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigB128cb :: Assertion+testKonstNotBigB128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002])+    (rev' @Double @4 conv2dB128c t128b)++testKonst5BigC128cb :: Assertion+testKonst5BigC128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])+    (rev' @Double @4 conv2dC128c t128b)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigC128cb :: Assertion+testKonstNotBigC128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])+    (rev' @Double @4 conv2dC128c t128b)+-}+++-- * A laborious version (meaning, out of bounds indexing is handled explicitly)++-- | Unpadded full convolution,+--   where the output size is the same as the input size.+--+-- It guards the out of bounds indexing behind a conditional+-- to prevent changed values after vectorization,+-- but the guarding is no longer needed, so this is only for testing.+--+-- BTW, the indexing lower bounds in the code are spurious,+-- so they get simplified away in the resulting AST program.+conv2dUnpaddedL+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpaddedL arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezL shK1 arrA [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rdot0 arrAt arrKt+    _ -> error "conv2dUnpaddedL: impossible pattern needlessly required"++-- | Slice a section out of a tensor,+--   given a base offset and shape of the section.+--+--   If the slice extends out side the source array then the corresponding+--   elements are set to zero.+slicezL+  :: (ADReady target, GoodScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicezL shOut d ixBase =+  rbuild shOut $ \ixResult -> indexz0L d (ixrZipWith (+) ixBase ixResult)++-- | Retrieve the element at the given index,+--   returning zero for out of range indices.+--+-- Warning: this uses ix twice and within0 again uses it twice,+-- so this variant without tlet should be used only when it's known+-- that ix is of small constant size (e.g., if it contains conditionals+-- that compare big tensors or their minimal elements, it likely is not,+-- unless the tensors are under tlet and only variables representing them+-- are used).+indexz0L+  :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)+  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)+indexz0L d ix = ifH (within0 @target (rshape @target d) ix) (d ! ix) (rscalar 0)++-- | Given an index and shape, check if the index is fully within the shape.+-- Note that @ix@ is used twice, so should be shared outside.+within0+  :: forall target n. (ADReady target, KnownNat n)+  => IShR n -> IxROf target n -> BoolOf target+within0 sh ix =+  let within :: IntOf target -> IntOf target -> BoolOf target+      within i dim = 0 <=. i &&* dim >. i+  in foldr (&&*) true+     $ zipWith within (toList ix) (map fromIntegral $ toList sh)++conv2d1Laborious+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2d1Laborious = conv2dUnpaddedL $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 1, 1, 1]) [-0.2]++conv2dALaborious+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dALaborious =+  conv2dUnpaddedL $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 2, 1, 1]) [-0.2, 25.0003]++conv2dBLaborious+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBLaborious = conv2dUnpaddedL (rconcrete $ unConcrete t16b)++conv2dCLaborious+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCLaborious = flip conv2dUnpaddedL (rconcrete $ unConcrete t16b)++conv2dBLaborious128b+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBLaborious128b = conv2dUnpaddedL (rconcrete $ unConcrete t128b)++conv2dCLaborious128b+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCLaborious128b = flip conv2dUnpaddedL (rconcrete $ unConcrete t128b)++conv2dBLaborious128c+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBLaborious128c = conv2dUnpaddedL (rconcrete $ unConcrete t128c)++conv2dCLaborious128c+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCLaborious128c = flip conv2dUnpaddedL (rconcrete $ unConcrete t128c)++testReplicate0RevLaborious :: Assertion+testReplicate0RevLaborious =+  assertEqualUpToEpsilon 1e-4+    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBLaborious) (rrepl [2, 2, 2, 2] 0))++testReplicate0Tiny1Laborious :: Assertion+testReplicate0Tiny1Laborious =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [-0.2])+    (rev' @Double @4 conv2d1Laborious (rrepl [1, 1, 1, 1] 0))++testReplicate0TinySLaborious :: Assertion+testReplicate0TinySLaborious =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 1, 1, 1] [582665.99432])+    (rev' @Double @4+          (conv2dUnpaddedL $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b)))+          (ringestData [1, 1, 1, 1] [0]))++testReplicate0TinyALaborious :: Assertion+testReplicate0TinyALaborious =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1, 2, 1, 1] [-0.2,25.0003])+    (rev' @Double @4 conv2dALaborious (rrepl [1, 2, 1, 1] 0))++testReplicate0LittleALaborious :: Assertion+testReplicate0LittleALaborious =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 2, 2, 2] [-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003])+    (rev' @Double @4 conv2dALaborious (rrepl [2, 2, 2, 2] 0))++-- with data t16++testKonst5LittleBLaborious :: Assertion+testKonst5LittleBLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dBLaborious (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5LittleCLaborious :: Assertion+testKonst5LittleCLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8])+    (rev' @Double @4 conv2dCLaborious (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5BigBLaborious :: Assertion+testKonst5BigBLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dBLaborious (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBLaborious :: Assertion+testKonstNotBigBLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,582597.1,582625.8943200001,582597.1,582625.8943200001])+    (rev' @Double @4 conv2dBLaborious+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++testKonst5BigCLaborious :: Assertion+testKonst5BigCLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dCLaborious (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCLaborious :: Assertion+testKonstNotBigCLaborious =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0,40.1,8.0,11.0,-3.0,0.0,0.0,0.0,0.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dCLaborious+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++-- with data t128b++testKonst5LittleBLaborious128b :: Assertion+testKonst5LittleBLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,112.3003,251.5006,209.49462,482.69492000000014,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004])+    (rev' @Double @4 conv2dBLaborious128b (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5LittleCLaborious128b :: Assertion+testKonst5LittleCLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987])+    (rev' @Double @4 conv2dCLaborious128b (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++{-+testKonst5BigBLaborious128b :: Assertion+testKonst5BigBLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dBLaborious128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBLaborious128b :: Assertion+testKonstNotBigBLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dBLaborious128b+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++testKonst5BigCLaborious128b :: Assertion+testKonst5BigCLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCLaborious128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCLaborious128b :: Assertion+testKonstNotBigCLaborious128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCLaborious128b+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))+-}++-- with data t128c++testKonst5LittleBLaborious128c :: Assertion+testKonst5LittleBLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,58.2,140.3,90.4,212.4])+    (rev' @Double @4 conv2dBLaborious128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5LittleCLaborious128c :: Assertion+testKonst5LittleCLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992])+    (rev' @Double @4 conv2dCLaborious128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++{-+testKonst5BigBLaborious128c :: Assertion+testKonst5BigBLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dBLaborious128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBLaborious128c :: Assertion+testKonstNotBigBLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dBLaborious128c+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++testKonst5BigCLaborious128c :: Assertion+testKonst5BigCLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dCLaborious128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCLaborious128c :: Assertion+testKonstNotBigCLaborious128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dCLaborious128c+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))+-}++-- with data t128b and t128c+{-+testKonst5LittleBLaborious128bc :: Assertion+testKonst5LittleBLaborious128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])+    (rev' @Double @4 conv2dBLaborious128b t128c)++testKonst5LittleCLaborious128bc :: Assertion+testKonst5LittleCLaborious128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1627.8210700004993,1571.2321300004994,1047.1431900004002,393.6715900002,1132.9261600005002,1188.6375200005,803.7488800004002,316.57160000019996,675.7488800003999,828.6545600004001,577.7659200003001,220.57728000019998,215.6659200003,388.5716000003,245.5772800002,94.68864000010001,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2725.0393200008984,1831.7390200008983,1259.3728000004999,568.6722000005001,2551.139320000898,1660.8390200008987,1151.3728000005,501.6722000005,1903.750080000699,1174.5497800006997,803.9778800004001,340.5775800004001,854.9778800004001,628.8778800004001,450.1892400002,198.8889400002,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @4 conv2dCLaborious128b t128c)++testKonst5BigBLaborious128bc :: Assertion+testKonst5BigBLaborious128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,112.3003,251.5006,417.79492000000005,494.89491000000015,209.49462,482.69492000000014,778.9778800001002,952.0721900001002,229.49462000000003,610.5892400000002,1113.1722000001,1412.1551500001997,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,56.58894000000004,580.6778800001001,1234.1494800003,1627.8210700004993,3.000000000000032,65.90000000000003,106.90000000000003,173.90000000000006,164.10000000000002,365.89432000010004,593.1946200001,821.2892400002004,667.2003000000001,1060.8778800002,1500.2781800001994,1870.0614400004986,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898,893.3003,1465.6665200003993,2156.3671200003987,2725.039320000898])+    (rev' @Double @4 conv2dBLaborious128b t128c)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBLaborious128cb :: Assertion+testKonstNotBigBLaborious128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002,54.100300000000004,111.20060000000001,191.4006,228.4006,119.09462,270.29492000000005,435.28356000009995,519.1778800001,109.09462000000002,318.19492,563.3835600001,687.2778800001003,174.08894000000004,477.28924000000006,774.2665200002001,931.9551600002003,58.2,140.3,226.39432,266.49431000000004,90.4,212.4,343.69432000000006,432.89431000000013,120.4,292.39432000000005,549.78864,724.8772700001001,-117.5,103.38864000010005,459.88296000009996,695.8659100003002])+    (rev' @Double @4 conv2dBLaborious128c t128b)++testKonst5BigCLaborious128cb :: Assertion+testKonst5BigCLaborious128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])+    (rev' @Double @4 conv2dCLaborious128c t128b)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCLaborious128cb :: Assertion+testKonstNotBigCLaborious128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003,2640.8154000007976,1836.3264600007988,1163.4488800005001,483.7716000003,2412.414800000798,1662.026160000799,1046.2488800005003,446.7716000003,2121.6375200006987,1436.2432000006995,914.5659200004003,399.8772800003,1953.5375200006988,1258.1432000006998,794.3659200004003,359.8772800003,1712.044990000598,1566.644690000599,1143.0671100004001,478.5721900004001,1445.5506800005985,1358.3503800005992,1016.8728000004002,438.47220000040005,1279.150680000599,1224.1503800005996,922.5728000004001,389.3722000004,987.1677200004992,962.1674200005002,710.5841600003,303.48356000030003])+    (rev' @Double @4 conv2dCLaborious128c t128b)+-}+++-- * A padded version (out of bounds indexing is not possible)++-- | Full convolution with just enough extra external zero padding+--   to ensure that the output size is the same as the input size+--   and all input points are read the same number of times.+--+--   The same result could be accomplished by tweaking indexes slightly+--   in conv2dUnpadded, but here additionally all bounds checks in the code+--   are spurious and will be simplified away in the resulting AST program.+conv2dPadded+  :: forall target r. (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dPadded arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]+      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case+        [iImg, iCinp, iPh, iPw] ->+          ifH (iPh <. fromIntegral (nKh `div` 2)+               ||* iPw <. fromIntegral (nKw `div` 2)+               ||* iPh >=. fromIntegral (nAh + nKh `div` 2)+               ||* iPw >=. fromIntegral (nAw + nKw `div` 2))+              (rscalar 0)+              (arrA ! [ iImg+                      , iCinp+                      , iPh - fromIntegral (nKh `div` 2)+                      , iPw - fromIntegral (nKw `div` 2) ])+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezL shK1 arrAPadded [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rdot0 arrAt arrKt+    _ -> error "conv2dPadded: impossible pattern needlessly required"++conv2d1Padded+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2d1Padded = conv2dPadded $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 1, 1, 1]) [-0.2]++conv2dAPadded+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dAPadded =+  conv2dPadded $ rconcrete $ Nested.rfromListPrimLinear (fromList [1, 2, 1, 1]) [-0.2, 25.0003]++conv2dBPadded+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBPadded = conv2dPadded (rconcrete $ unConcrete t16b)++conv2dCPadded+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCPadded = flip conv2dPadded (rconcrete $ unConcrete t16b)++conv2dBPadded128b+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dBPadded128b = conv2dPadded (rconcrete $ unConcrete t128b)++conv2dCPadded128b+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dCPadded128b = flip conv2dPadded (rconcrete $ unConcrete t128b)++_conv2dBPadded128c+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+_conv2dBPadded128c = conv2dPadded (rconcrete $ unConcrete t128c)++_conv2dCPadded128c+  :: (ADReady target, GoodScalar r, Differentiable r)+  => target (TKR 4 r) -> target (TKR 4 r)+_conv2dCPadded128c = flip conv2dPadded (rconcrete $ unConcrete t128c)++-- TODO: OOMs+_testReplicate0RevPadded :: Assertion+_testReplicate0RevPadded =+  assertEqualUpToEpsilon 1e-4+    (rconcrete $ Nested.rfromListPrimLinear [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.89432,28.79432,-309.09999999999997,25.8])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded) (rrepl [2, 2, 2, 2] 0))++testReplicate0Tiny1Padded :: Assertion+testReplicate0Tiny1Padded =+  assertEqualUpToEpsilon 1e-10+    (ringestData [1, 1, 1, 1] [-0.2])+    (cgrad (kfromR . rsum0 @4 @(TKScalar Double) . conv2d1Padded) (rrepl [1, 1, 1, 1] 0))++testReplicate0TinySPadded :: Assertion+testReplicate0TinySPadded =+  assertEqualUpToEpsilon 1e-10+    (ringestData [1, 1, 1, 1] [582665.99432])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) .+          (conv2dPadded $ rreplicate0N [1, 1, 1, 1] (rsum0 (rconcrete $ unConcrete t16b))))+          (ringestData [1, 1, 1, 1] [0]))++testReplicate0TinyAPadded :: Assertion+testReplicate0TinyAPadded =+  assertEqualUpToEpsilon 1e-10+    (ringestData [1, 2, 1, 1] [-0.2,25.0003])+    (cgrad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dAPadded) (rrepl [1, 2, 1, 1] 0))++testReplicate0LittleAPadded :: Assertion+testReplicate0LittleAPadded =+  assertEqualUpToEpsilon 1e-10+    (ringestData [2, 2, 2, 2] [-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dAPadded) (rrepl [2, 2, 2, 2] 0))++-- with data t16++-- TODO: OOMs+_testKonst5LittleBPadded :: Assertion+_testKonst5LittleBPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [2, 2, 2, 2] [40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,-309.09999999999997,25.8])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded) (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++-- TODO: OOMs+_testKonst5LittleCPadded :: Assertion+_testKonst5LittleCPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [2, 2, 2, 2] [18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001,18.1,29.1,32.1,40.1,582932.0,582934.99432,582597.1,582625.8943200001])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dCPadded) (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++-- TODO: OOMs+_testKonst5BigBPadded :: Assertion+_testKonst5BigBPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded) (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- TODO: OOMs+-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+_testKonstNotBigBPadded :: Assertion+_testKonstNotBigBPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [3, 2, 4, 2] [40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8,40.1,8.0,40.1,8.0,40.1,8.0,11.0,-3.0,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,582625.8943200001,28.794320000000003,-309.09999999999997,25.8])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded)+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++-- TODO: OOMs+_testKonst5BigCPadded :: Assertion+_testKonst5BigCPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [3, 2, 4, 2] [0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dCPadded) (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- TODO: OOMs+-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+_testKonstNotBigCPadded :: Assertion+_testKonstNotBigCPadded =+  assertEqualUpToEpsilon 1e-8+    (ringestData [3, 2, 4, 2] [0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997,0.0,0.0,18.1,29.1,32.1,40.1,14.0,11.0,0.0,0.0,582932.0,582934.99432,582597.1,582625.8943200001,-334.9,-309.09999999999997])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dCPadded)+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++-- with data t128b++-- TODO: OOMs+_testKonst5LittleBPadded128b :: Assertion+_testKonst5LittleBPadded128b =+  assertEqualUpToEpsilon 1e-8+    (ringestData [2, 2, 2, 2] [578.1829600001,558.1716000002,608.0772800002001,577.7659200003001,729.1778800002002,701.1835600003001,833.9722000003002,803.9778800004001,578.1829600001,558.1716000002,608.0772800002001,577.7659200003001,729.1778800002002,701.1835600003001,833.9722000003002,803.9778800004001])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dBPadded128b) (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++-- TODO: OOMs+_testKonst5LittleCPadded128b :: Assertion+_testKonst5LittleCPadded128b =+  assertEqualUpToEpsilon 1e-8+    (ringestData [2, 2, 2, 2] [1113.1722000001,1412.1551500001997,1234.1494800003002,1627.8210700004993,1500.2781800001994,1870.0614400004986,2156.3671200003987,2725.0393200008984,1113.1722000001,1412.1551500001997,1234.1494800003002,1627.8210700004993,1500.2781800001994,1870.0614400004986,2156.3671200003987,2725.0393200008984])+    (grad (kfromR . rsum0 @4 @(TKScalar Double) . conv2dCPadded128b) (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++{-+testKonst5BigBPadded128b :: Assertion+testKonst5BigBPadded128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dBPadded128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBPadded128b :: Assertion+testKonstNotBigBPadded128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993,112.3003,251.5006,209.49462,482.69492000000014,229.49462000000003,610.5892400000002,56.58894000000004,580.6778800001001,3.000000000000032,65.90000000000003,164.10000000000002,365.89432000010004,667.2003000000001,1060.8778800002,893.3003,1465.6665200003993])+    (rev' @Double @4 conv2dBPadded128b+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++testKonst5BigCPadded128b :: Assertion+testKonst5BigCPadded128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCPadded128b (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCPadded128b :: Assertion+testKonstNotBigCPadded128b =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001,1627.8210700004993,1571.2321300004994,1132.9261600005002,1188.6375200005,675.7488800003999,828.6545600004001,215.6659200003,388.5716000003,2725.0393200008984,1831.7390200008983,2551.139320000898,1660.8390200008987,1903.750080000699,1174.5497800006997,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCPadded128b+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++-- with data t128c++testKonst5LittleBPadded128c :: Assertion+testKonst5LittleBPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [186.7886400001,121.7829600001,269.09432000009997,261.3943200001,210.9943200001,231.79432000010002,160.00030000000004,194.00060000000005,186.7886400001,121.7829600001,269.09432000009997,261.3943200001,210.9943200001,231.79432000010002,160.00030000000004,194.00060000000005])+    (rev' @Double @4 conv2dBPadded128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5LittleCPadded128c :: Assertion+testKonst5LittleCPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2, 2, 2, 2] [1772.649480000399,2138.4267600005987,2157.0438000004983,2640.8154000007976,961.7781800001002,1359.4557500003987,1233.4728000001987,1712.044990000598,1772.649480000399,2138.4267600005987,2157.0438000004983,2640.8154000007976,961.7781800001002,1359.4557500003987,1233.4728000001987,1712.044990000598])+    (rev' @Double @4 conv2dCPadded128c (rreplicate0N [2, 2, 2, 2] (rscalar 5)))++testKonst5BigBPadded128c :: Assertion+testKonst5BigBPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dBPadded128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBPadded128c :: Assertion+testKonstNotBigBPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005,54.100300000000004,111.20060000000001,119.09462,270.29492000000005,109.09462000000002,318.19492,174.08894000000004,477.28924000000006,58.2,140.3,90.4,212.4,120.4,292.39432000000005,-117.5,103.38864000010005])+    (rev' @Double @4 conv2dBPadded128c+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))++testKonst5BigCPadded128c :: Assertion+testKonst5BigCPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dCPadded128c (rreplicate0N [3, 2, 4, 2] (rscalar 5)))++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCPadded128c :: Assertion+testKonstNotBigCPadded128c =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [3, 2, 4, 2] [2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002,2640.8154000007976,1836.3264600007988,2412.414800000798,1662.026160000799,2121.6375200006987,1436.2432000006995,1953.5375200006988,1258.1432000006998,1712.044990000598,1566.644690000599,1445.5506800005985,1358.3503800005992,1279.150680000599,1224.1503800005996,987.1677200004992,962.1674200005002])+    (rev' @Double @4 conv2dCPadded128c+          (rfromList0N [3, 2, 4, 2] (map rscalar [37, 36 .. -10])))+-}++-- with data t128b and t128c+{-+testKonst5LittleBPadded128bc :: Assertion+testKonst5LittleBPadded128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001,1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001])+    (rev' @Double @4 conv2dBPadded128b t128c)++testKonst5LittleCPadded128bc :: Assertion+testKonst5LittleCPadded128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [0.0,0.0,0.0,0.0,251.5006,417.7949200000001,494.8949100000001,382.59461000000005,482.69492000000014,778.9778800001002,952.0721900001002,742.5775700001002,610.5892400000002,1113.1722000001,1412.1551500001997,1182.6605300002002,580.6778800001001,1234.1494800003002,1627.8210700004993,1571.2321300004994,329.17728000010004,816.3545600003002,1132.9261600005002,1188.6375200005,97.98296000010004,455.17160000020016,675.7488800003999,828.6545600004001,-29.9113599999,120.97728000019995,215.6659200003,388.5716000003,0.0,0.0,0.0,0.0,65.90000000000003,106.90000000000003,173.90000000000006,170.90000000000003,365.89432000010004,593.1946200001,821.2892400002003,657.1892400002001,1060.8778800002,1500.2781800001994,1870.0614400004986,1202.8611400005,1465.6665200003995,2156.3671200003987,2725.0393200008984,1831.7390200008983,1399.7665200003996,2049.4671200003986,2551.139320000898,1660.8390200008987,1099.7722000003,1563.1725000002994,1903.750080000699,1174.5497800006997,404.7886400002001,656.0889400002,854.9778800004001,628.8778800004001,0.0,0.0,0.0,0.0,251.5006,417.7949200000001,494.8949100000001,382.59461000000005,482.69492000000014,778.9778800001002,952.0721900001002,742.5775700001002,610.5892400000002,1113.1722000001,1412.1551500001997,1182.6605300002002,580.6778800001001,1234.1494800003002,1627.8210700004993,1571.2321300004994,329.17728000010004,816.3545600003002,1132.9261600005002,1188.6375200005,97.98296000010004,455.17160000020016,675.7488800003999,828.6545600004001,-29.9113599999,120.97728000019995,215.6659200003,388.5716000003,0.0,0.0,0.0,0.0,65.90000000000003,106.90000000000003,173.90000000000006,170.90000000000003,365.89432000010004,593.1946200001,821.2892400002003,657.1892400002001,1060.8778800002,1500.2781800001994,1870.0614400004986,1202.8611400005,1465.6665200003995,2156.3671200003987,2725.0393200008984,1831.7390200008983,1399.7665200003996,2049.4671200003986,2551.139320000898,1660.8390200008987,1099.7722000003,1563.1725000002994,1903.750080000699,1174.5497800006997,404.7886400002001,656.0889400002,854.9778800004001,628.8778800004001])+    (rev' @Double @4 conv2dCPadded128b t128c)++testKonst5BigBPadded128bc :: Assertion+testKonst5BigBPadded128bc =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,8,4] [1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001,1113.1722000001,1412.1551500001997,1182.6605300001997,801.5659100002002,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,1234.1494800003,1627.8210700004993,1571.2321300004992,1047.1431900004,816.3545600003002,1132.9261600005,1188.6375200004998,803.7488800004002,455.17160000019993,675.7488800004002,828.6545600004,577.7659200003001,1500.2781800001994,1870.0614400004986,1202.8611400004997,809.1835600003001,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2156.3671200003987,2725.039320000898,1831.7390200008986,1259.3728000004996,2049.4671200003986,2551.139320000898,1660.8390200008987,1151.3728000004999,1563.172500000299,1903.7500800006983,1174.5497800006997,803.9778800004001])+    (rev' @Double @4 conv2dBPadded128b t128c)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigBPadded128cb :: Assertion+testKonstNotBigBPadded128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002,606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002,606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002,606.6659200002001,754.4545600002001,651.5659200002001,373.6659200002,720.1772800002001,917.0659200003003,749.9716000003002,467.1772800001999,1209.2659200003,1451.1488800004995,884.9545600005001,547.1716000002999,1382.7772800002997,1708.860240000599,1078.4602400006002,708.7829600003,316.58864000010004,552.3716000003001,707.9716000003,538.0829600002,328.18894000010005,579.9722000003001,735.8722000002999,565.9835600002,411.9895400001,634.5784800003,706.4781800003,507.58924000020005,773.5898400001,1016.1790800003,753.2787800002999,550.5898400002])+    (rev' @Double @4 conv2dBPadded128c t128b)++testKonst5BigCPadded128cb :: Assertion+testKonst5BigCPadded128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992])+    (rev' @Double @4 conv2dCPadded128c t128b)++-- The gradient is the same as above, because one argument is the same+-- and convolution is linear.+testKonstNotBigCPadded128cb :: Assertion+testKonstNotBigCPadded128cb =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [4,2,4,4] [720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992,720.1835600001002,1155.4608400002999,1436.2438000003995,1150.0548600004,1222.1722000002,1772.649480000399,2138.4267600005987,1463.1378200005997,1477.3665200002997,2157.0438000004983,2640.8154000007976,1836.3264600007988,1366.1659200002998,1965.6432000004988,2412.414800000798,1662.026160000799,226.3886400001001,671.8832600001001,1012.8665100002999,1078.3665100003,419.3835600001001,961.7781800001002,1359.4557500003987,1310.9554500003997,568.9778800002001,1233.4728000001987,1712.044990000598,1566.644690000599,428.6778800002001,1007.0784800001993,1445.5506800005985,1358.3503800005992])+    (rev' @Double @4 conv2dCPadded128c t128b)+-}+++-- * Disparity and misc++-- | Disparity cost volume.+--+--   Take two arrays of multi channel 2d images, where the first contains+--   left views of the scene and the second contains right views.+--+--   For each pair of images, slice the right image over the left image,+--   and for each offset produce the L1 distance indicating how well+--   correponding+--   multi-channel image elements in the right image match those in the left.+--+--   Described in:+--    Anytime Stereo Image Depth Estimation on Mobile Devices+--    Wang, Lai et al, ICRA 2019+--    https://arxiv.org/abs/1810.11408+--    Section III b).+--+costVolume+  :: forall r target. (ADReady target, GoodScalar r)+  => Int -> Int -> target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+costVolume iStart nCount arrL arrR =+  let [nImgs, nChas, nRows, nCols] = rshape arrL+      shO = [nImgs, nCount, nRows, nCols]+  in rbuild shO $ \[iImg, iDisp, iRow, iCol] ->+       let arrVecL = rbuild (nChas :$: ZSR) $ \[iCha] ->+                       rindex0 arrL [iImg, iCha, iRow, iCol]+           iSrc = iCol - fromIntegral iStart - iDisp+           arrVecR = rbuild [nChas] $ \[iCha] ->+                       rindex0 arrR [iImg, iCha, iRow, iSrc]+       in rsum0 $ rzipWith1 (\xL xR -> abs (xL - xR)) arrVecL arrVecR++test_disparityKonst :: Assertion+test_disparityKonst = do+  let arrL :: ADReady target => target (TKR 4 Double)+      arrL = rreplicate0N [1, 2, 4, 6] (rscalar (-0.2))+      arrR :: ADReady target => target (TKR 4 Double)+      arrR = rreplicate0N [1, 2, 4, 6] (rscalar 0.3)+      arrO = costVolume @Double 0 4 arrL arrR+      arrDL = vjp (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL arrO+      arrDR = vjp (\aR -> costVolume 0 4 (rfromPrimal arrL) aR) arrR arrO+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,4,4,6] [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,1.0,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,1.0,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0])+    arrO+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0,-2.2,-2.8,-3.4,-4.0,-4.0,-4.0])+    arrDL+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0])+   arrDR+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,4,6] [4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0,4.0,4.0,4.0,3.0,2.0,1.0])+    (rev' @Double @4 (costVolume 0 4 arrL) arrR)+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,4,6] [-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0,-2.0])+    (rev' @Double @4 (\aL -> costVolume 0 2 aL arrR) arrL)+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,4,6] [2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0,2.0,2.0,2.0,2.0,2.0,1.0])+    (rev' @Double @4 (costVolume 0 2 arrL) arrR)++test_disparityKonst2 :: Assertion+test_disparityKonst2 = do+  let arrL :: (BaseTensor target, GoodScalar r, Differentiable r) => target (TKR 4 r)+      arrL = ringestData [1, 2, 4, 6] [0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0,0.4,0.4,0.4,1.0,1.0,1.0, 1.7041241452319316,1.21999,0.21355339059327375,0.7867666666666666,0.7331698975466578,0.6964466094067263,1.1,1.1041141452319316,0.42000000000000004,0.3536533905932737,0.78,1.253169897546658,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3,2.808238290463863,1.21999,-0.5672067811865474,0.7867666666666666,1.986339795093316,0.6964466094067263]+      arrR :: (BaseTensor target, GoodScalar r, Differentiable r) => target (TKR 4 r)+      arrR = ringestData [1, 2, 4, 6] [0.2, 0.5, -0.2, 0.0001, 0.44, 0.9, -0.9, 0.00001, -0.22, -0.28, -0.34, -0.40, -0.40,-0.22,-0.28,-0.34, 0.22360679774997896,0.35355339059327373,0.20412414523193154,0.5, -0.35355339059327373,0.16666666666666666,0.17677669529663687,-0.25, -2.808238290463863,-1.21999,-0.5672067811865474,-0.7867666666666666,-1.986339795093316,-0.6964466094067263,2.808238290463863,1.21999,-0.5672067811865474,0.7867666666666666,0.6964466094067263,0.42000000000000004,0.3536533905932737,0.78,1.253169897546658,0.50001,0.42000000000000004,0.2801,0.78,1.1,0.50001,0.42000000000000004,0.2801,0.78]+      arrO = rreplicate0N [1, 4, 4, 6] (rscalar (1 :: Double))+      res1 = rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [4.0,2.0,2.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,2.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,2.0,0.0,0.0,-2.0,0.0,4.0,4.0,2.0,0.0,-4.0,1.0,4.0,4.0,4.0,-4.0,2.0,4.0,2.0]+      res2 = rconcrete $ Nested.rfromListPrimLinear [1,2,4,6] [-4.0,0.0,-4.0,-3.0,-2.0,-1.0,-4.0,-4.0,-4.0,-3.0,-2.0,-1.0,-4.0,-4.0,-4.0,-3.0,-2.0,-1.0,-4.0,-2.0,-4.0,-3.0,-2.0,-1.0,-4.0,-4.0,-4.0,-3.0,-2.0,-1.0,4.0,4.0,-4.0,1.0,-2.0,-1.0,-2.0,3.0,2.0,-1.0,-2.0,-1.0,-2.0,0.0,-2.0,-3.0,-2.0,1.0]+      arrDL :: Concrete (TKR 4 Double)+      arrDL = vjp (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL arrO+      arrDR :: Concrete (TKR 4 Double)+      arrDR = vjp (costVolume 0 4 (rfromPrimal arrL)) arrR arrO+  assertEqualUpToEpsilon 1e-7+    res1+    arrDL+  assertEqualUpToEpsilon 1e-7+    res2+    arrDR+  assertEqualUpToEpsilon' 1e-7+    res1+    (rev' @Double @4 (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL)+  assertEqualUpToEpsilon' 1e-7+    res2+    (rev' @Double @4 (costVolume 0 4 (rfromPrimal arrL)) arrR)++test_disparitySmall :: Assertion+test_disparitySmall = do+  let arrL :: ADReady target => target (TKR 4 Double)+      arrL = ringestData [1, 2, 3, 2] [0.2, 0.5, -0.2, 0.0001, 0.44, 0.9, -0.9, 0.00001, -0.22, -0.28, -0.34, -0.40]+      arrR :: ADReady target => target (TKR 4 Double)+      arrR = ringestData [1, 2, 3, 2] [-0.40,-0.22,-0.28,-0.34, 0.22360679774997896,0.35355339059327373,0.20412414523193154,0.5, -0.35355339059327373,0.16666666666666666,0.17677669529663687,-0.25]+      arrO = costVolume @Double 0 4 arrL arrR+      arrDL = vjp (\aL -> costVolume 0 4 aL (rfromPrimal arrR)) arrL arrO+      arrDR = vjp (\aR -> costVolume 0 4 (rfromPrimal arrL) aR) arrR arrO+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,4,3,2] [1.7041241452319316,1.21999,0.21355339059327375,0.7867666666666666,0.7331698975466578,0.6964466094067263,1.1,1.1041141452319316,0.42000000000000004,0.3536533905932737,0.78,1.253169897546658,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3,1.1,0.50001,0.42000000000000004,0.2801,0.78,1.3])+    arrO+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,3,2] [-2.0,-1.0,-2.0,-1.0,-2.0,-1.0,2.0,1.0,-2.0,1.0,2.0,1.0])+    (rev' @Double @4 (costVolume 0 4 arrL) arrR)+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,2,3,2] [5.004124145231932,3.3241241452319317,-1.0464466094067264,1.7006200572599404,3.0731698975466575,4.5496165069533845,-5.004124145231932,-1.3240841452319316,-1.0464466094067264,-0.9933132760733929,-3.0731698975466575,-4.5496165069533845])+    arrDL+  assertEqualUpToEpsilon 1e-7+    (rconcrete $ Nested.rfromListPrimLinear [1,2,3,2] [-2.808238290463863,-1.21999,-0.5672067811865474,-0.7867666666666666,-1.986339795093316,-0.6964466094067263,2.808238290463863,1.21999,-0.5672067811865474,0.7867666666666666,1.986339795093316,0.6964466094067263])+   arrDR+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,3,2] [-1.0,0.0,-1.0,0.0,-1.0,0.0,1.0,0.0,-1.0,0.0,1.0,0.0])+    (rev' @Double @4 (costVolume 1 4 arrL) arrR)+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,3,2] [2.0,2.0,-2.0,2.0,2.0,2.0,-2.0,2.0,-2.0,-2.0,-2.0,-2.0])+    (rev' @Double @4 (\aL -> costVolume 2 2 aL arrR) arrL)+  assertEqualUpToEpsilon' 1e-7+    (ringestData [1,2,3,2] [-1.0,0.0,-1.0,0.0,-1.0,0.0,1.0,0.0,-1.0,0.0,1.0,0.0])+    (rev' @Double @4 (costVolume 1 2 arrL) arrR)++codeTomsSlice :: ADReady target+              => target (TKR 2 Double) -> target (TKR 0 Double)+codeTomsSlice a =+  let (n, m) = case rshape a of+        [n', m'] -> (n', m')+        _ -> error "codeTomsSlice"+      a1 = rbuild @2 @0 [n,m-1] (\[i',j'] -> rindex0 a [i',j'])+      a2 = rbuild [n,m-1] (\[i',j'] -> rindex0 a [i',j' + 1])+  in rsum0 @2 $ rbuild [n,m] $ \[i, _j] ->+       rfromIndex0 i * rsum0 (a1 * a2)++testTomsSliceRev :: Assertion+testTomsSliceRev = do+  assertEqualUpToEpsilon 1e-5+    (ringestData [32,4] [63686.39999999999,137292.80000000002,121222.4,79558.40000000002,192646.40000000005,223971.0617601984,228556.80000000005,116846.33088019838,63686.39999999999,137292.80000000002,127174.4,79558.40000000002,192646.40000000005,158499.06176019844,202566.40000000005,51374.330880198424,11904.0,5952.0,7936.0,1984.0,116846.33088019838,385292.8000000001,227740.66176039676,192646.40000000005,116846.33088019838,228556.80000000005,174580.73088019836,35910.399999999994,79558.40000000002,127372.79999999997,143244.80000000002,63686.39999999999,105152.0,186683.13088000007,105151.98016,107124.73088000003,-396.79999999999995,26188.8,17459.2,25990.399999999998,-7936.0,73408.0,-1995.2691200000017,57536.0,51584.0,-660672.0,55552.0,3968.0,3968.0,3571.2,3571.2,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,49203.79519999998,49203.79519999998,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,129158.9952,65472.59519999998,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,1984.0000000000146,67059.20000000001,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,-21823.99999999993,108921.6,16070.400000000005,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,771974.4,218019.0617601984,192646.40000000005,170414.3308801984,385292.8000000001,340828.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,236294.40000000005,271587.0617601984,192646.40000000005,45422.33088019842,385292.8000000001,162268.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,369222.4,220003.0617601984,192646.40000000005,104942.33088019838,385292.8000000001,215836.66176039676,192646.40000000005])+    (grad (kfromR . codeTomsSlice) (rreshape [32, 4] t128))++testTomsSlice :: Assertion+testTomsSlice = do+  assertEqualUpToEpsilon' 1e-5+    (ringestData [32,4] [63686.39999999999,137292.80000000002,121222.4,79558.40000000002,192646.40000000005,223971.0617601984,228556.80000000005,116846.33088019838,63686.39999999999,137292.80000000002,127174.4,79558.40000000002,192646.40000000005,158499.06176019844,202566.40000000005,51374.330880198424,11904.0,5952.0,7936.0,1984.0,116846.33088019838,385292.8000000001,227740.66176039676,192646.40000000005,116846.33088019838,228556.80000000005,174580.73088019836,35910.399999999994,79558.40000000002,127372.79999999997,143244.80000000002,63686.39999999999,105152.0,186683.13088000007,105151.98016,107124.73088000003,-396.79999999999995,26188.8,17459.2,25990.399999999998,-7936.0,73408.0,-1995.2691200000017,57536.0,51584.0,-660672.0,55552.0,3968.0,3968.0,3571.2,3571.2,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,49203.79519999998,49203.79519999998,-396.79999999999995,-396.79999999999995,49203.79519999998,49203.79519999998,49600.59519999998,49600.59519999998,129158.9952,65472.59519999998,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,1984.0000000000146,67059.20000000001,79558.40000000002,-5952.0,73198.33087999995,51175.930880000036,51374.33087999995,51187.20000000001,-21823.99999999993,108921.6,16070.400000000005,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,771974.4,218019.0617601984,192646.40000000005,170414.3308801984,385292.8000000001,340828.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,236294.40000000005,271587.0617601984,192646.40000000005,45422.33088019842,385292.8000000001,162268.6617603968,192646.40000000005,57734.399999999994,99596.79999999999,137292.80000000002,63686.39999999999,79558.40000000002,127372.79999999997,159116.80000000005,63686.39999999999,107124.73088000003,369222.4,220003.0617601984,192646.40000000005,104942.33088019838,385292.8000000001,215836.66176039676,192646.40000000005])+    (rev' codeTomsSlice (rreshape [32, 4] t128))+++-- * PP Tests++testTomsSlicePP :: Assertion+testTomsSlicePP = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent codeTomsSlice (FTKR [32, 4] FTKScalar)+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (sscalar 4.0 * sdot0 (sconcrete (sfromListLinear [32] [0.0,1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0,13.0,14.0,15.0,16.0,17.0,18.0,19.0,20.0,21.0,22.0,23.0,24.0,25.0,26.0,27.0,28.0,29.0,30.0,31.0])) (sreplicate @32 (sdot0 (sslice (SNat @0) (SNat @3) (str (sfromR m1))) (sslice (SNat @1) (SNat @3) (str (sfromR m1))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> let v8 = sreplicate @32 (ssum @96 (sreshape @[96] (str (sslice (SNat @0) (SNat @3) (str (sfromR m1))) * str (sslice (SNat @1) (SNat @3) (str (sfromR m1)))))) in rfromS (ssum @128 (sreshape @[128] (str (sreplicate @4 (siota (SNat @32) * v8)))))"+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let m10 = sreshape @[32,3] (sreplicate @96 (ssum @32 (siota (SNat @32) * ssum @4 (str (sreshape @[32,4] (sreplicate @128 (sfromR dret))))))) in rfromS (str (sappend (sconcrete (sfromListLinear [0,32] [])) (sappend (str (str (sslice (SNat @1) (SNat @3) (str (sfromR m1))) * m10)) (sconcrete (sreplicate [1,32] 0.0)))) + str (sappend (sconcrete (sreplicate [1,32] 0.0)) (sappend (str (str (sslice (SNat @0) (SNat @3) (str (sfromR m1))) * m10)) (sconcrete (sfromListLinear [0,32] [])))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> rfromS (let x10 = sdot0 (sconcrete (sfromListLinear [32] [0.0,4.0,8.0,12.0,16.0,20.0,24.0,28.0,32.0,36.0,40.0,44.0,48.0,52.0,56.0,60.0,64.0,68.0,72.0,76.0,80.0,84.0,88.0,92.0,96.0,100.0,104.0,108.0,112.0,116.0,120.0,124.0])) (sreplicate @32 (sfromR dret)) in str (sappend (sslice (SNat @1) (SNat @3) (str (sfromR m1)) * sreplicate @3 (sreplicate @32 x10)) (sconcrete (sreplicate [1,32] 0.0))) + str (sappend (sconcrete (sreplicate [1,32] 0.0)) (sslice (SNat @0) (SNat @3) (str (sfromR m1)) * sreplicate @3 (sreplicate @32 x10))))"++testCNNOPP0c :: Assertion+testCNNOPP0c = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dCLaborious (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i86, i88] -> [i86 + i88]))) (\\[i40, i41] -> [i40 + i41])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i42, i43, i44, i45] -> [ifH (notB (2 <=. i42 + i44) &&* notB (2 <=. i43 + i45)) 0 1, i42, i43, i44, i45])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w46 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i38, i39] -> [i38 + i39]))) (\\[i40, i41] -> [i40 + i41])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i42, i43, i44, i45] -> [ifH (notB (2 <=. i42 + i44) &&* notB (2 <=. i43 + i45)) 0 1, i42, i43, i44, i45])))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w46 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w46 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i38, i39] -> [i38 + i39]))) (\\[i40, i41] -> [i40 + i41])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i42, i43, i44, i45] -> [ifH (notB (2 <=. i42 + i44) &&* notB (2 <=. i43 + i45)) 0 1, i42, i43, i44, i45])))))) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w46 * sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (ssum @2 (ssum @2 (sdot1In (stranspose @[0,1,2,5,3,4] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,4,2] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i112, i114] -> [i112 + i114]))) (\\[i40, i41] -> [i40 + i41])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i96, i97, i98, i103, i104] -> [ifH (notB (2 <=. i96 + i103) &&* notB (2 <=. i97 + i104)) 0 1, i96, i97, i103, i104]))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))) !$ [0]))))"++testCNNOPP0b :: Assertion+testCNNOPP0b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dBLaborious (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i108, i110] -> [i108 + i110]))) (\\[i47, i48] -> [i47 + i48])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i49, i50, i51, i52] -> [ifH (notB (2 <=. i49 + i51) &&* notB (2 <=. i50 + i52)) 0 1, i49, i50, i51, i52]))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w53 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i45, i46] -> [i45 + i46]))) (\\[i47, i48] -> [i47 + i48])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i49, i50, i51, i52] -> [ifH (notB (2 <=. i49 + i51) &&* notB (2 <=. i50 + i52)) 0 1, i49, i50, i51, i52])))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * w53))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w59 = sscatter (stranspose @[1,2,4,5,0,3] (ssum @1 (stranspose @[3,0,1,2] (ssum @2 (str (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))))))))) (\\[i55, i56, i57, i58] -> [ifH (notB (2 <=. i55 + i57) &&* notB (2 <=. i56 + i58)) 0 1, i55, i56, i57, i58]) in rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[1,3,4,0,5,2] (w59 !$ [0])) (\\[i60, i61] -> [i60 + i61]))) (\\[i62, i63] -> [i62 + i63])))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[0,2,4,5,1,6,3] (sscatter (sdot1In (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0,5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0,5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0,5.0,13.1,-2.0,582934.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,2.0,9.0,0.0,2.99432,6.0,8.0,0.1,-335.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,1.0,-4.0,-0.2,26.0])) (stranspose @[4,2,3,6,7,0,5,1] (sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))) !$ [0])) (\\[i55, i56, i57, i58] -> [ifH (notB (2 <=. i55 + i57) &&* notB (2 <=. i56 + i58)) 0 1, i55, i56, i57, i58])) !$ [0]) (\\[i60, i61] -> [i60 + i61]))) (\\[i62, i63] -> [i62 + i63])))"++testCNNOPP1e :: Assertion+testCNNOPP1e = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dUnpaddedL (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i93, i95] -> [i93 + i95]))) (\\[i30, i31] -> [i30 + i31])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i32, i33, i34, i35] -> [ifH (notB (2 <=. i32 + i34) &&* notB (2 <=. i33 + i35)) 0 1, i32, i33, i34, i35])))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w36 = str (sreplicate @2 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i28, i29] -> [i28 + i29]))) (\\[i30, i31] -> [i30 + i31])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i32, i33, i34, i35] -> [ifH (notB (2 <=. i32 + i34) &&* notB (2 <=. i33 + i35)) 0 1, i32, i33, i34, i35])))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w36 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w38 = sreshape @[2,2,2,2,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (stranspose @[4,0,1,5,2,3] (sgather (sfromVector (fromList [stranspose @[3,0,5,1,2,4] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i124, i126] -> [i124 + i126]))) (\\[i30, i31] -> [i30 + i31])), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i32, i33, i34, i35] -> [ifH (notB (2 <=. i32 + i34) &&* notB (2 <=. i33 + i35)) 0 1, i32, i33, i34, i35]))))) (stranspose @[2,3,1,4,5,6,0] w38)))) (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[0,2,4,5,1,6,3] (sscatter (sdot1In (stranspose @[2,3,5,6,0,4,1] (sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1)))))))) (stranspose @[2,3,5,6,0,4,1] w38)) (\\[i39, i40, i41, i42] -> [ifH (notB (2 <=. i39 + i41) &&* notB (2 <=. i40 + i42)) 0 1, i39, i40, i41, i42])) !$ [0]) (\\[i44, i45] -> [i44 + i45]))) (\\[i46, i47] -> [i46 + i47]))))"++testCNNOPP2 :: Assertion+testCNNOPP2 = do+  resetVarCounter+  let t = maxPool2dUnpadded2+            (rconcrete $ Nested.rreplicateScal (1 :$: 1 :$: 2 :$: 2 :$: ZSR) 1)+  printAstPretty (simplifyInlineContract t)+    @?= "rfromS (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sappend (sreplicate @1 (sgather (sreplicate @1 (stranspose @[2,0,1] (sgather (sconcrete (sfromListLinear [2,2] [1.0,1.0,1.0,1.0])) (\\[i68, i69] -> [i69 + i68])))) (\\[i44, i35, i8] -> [i8, i8, i8, 2 * i44 + i35]))) (sconcrete (sreplicate [1,2,2,2] 0.0))) !$ [0, 0])))"+  printAstPretty t+    @?= "rfromS (sreplicate @2 (sreplicate @2 (let u36 = let u41 = sgather (sgather (sreplicate @1 (let w32 = sgather (stranspose @[3,2,0,1] (sgather (sconcrete (sfromListLinear [2,3,2] [1.0,1.0,0.0,0.0,0.0,0.0,1.0,1.0,0.0,0.0,0.0,0.0])) (\\[i26, i15] -> [i26 + i15]))) (\\[i22, i16] -> [i22 + i16]) in stranspose @[1,2,3,0] (sappend (sreplicate @1 (stranspose @[2,0,4,1,3] w32 !$ [0])) (sconcrete (sreplicate [2,2,2,2,2] 0.0))))) (\\[i20] -> [i20, i20, i20, 0])) (\\[i44, i39, i35, i8] -> [2 * i39 + i8, i39, 2 * i44 + i35]) in str (sappend (sreplicate @1 (str u41 !$ [0])) (sconcrete (sreplicate [1,2,2,2] 0.0))) in stranspose @[2,3,0,1] u36 !$ [0, 0])))"++testCNNOPP2b :: Assertion+testCNNOPP2b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded2 (FTKR [1, 1, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sappend (sreplicate @1 (sgather (sreplicate @1 (stranspose @[2,0,1] (sgather (sfromR u1 !$ [0, 0]) (\\[i92, i93] -> [i93 + i92])))) (\\[i94, i95, i96] -> [i96, i96, i96, 2 * i94 + i95]))) (sconcrete (sreplicate [1,2,2,2] 0.0))) !$ [0, 0])))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sappend (sreplicate @1 (sgather (sreplicate @1 (stranspose @[2,0,1] (sgather (sfromR u1 !$ [0, 0]) (\\[i92, i93] -> [i93 + i92])))) (\\[i94, i95, i96] -> [i96, i96, i96, 2 * i94 + i95]))) (sconcrete (sreplicate [1,2,2,2] 0.0))) !$ [0, 0])))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let u98 = stranspose @[3,2,0,1] (soneHot (ssum @2 (ssum @2 (sfromR dret))) [0, 0]) in rfromS (soneHot (sscatter (stranspose @[1,2,0] (ssum @1 (sscatter (ssum @1 (sslice (SNat @0) (SNat @1) u98)) (\\[i99, i100, i101] -> [i101, i101, i101, 2 * i99 + i100])))) (\\[i102, i103] -> [i103 + i102])) [0, 0])"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (sreplicate @1 (sreplicate @1 (sscatter (sscatter (stranspose @[3,2,0,1] (soneHot (ssum @2 (ssum @2 (sfromR dret))) [0, 0]) !$ [0]) (\\[i99, i100, i101] -> [i101, i101, 2 * i99 + i100, i101]) !$ [0]) (\\[i102, i103] -> [i103 + i102]))))"++maxPool2dUnpadded2+  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded2 a =+  rbuild [2, 2, 2, 2] $ \case+    [_, _, iBh, iBw] ->+      let arrt = slicez2 (conv2dUnpadded2 a) [iBw, 1, 2 * iBh, 2 * iBw]+      in rmaximum2 arrt+    _ -> error "maxPool2dUnpadded2: impossible pattern needlessly required"++conv2dUnpadded2+  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded2 a =+  rbuild [3, 3, 2, 2] $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez2 a [iImg, 0, iBh, iBw]+      in rindex0 arrAt [0, iBw, iBw, 0]+    _ -> error "conv2dUnpadded2: impossible pattern needlessly required"++slicez2+  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double, n ~ 4)+  => target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez2 d ixBase =+  rbuild [1, 1, 2, 2] $ \ixResult -> indexz02 d (ixrZipWith (+) ixBase ixResult)++indexz02+  :: forall target r n.+     (target ~ AstTensor AstMethodLet FullSpan, r ~ Double, n ~ 4)+  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)+indexz02 d ix = ifH (1 >. (toList ix !! 0)) (d ! ix) (rscalar 0)++rmaximum2 :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)+         => target (TKR 4 r) -> target (TKR 0 r)+rmaximum2 t0 = tlet t0 $ \t -> rindex0 t [0, 0, 0, 0]++{- TODO: divergent result; bring back when GHC 9.10 dropped:+testCNNOPP3 :: Assertion+testCNNOPP3 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded33 $ conv2dUnpadded3 blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (sreplicate @2 (sgather (stranspose @[2,1,0,4,3] (sappend (sreplicate @1 (sgather (sconcrete (sfromListLinear [2] [7.0,0.0])) (\\[i18, i22, i17, i15] -> [ifH (notB (2 <=. remH i22 4 + i18) &&* (notB (2 <=. i22 + i17) &&* notB (2 <=. i22 + i15))) 0 1]))) (sconcrete (sfromListLinear [1,2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])))) (\\[i52, i51] -> [remH i51 4, i52, i52, remH i51 4])))"+  printAstPretty afcnn2T+    @?= "rfromS (let w30 = sgather (sfromVector (fromList [stranspose @[4,0,1,2,5,3] (sgather (stranspose @[1,2,4,5,0,3] (sgather (sappend (sreplicate @1 (sgather (sconcrete (sfromListLinear [2] [7.0,0.0])) (\\[i18, i22, i17, i15] -> [ifH (notB (2 <=. remH i22 4 + i18) &&* (notB (2 <=. i22 + i17) &&* notB (2 <=. i22 + i15))) 0 1]))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sscalar 0.0))))))) (\\[i43, i38, i29, i7] -> [i43 + i7, i43 + i7, remH i38 4 + i29]))) (\\[i37, i33, i28, i8] -> [i37, i28, i33 + i8, remH i37 4 + i28])), sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sscalar 0.0))))))])) (\\[i46, i41, i36, i32, i27, i26, i24, i23] -> [ifH (notB (2 <=. remH i36 4 + i27) &&* (notB (2 <=. i46 + i26) &&* (notB (2 <=. i41 + i24) &&* notB (2 <=. i32 + i23)))) 0 1, i41, i36, i32, i27, i24, i23]) in stranspose @[4,5,6,7,0,1,2,3] w30 !$ [0, 0, 0, 0])"+-}++testCNNOPP3b :: Assertion+testCNNOPP3b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded33 . conv2dUnpadded3) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[2,1,0] (sgather (sgather (sappend (sreplicate @1 (stranspose @[0,4,5,1,2,3] (sgather (sfromVector (fromList [stranspose @[5,2,3,4,0,1] (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sgather (stranspose @[3,0,2,1] (sgather (stranspose @[3,0,1,2] (sgather (stranspose @[3,0,2,1] (sfromR u1) !$ [1]) (\\[i191, i193] -> [remH i191 4 + i193]))) (\\[i195, i197] -> [i195 + i197, i195]))) (\\[i128, i129] -> [i128 + i129, i128]))))), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i130, i131, i132, i133] -> [ifH (notB (2 <=. remH i130 4 + i131) &&* (notB (2 <=. i130 + i132) &&* notB (2 <=. i130 + i133))) 0 1, i130, i131, i132, i133])))) (sconcrete (sreplicate [1,2,2,2,2,2,2] 0.0))) (\\[i134, i135, i136, i137] -> [i135, i134, i135, i137, i135, i137, i134])) (\\[i138] -> [remH i138 4])))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (stranspose @[2,1,0] (sgather (sgather (sappend (sreplicate @1 (stranspose @[0,4,5,1,2,3] (sgather (sfromVector (fromList [stranspose @[5,2,3,4,0,1] (sreplicate @2 (sreplicate @2 (stranspose @[2,3,1,0] (sgather (stranspose @[3,0,2,1] (sgather (stranspose @[3,0,1,2] (sgather (stranspose @[3,0,2,1] (sfromR u1) !$ [1]) (\\[i124, i125] -> [remH i124 4 + i125]))) (\\[i126, i127] -> [i126 + i127, i126]))) (\\[i128, i129] -> [i128 + i129, i128]))))), sconcrete (sreplicate [2,2,2,2,2,2] 0.0)])) (\\[i130, i131, i132, i133] -> [ifH (notB (2 <=. remH i130 4 + i131) &&* (notB (2 <=. i130 + i132) &&* notB (2 <=. i130 + i133))) 0 1, i130, i131, i132, i133])))) (sconcrete (sreplicate [1,2,2,2,2,2,2] 0.0))) (\\[i134, i135, i136, i137] -> [i135, i134, i135, i137, i135, i137, i134])) (\\[i138] -> [remH i138 4])))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w145 = sscatter (sscatter (stranspose @[2,1,0] (sfromR dret)) (\\[i140] -> [remH i140 4])) (\\[i141, i142, i143, i144] -> [i142, i141, i142, i144, i142, i144, i141]) ; w150 = sscatter (stranspose @[0,3,4,5,1,2] (ssum @1 (sslice (SNat @0) (SNat @1) w145))) (\\[i146, i147, i148, i149] -> [ifH (notB (2 <=. remH i146 4 + i147) &&* (notB (2 <=. i146 + i148) &&* notB (2 <=. i146 + i149))) 0 1, i146, i147, i148, i149]) in rfromS (stranspose @[1,3,2,0] (soneHot (sscatter (stranspose @[1,2,3,0] (sscatter (stranspose @[1,3,2,0] (sscatter (stranspose @[3,2,0,1] (ssum @2 (ssum @2 (stranspose @[4,5,1,2,3,0] (w150 !$ [0]))))) (\\[i151, i152] -> [i151 + i152, i151]))) (\\[i153, i154] -> [i153 + i154, i153]))) (\\[i155, i156] -> [remH i155 4 + i156])) [1]))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[1,3,2,0] (soneHot (sscatter (stranspose @[1,2,3,0] (sscatter (stranspose @[1,3,2,0] (sscatter (ssum @2 (ssum @2 (stranspose @[0,5,6,1,4,2,3] (sscatter (sscatter (sscatter (stranspose @[2,1,0] (sfromR dret)) (\\[i140] -> [remH i140 4])) (\\[i141, i142, i143, i144] -> [i142, i141, i142, i144, i141, i142, i144]) !$ [0]) (\\[i146, i147, i148, i149] -> [ifH (notB (2 <=. remH i146 4 + i147) &&* (notB (2 <=. i146 + i148) &&* notB (2 <=. i146 + i149))) 0 1, i146, i147, i148, i149])) !$ [0]))) (\\[i151, i152] -> [i151 + i152, i151]))) (\\[i153, i154] -> [i153 + i154, i153]))) (\\[i155, i156] -> [remH i155 4 + i156])) [1]))"++maxPool2dUnpadded3+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded3 arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez3 [2, 2, 2, 2] arr [iBh `quotH` 4, aa, bb, iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded3: impossible pattern needlessly required"++maxPool2dUnpadded33+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded33 arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez33 [2, 2, 2, 2] arr [iBh `remH` 4, aa, bb, iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded33: impossible pattern needlessly required"++conv2dUnpadded3+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded3 arrA =+  let shB = [2, 2, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez33 shB arrA [iImg `remH` 4, iImg, iImg, 1]+      in rindex0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dUnpadded3: impossible pattern needlessly required"++slicez3+  :: (ADReady target, GoodScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez3 shOut d ixBase =+  rbuild shOut $ \_ -> indexz03 d (ixrZipWith (+) ixBase ixBase)++slicez33+  :: (ADReady target, GoodScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez33 shOut d ixBase =+  rbuild shOut $ \ixResult -> indexz03 d (ixrZipWith (+) ixBase ixResult)++indexz03+  :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)+  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)+indexz03 d ix = ifH (within0 @target (rshape @target d) ix) (d ! ix) (rscalar 0)++rmaximum3 :: (BaseTensor target, LetTensor target, KnownNat n, GoodScalar r)+         => target (TKR n r) -> target (TKR 0 r)+rmaximum3 t0 = tlet t0 $ \t -> rindex0 t [0, 0, 0, 0]++testCNNOPP4 :: Assertion+testCNNOPP4 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @3) knownSTK+                   $ AstReplicate (SNat @3) knownSTK+                   $ AstReplicate (SNat @3) knownSTK+                   $ AstReplicate (SNat @3) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded4 blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (str (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sreplicate @1 (stranspose @[2,1,0] (sgather (stranspose @[3,4,5,2,6,1,0] (sgather (stranspose @[4,1,3,0,2] (sgather (stranspose @[3,0,4,1,2] (sgather (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i61, i64] -> [i61 + i64]))) (\\[i67, i69] -> [3 + (negate i69 + i67), i69]))) (\\[i71, i73, i76] -> [i71 * i73 + i76])) !$ [1, 0, 0, 0]) (\\[i84] -> [2 * i84]))))))"+      -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,7.0,7.0,7.0,7.0,0.0,0.0,0.0,0.0,7.0,7.0,7.0,7.0]))"+  printAstPretty afcnn2T+    @?= "rfromS (let w19 = sgather (sfromVector (fromList [stranspose @[3,0,5,6,1,2,4] (sgather (stranspose @[6,0,3,1,4,5,2] (sgather (stranspose @[3,0,2,1] (sgather (stranspose @[0,2,1] (sgather (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i32, i5] -> [i32 + i5]))) (\\[i31, i6] -> [i31, 3 + (negate i31 + i6)]))) (\\[i36, i26, i7] -> [i36 * i26 + i7]))) (\\[i22, i8] -> [2 * i22 + i8])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) (\\[i28, i21, i15, i12, i9] -> [ifH (notB (2 <=. i28 + i15) &&* (notB (0 <=. negate i28 + i12) &&* notB (3 <=. 2 * i21 + i9))) 0 1, i28, i21, i15, i12, i9]) in stranspose @[2,3,4,7,5,0,6,1] w19 !$ [0, 0, 0, 0])"++testCNNOPP4b :: Assertion+testCNNOPP4b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded4 (FTKR [3, 3, 3, 3] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (str (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sreplicate @1 (stranspose @[2,1,0] (sgather (stranspose @[2,3,0,1] (sgather (stranspose @[1,0,3,2] (sreplicate @2 (stranspose @[2,3,0,1] (sreplicate @2 (stranspose @[2,1,0] (sreplicate @2 (sfromR u1 !$ [2, 2]))))))) (\\[i194, i195] -> [i195 * i194, i194, i195]))) (\\[i125] -> [i125, 2 * i125]))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (str (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sreplicate @1 (stranspose @[2,1,0] (sgather (stranspose @[3,5,6,2,4,7,1,0] (sgather (stranspose @[3,4,7,1,5,6,0,2] (sgather (stranspose @[6,0,7,4,3,2,1,5] (sgather (sslice (SNat @1) (SNat @2) (stranspose @[0,2,1] (sfromR u1))) (\\[i115, i116, i117, i118, i119] -> [i115 + i116]))) (\\[i120, i121] -> [3 + (negate i121 + i120), i121]))) (\\[i122, i123, i124] -> [i122, i123, i122 * i123 + i124])) !$ [1, 0, 0, 0]) (\\[i125] -> [i125, 2 * i125]))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (stranspose @[0,2,1] (sappend (sconcrete (sreplicate [1,3,3,3] 0.0)) (sappend (sscatter (stranspose @[1,6,5,4,3,7,0,2] (sscatter (stranspose @[6,3,7,0,1,4,5,2] (sscatter (stranspose @[7,6,3,0,4,1,2,5] (soneHot (sscatter (stranspose @[2,1,0] (ssum @1 (sslice (SNat @1) (SNat @1) (str (sfromR dret))))) (\\[i127] -> [i127, 2 * i127])) [1, 0, 0, 0])) (\\[i128, i129, i130] -> [i128, i129, i128 * i129 + i130]))) (\\[i131, i132] -> [3 + (negate i132 + i131), i132]))) (\\[i133, i134, i135, i136, i137] -> [i133 + i134])) (sconcrete (sfromListLinear [0,3,3,3] [])))))"+      -- TODO: was once "\\dret u1 -> rfromS (soneHot (sscatter (ssum @1 (sslice (SNat @1) (SNat @1) (str (sfromR dret)))) (\\[i86, i87, i88] -> [i86 * i87, 2 * i88])) [2, 2])"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (sappend (sconcrete (sreplicate [1,3,3,3] 0.0)) (stranspose @[0,2,1] (sscatter (stranspose @[1,6,5,4,3,7,0,2] (sscatter (stranspose @[6,3,7,0,1,4,5,2] (sscatter (stranspose @[7,6,3,0,4,1,2,5] (soneHot (sscatter (stranspose @[1,3,2,0] (sfromR dret) !$ [1]) (\\[i127] -> [i127, 2 * i127])) [1, 0, 0, 0])) (\\[i128, i129, i130] -> [i128, i129, i128 * i129 + i130]))) (\\[i131, i132] -> [3 + (negate i132 + i131), i132]))) (\\[i133, i134, i135, i136, i137] -> [i133 + i134]))))"+      -- TODO: was once "\\dret u1 -> rfromS (soneHot (sscatter (str (sfromR dret) !$ [1]) (\\[i86, i87, i88] -> [i86 * i87, 2 * i88])) [2, 2])"++testCNNOPP5 :: Assertion+testCNNOPP5 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @6) knownSTK+                   $ AstReplicate (SNat @6) knownSTK+                   $ AstReplicate (SNat @6) knownSTK+                   $ AstReplicate (SNat @6) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = conv2dUnpadded4 blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (sconcrete (sreplicate [1,1,2,2] 7.0))"+  printAstPretty afcnn2T+    @?= "rfromS (sconcrete (sreplicate [1,1,2,2] 7.0))"++testCNNOPP5b :: Assertion+testCNNOPP5b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dUnpadded4 (FTKR [5, 5, 5, 5] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (sreplicate @1 (sreplicate @1 (str (sslice (SNat @0) (SNat @2) (str (sslice (SNat @0) (SNat @2) (sfromR u1 !$ [0, 0])))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (sreplicate @1 (sreplicate @1 (str (sslice (SNat @0) (SNat @2) (str (sslice (SNat @0) (SNat @2) (sfromR u1 !$ [0, 0])))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (soneHot (sappend (sconcrete (sfromListLinear [0,5] [])) (sappend (str (sappend (sconcrete (sfromListLinear [0,2] [])) (sappend (str (ssum @1 (ssum @1 (sfromR dret)))) (sconcrete (sreplicate [3,2] 0.0))))) (sconcrete (sreplicate [3,5] 0.0)))) [0, 0])"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (sappend (str (sappend (stranspose @[0,1,3,2] (sfromR dret) !$ [0, 0]) (sconcrete (sreplicate [3,2] 0.0)))) (sconcrete (sreplicate [3,5] 0.0))) [0, 0])"++maxPool2dUnpadded4+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded4 arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez4 [2, 2, 2, 2] arr [bb + 1, 3 - bb, aa * iBh, 2 * iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded4: impossible pattern needlessly required"++conv2dUnpadded4+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded4 arrA =+  let shB = [1, 1, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez4 shB arrA [iImg, 0, iBh, iBw]+      in rindex0 arrAt [0, 0, 0, 0]+    _ -> error "conv2dUnpadded4: impossible pattern needlessly required"++slicez4+  :: (ADReady target, GoodScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez4 shOut d ixBase =+  rbuild shOut $ \ixResult -> indexz03 d (ixrZipWith (+) ixBase ixResult)++testCNNOPP6 :: Assertion+testCNNOPP6 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded3 $ conv2dUnpadded3z blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printAstPretty afcnn2T+    @?= "rfromS (stranspose @[1,2,0] (sreplicate @2 (let t30 = sgather (stranspose @[2,1,0] (sgather (str (sgather (sreplicate @2 (str (sreplicate @2 (let m21 = sgather (str (sgather (sconcrete (sreplicate [2,2,2,2] 7.0)) (\\[i9] -> [2 * i9, 2 * i9, 2 * i9]))) (\\[i12] -> [2 * i12]) in sappend (sreplicate @1 (sappend (sreplicate @1 (m21 !$ [0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))))) (\\[i1] -> [2 * i1, 0]))) (\\[i2] -> [2 * i2]))) (\\[i4] -> [2 * i4]) in sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (t30 !$ [0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+++testCNNOPP6b :: Assertion+testCNNOPP6b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3 . conv2dUnpadded3z) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let t34 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) in rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t34)))))) [0, 0, 0, 0])"+++  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (ssum0 (stranspose @[0,1,3,2] (sfromR dret) !$ [0, 0, 0])) [0, 0, 0, 0])"++conv2dUnpadded3z+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded3z arrA =+  let shB = [2, 2, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBw]+      in rindex0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dUnpadded3z: impossible pattern needlessly required"++testCNNOPP7 :: Assertion+testCNNOPP7 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded3y $ conv2dUnpadded3y blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printAstPretty afcnn2T+    @?= "rfromS (let u27 = sgather (stranspose @[3,2,0,1] (sgather (stranspose @[1,2,0] (sgather (sreplicate @2 (stranspose @[1,2,0] (sreplicate @2 (let m21 = sgather (str (sgather (sconcrete (sreplicate [2,2,2,2] 7.0)) (\\[i9] -> [2 * i9, 2 * i9, 2 * i9]))) (\\[i11] -> [2 * i11]) in sappend (sreplicate @1 (sappend (sreplicate @1 (m21 !$ [0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))))) (\\[i1] -> [2 * i1]))) (\\[i31, i3] -> [2 * i3, 2 * i31]))) (\\[i4] -> [2 * i4]) in stranspose @[1,2,0] (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (u27 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"++testCNNOPP7b :: Assertion+testCNNOPP7b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3y . conv2dUnpadded3y) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (stranspose @[2,0,1] (sfromR dret)))))))))) [0, 0, 0, 0])"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (sfromR dret !$ [0, 0, 0, 0]) [0, 0, 0, 0])"++maxPool2dUnpadded3y+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded3y arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez3 [2, 2, 2, 2] arr [iBh, aa, bb, iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded3y: impossible pattern needlessly required"++conv2dUnpadded3y+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded3y arrA =+  let shB = [2, 2, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBh]+      in rindex0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dUnpadded3y: impossible pattern needlessly required"++-- TODO: OOMs+_testPaddedCNNOPP0c :: Assertion+_testPaddedCNNOPP0c = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dCPadded (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[5,0,1,4,2,3] (sgather (sconcrete (sfromListLinear [4,4,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i37, i38, i39, i40] -> [i37 + i39, i38 + i40])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w41 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[5,0,1,4,2,3] (sgather (sconcrete (sfromListLinear [4,4,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i37, i38, i39, i40] -> [i37 + i39, i38 + i40])))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w41 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w41 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[5,0,1,4,2,3] (sgather (sconcrete (sfromListLinear [4,4,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,13.1,-2.0,582934.0,2.0,9.0,0.0,2.99432,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,6.0,8.0,0.1,-335.0,1.0,-4.0,-0.2,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i37, i38, i39, i40] -> [i37 + i39, i38 + i40])))))) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w41 * sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (ssum @2 (ssum @2 (sdot1In (stranspose @[0,1,2,6,3,4,5] (sgather (sconcrete (sfromListLinear [4,4,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0,-2.0,13.1,582934.0,2.0,0.0,9.0,2.99432,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,6.0,0.1,8.0,-335.0,1.0,-0.2,-4.0,26.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])) (\\[i66, i67, i68, i73, i74] -> [i66 + i73, i67 + i74]))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))) !$ [0]))))"++-- TODO: OOMs+_testPaddedCNNOPP0b :: Assertion+_testPaddedCNNOPP0b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dBPadded (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR u1)) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))))))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i50, i51, i52, i53, i54, i55] -> [i51 + i54, i51, i54, i50, i52, i53, i55, i50, i53, i52 + i55])))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w56 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR u1)) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))))))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i50, i51, i52, i53, i54, i55] -> [i51 + i54, i51, i54, i50, i52, i53, i55, i50, i53, i52 + i55]))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * w56))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w64 = sscatter (ssum @1 (stranspose @[3,0,1,2] (ssum @2 (str (sconcrete (sfromListLinear [2,2,2,2,1,2,2,2] [5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,5.0,2.0,6.0,1.0,-2.0,0.0,0.1,-0.2,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0,13.1,9.0,8.0,-4.0,582934.0,2.99432,-335.0,26.0]) * sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))))))) (\\[i58, i59, i60, i61, i62, i63] -> [i59 + i62, i59, i62, i58, i60, i61, i63, i58, i61, i60 + i63]) ; u65 = ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (stranspose @[6,1,2,3,4,5,9,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)) (sappend (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @3) (stranspose @[9,3,1,4,5,2,6,7,8,0] (sslice (SNat @1) (SNat @3) w64))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)))))))))) in rfromS (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (sslice (SNat @1) (SNat @3) u65))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (stranspose @[6,1,2,3,4,5,9,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)) (sappend (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @3) (stranspose @[9,3,1,4,5,2,6,7,8,0] (sslice (SNat @1) (SNat @3) (sscatter (sdot1In (sconcrete (sfromListLinear [2,2,2,2,2,2,2] [5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0,5.0,13.1,2.0,9.0,6.0,8.0,1.0,-4.0,-2.0,582934.0,0.0,2.99432,0.1,-335.0,-0.2,26.0])) (stranspose @[4,0,2,3,5,6,7,1] (sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret)))) !$ [0])) (\\[i58, i59, i60, i61, i62, i63] -> [i59 + i62, i59, i62, i58, i60, i61, i63, i58, i61, i60 + i63]))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)))))))))))))))"++-- TODO: OOMs+_testPaddedCNNOPP1e :: Assertion+_testPaddedCNNOPP1e = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dPadded (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))))))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i44, i75, i45, i46, i76, i47] -> [i75 + i76, i75, i76, i44, i45, i46, i47, i44, i46, i45 + i47]))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let u41 = sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0))) ; w48 = str (sreplicate @2 (stranspose @[0,4,1,2,5,3] (sgather (stranspose @[2,3,4,5,6,7,8,0,1] (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 u41)))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i42, i43] -> [i42 + i43, i42, i43]))) (\\[i44, i45, i46, i47] -> [i44, i45, i46, i47, i44, i46, i45 + i47])))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w48 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w50 = sreshape @[2,2,2,2,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (sgather (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (stranspose @[9,2,5,1,3,4,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0)) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @2) (stranspose @[9,1,2,3,4,5,0,7,8,6] (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))))))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,3] 0.0))))))) (\\[i44, i84, i45, i46, i85, i47] -> [i84 + i85, i84, i85, i44, i45, i46, i47, i44, i46, i45 + i47])))) (stranspose @[2,3,1,4,5,6,0] w50)))) (stranspose @[1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (ssum @2 (stranspose @[6,1,2,3,4,5,9,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0)) (sappend (stranspose @[9,1,2,3,4,5,6,7,8,0] (sappend (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0)) (sappend (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @0) (SNat @2) (stranspose @[9,1,2,3,4,5,6,7,8,0] (sslice (SNat @1) (SNat @3) (stranspose @[9,3,1,4,5,2,6,7,8,0] (sslice (SNat @1) (SNat @3) (sscatter (stranspose @[7,8,0,1,2,3,4,5,6] (sscatter (sdot1In (sreplicate @2 (stranspose @[2,3,5,0,4,1] (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))) (stranspose @[0,3,4,6,2,5,1] w50)) (\\[i51, i52, i53, i54] -> [i51, i52, i53, i54, i51, i53, i52 + i54]))) (\\[i55, i56] -> [i55 + i56, i55, i56]))))))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,2,2,2,2,2,2,4] 0.0))))))))))))))))"++-- This is fragile due to indexing out of bounds, see above.+testPaddedCNNOPP1b :: Assertion+testPaddedCNNOPP1b = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dShrinking (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (6 :$: 2 :$: 6 :$: 6 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[6,2,4,4,8] (str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i64, i66] -> [i64 + i66]))) (\\[i24, i25] -> [i24 + i25])))) * sreplicate @6 (str (sreplicate @4 (str (sreplicate @4 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w26 = str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[6,2,4,4,8] (w26 * sreplicate @6 (str (sreplicate @4 (str (sreplicate @4 (sfromR (tproject1 u1))))))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w26 = str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i22, i23] -> [i22 + i23]))) (\\[i24, i25] -> [i24 + i25])))) ; w28 = sreshape @[6,2,4,4,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (rfromS (ssum @4 (str (ssum @4 (str (ssum @6 (w26 * w28))))))) (rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[2,5,0,1,3,4] (ssum @2 (str (sreplicate @6 (str (sreplicate @4 (str (sreplicate @4 (sfromR (tproject1 u1)))))) * w28)))) (\\[i29, i30] -> [i29 + i30]))) (\\[i31, i32] -> [i31 + i32]))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [6,2,6,6] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [6,2,6,6] STKScalar)) (let w28 = sreshape @[6,2,4,4,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @4 (ssum @4 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i81, i83] -> [i81 + i83]))) (\\[i24, i25] -> [i24 + i25]))))) (stranspose @[2,3,1,4,5,6,0] w28)))) (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (sdot1In (stranspose @[3,6,0,2,4,5,1] (sreplicate @6 (str (sreplicate @4 (str (sreplicate @4 (sfromR (tproject1 u1)))))))) (stranspose @[3,6,0,2,4,5,1] w28)) (\\[i29, i30] -> [i29 + i30]))) (\\[i31, i32] -> [i31 + i32]))))"++testPaddedCNNOPPLet :: Assertion+testPaddedCNNOPPLet = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dPaddedLet (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))), sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i36, i37] -> [i36, i37, ifH (notB (2 <=. i37) &&* notB (2 <=. i36)) 0 1])))))) (\\[i81, i83] -> [i81 + i83]))) (\\[i41, i42] -> [i41 + i42])))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let u35 = sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0))) ; u38 = stranspose @[1,2,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] u35, sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i36, i37] -> [i36, i37, ifH (notB (2 <=. i37) &&* notB (2 <=. i36)) 0 1])))))) ; w43 = str (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] u38) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w43 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w45 = sreshape @[2,2,2,2,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))), sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i36, i37] -> [i36, i37, ifH (notB (2 <=. i37) &&* notB (2 <=. i36)) 0 1])))))) (\\[i100, i102] -> [i100 + i102]))) (\\[i41, i42] -> [i41 + i42]))))) (stranspose @[2,3,1,4,5,6,0] w45)))) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,2,1,4,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,2,1,0] (sappend (sconcrete (sreplicate [1,2,2,4,2] 0.0)) (stranspose @[1,3,4,0,2] (sappend (sconcrete (sreplicate [1,3,2,2,2] 0.0)) (sscatter (sslice (SNat @1) (SNat @3) (stranspose @[3,0,1,2] (sslice (SNat @1) (SNat @3) (sscatter (stranspose @[2,4,1,3,0] (sscatter (sdot1In (stranspose @[3,6,0,2,4,5,1] (sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1)))))))) (stranspose @[3,6,0,2,4,5,1] w45)) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49]))))) (\\[i51, i52] -> [i51, i52, ifH (notB (2 <=. i52) &&* notB (2 <=. i51)) 0 1]))))))))) !$ [0]))"++conv2dPaddedLet+  :: forall target r. (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dPaddedLet arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]+      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case+        [iImg, iCinp, iPh, iPw] ->+          ifH (iPh <. fromIntegral (nKh `div` 2)+               ||* iPw <. fromIntegral (nKw `div` 2)+               ||* iPh >=. fromIntegral (nAh + nKh `div` 2)+               ||* iPw >=. fromIntegral (nAw + nKw `div` 2))+              (rscalar 0)+              (arrA ! [ iImg+                      , iCinp+                      , iPh - fromIntegral (nKh `div` 2)+                      , iPw - fromIntegral (nKw `div` 2) ])+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in tlet arrAPadded $ \arrAPadded2 -> rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezL shK1 arrAPadded2 [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rdot0 arrAt arrKt+    _ -> error "conv2dPaddedLet: impossible pattern needlessly required"++testPaddedCNNOPPLet2 :: Assertion+testPaddedCNNOPPLet2 = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dPaddedLet2 (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[3,5,2,0,4,1] (sgather (stranspose @[1,2,0] (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))), sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i53, i54] -> [i53, i54, ifH (notB (2 <=. i54) &&* notB (2 <=. i53)) 0 1])))))))) (\\[i180, i182] -> [i180 + i182]))) (\\[i62, i63] -> [i62, i62 + i63])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR (tproject1 u1))))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let u52 = sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0))) ; u55 = stranspose @[1,2,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] u52, sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i53, i54] -> [i53, i54, ifH (notB (2 <=. i54) &&* notB (2 <=. i53)) 0 1])))))) ; w64 = str (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[3,5,2,0,4,1] (sgather (stranspose @[1,2,0] (sgather (stranspose @[1,5,0,2,3,4] (sreplicate @2 (stranspose @[3,0,1,2] (sgather (stranspose @[2,3,0,1] (sreplicate @2 u55)) (\\[i56, i57, i58] -> [i58, i56]))))) (\\[i59] -> [i59, i59]))) (\\[i60, i61] -> [i60, i60 + i61]))) (\\[i62, i63] -> [i62, i62 + i63])))))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w64 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR (tproject1 u1))))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w66 = sreshape @[2,2,2,2,1,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[5,0,1,4,3,2] (sgather (stranspose @[2,4,0,3,1] (sgather (str (sreplicate @2 (sappend (sconcrete (sreplicate [1,2,4,2] 0.0)) (stranspose @[3,2,0,1] (sappend (sconcrete (sreplicate [1,2,2,3] 0.0)) (stranspose @[0,2,3,1] (sgather (sslice (SNat @1) (SNat @3) (stranspose @[3,0,4,1,2] (sslice (SNat @1) (SNat @3) (stranspose @[3,1,2,4,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))), sconcrete (sreplicate [2,2,4,4] 0.0)])))))) (\\[i53, i54] -> [i53, i54, ifH (notB (2 <=. i54) &&* notB (2 <=. i53)) 0 1])))))))) (\\[i253, i255] -> [i255 + i253]))) (\\[i222, i223, i229] -> [i222, i222 + i229]))) (stranspose @[4,2,3,1,5,6,7,0] w66 !$ [0])))) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,2,1,4,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,2,1,0] (sappend (sconcrete (sreplicate [1,2,2,4,2] 0.0)) (stranspose @[1,3,4,0,2] (sappend (sconcrete (sreplicate [1,3,2,2,2] 0.0)) (sscatter (sslice (SNat @1) (SNat @3) (stranspose @[3,0,1,2] (sslice (SNat @1) (SNat @3) (ssum @2 (stranspose @[2,1,3,0] (sscatter (ssum @2 (stranspose @[2,3,4,5,0,1] (sscatter (stranspose @[2,0,1] (sscatter (stranspose @[3,5,2,0,4,1] (sscatter (sdot1In (sreplicate @2 (stranspose @[2,0,1,4,3] (sreplicate @2 (sreplicate @2 (stranspose @[3,2,1,0] (sfromR (tproject1 u1))))))) (stranspose @[4,3,7,0,2,5,6,1] w66 !$ [0])) (\\[i67, i68] -> [i67, i67 + i68]))) (\\[i69, i70] -> [i69, i69 + i70]))) (\\[i71] -> [i71, i71])))) (\\[i72, i73, i74] -> [i74, i72]))))))) (\\[i76, i77] -> [i76, i77, ifH (notB (2 <=. i77) &&* notB (2 <=. i76)) 0 1]))))))))) !$ [0]))"++conv2dPaddedLet2+  :: forall target r. (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dPaddedLet2 arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]+      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case+        [iImg, iCinp, iPh, iPw] ->+          ifH (iPh <. fromIntegral (nKh `div` 2)+               ||* iPw <. fromIntegral (nKw `div` 2)+               ||* iPh >=. fromIntegral (nAh + nKh `div` 2)+               ||* iPw >=. fromIntegral (nAw + nKw `div` 2))+              (rscalar 0)+              (arrA ! [ iImg+                      , iCinp+                      , iPh - fromIntegral (nKh `div` 2)+                      , iPw - fromIntegral (nKw `div` 2) ])+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = tlet arrAPadded $ \arrAPadded2 ->+                    slicezL shK1 arrAPadded2 [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rdot0 arrAt arrKt+    _ -> error "conv2dPaddedLet2: impossible pattern needlessly required"++-- TODO: OOMs+_testPaddedCNNOPP2 :: Assertion+_testPaddedCNNOPP2 = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = conv2dPadded2 (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+                       (FTKR (2 :$: 2 :$: 2 :$: 2 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (str (sreplicate @2 (stranspose @[5,0,1,4,2,3] (sgather (sappend (sconcrete (sreplicate [1,4,2,2] 0.0)) (sappend (stranspose @[3,0,2,1] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,4,2,2] 0.0)))) (\\[i24, i25, i27, i28] -> [i25 + i28, i24 + i27])))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w29 = str (sreplicate @2 (sgather (sappend (sconcrete (sreplicate [1,2,2,4] 0.0)) (sappend (stranspose @[3,1,2,0] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,2,2,4] 0.0)))) (\\[i23, i24, i25, i26, i27, i28] -> [i25 + i28, i23, i26, i24 + i27]))) in rfromS (ssum @8 (stranspose @[4,0,1,2,3] (sreshape @[2,2,2,2,8] (w29 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (sfromR (tproject1 u1))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,2,2,2] FTKScalar)) ConvSX))) (STKProduct (STKS [2,2,2,2] STKScalar) (STKS [2,2,2,2] STKScalar)) (let w31 = sreshape @[2,2,2,2,2,2,2] (stranspose @[1,2,3,4,0] (sreplicate @8 (sfromR dret))) in tpair (ssum @2 (ssum @2 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @2 (stranspose @[5,0,1,4,2,3] (sgather (sappend (sconcrete (sreplicate [1,4,2,2] 0.0)) (sappend (stranspose @[3,0,2,1] (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sappend (stranspose @[2,0,1] (sfromR (tproject2 u1))) (sconcrete (sreplicate [1,2,2,2] 0.0))))) (sconcrete (sreplicate [1,4,2,2] 0.0)))) (\\[i24, i25, i27, i28] -> [i25 + i28, i24 + i27]))))) (stranspose @[2,3,1,4,5,6,0] w31)))) (stranspose @[1,2,0] (sslice (SNat @1) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @1) (SNat @2) (sscatter (sdot1In (sreplicate @2 (sreplicate @2 (sreplicate @2 (stranspose @[1,2,3,0] (sfromR (tproject1 u1)))))) (stranspose @[0,2,3,4,5,6,1] w31)) (\\[i32, i33, i34, i35, i36, i37] -> [i34 + i37, i32, i35, i33 + i36])))))))"++conv2dPadded2+  :: forall target r. (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r) -> target (TKR 4 r)+conv2dPadded2 arrK arrA =+  let [nImgs, nCinpA, nAh, nAw] = rshape arrA+      [nCoutK, nCinpK, nKh, nKw] = rshape arrK+      shAPadded = [nImgs, nCinpA, nAh + nKh, nAw + nKw]+      arrAPadded = rbuild @4 @0 @(TKScalar r) @target shAPadded $ \case+        [iImg, iCinp, iPh, iPw] ->+               arrA ! [ iImg+                      , iCinp+                      , iPh - fromIntegral (nKh `div` 2)+                      , iPw - fromIntegral (nKw `div` 2) ]+      nCinp = assert (nCinpA == nCinpK `blame` (nCinpA, nCinpK)) nCinpA+      shB = [nImgs, nCoutK, nAh, nAw]+      shK1 = [1, nCinp, nKh, nKw]+  in rbuild shB $ \case+    [iImg, iCout, iBh, iBw] ->+      let arrAt = slicezL shK1 arrAPadded [iImg, 0, iBh, iBw]+          arrKt = slicezL shK1 arrK [iCout, 0, 0, 0]+      in rdot0 arrAt arrKt+    _ -> error "conv2dPadded2: impossible pattern needlessly required"+++-- * Non-laborious CNN PP tests++-- Convolution differentiated wrt the kernel.+testCNNOPP0cW :: Assertion+testCNNOPP0cW = do+  resetVarCounter+  let ftk = FTKR (7 :$: 5 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      ftk2 = FTKR (5 :$: 5 :$: 5 :$: 5 :$: ZSR) (FTKScalar @Double)+      f = simplifyInline . flip conv2dUnpadded var+      env =+        extendEnv varName (dDnotShared (AstRaw var) (DeltaZero ftk)) emptyEnv+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f env) ftk2+  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u0 -> \\u1 -> rfromS (ssum @125 (stranspose @[4,0,1,2,3] (sreshape @[7,5,7,7,125] (str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u0)) (\\[i81, i83] -> [i81 + i83]))) (\\[i41, i42] -> [i41 + i42])))) * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u1)))))))))"+  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev+    @?= "\\u0 -> \\u1 -> let w43 = str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u0)) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) in rfromS (ssum @125 (stranspose @[4,0,1,2,3] (sreshape @[7,5,7,7,125] (w43 * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u1)))))))))"+  "\\u0 -> " ++ printArtifactPretty artifactRev+    @?= "\\u0 -> \\dret u1 -> let w43 = str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u0)) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) ; w45 = sreshape @[7,5,7,7,5,5,5] (stranspose @[1,2,3,4,0] (sreplicate @125 (sfromR dret))) in rfromS (ssum @7 (str (ssum @7 (str (ssum @7 (w43 * w45))))))"+  "\\u0 -> " ++ printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\u0 -> \\dret u1 -> rfromS (ssum @7 (ssum @7 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u0)) (\\[i98, i100] -> [i98 + i100]))) (\\[i41, i42] -> [i41 + i42]))))) (stranspose @[2,3,1,4,5,6,0] (sreshape @[7,5,7,7,5,5,5] (stranspose @[1,2,3,4,0] (sreplicate @125 (sfromR dret))))))))"++-- Convolution differentiated wrt the data.+testCNNOPP0bW :: Assertion+testCNNOPP0bW = do+  resetVarCounter+  let ftk = FTKR (5 :$: 5 :$: 5 :$: 5 :$: ZSR) (FTKScalar @Double)+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      ftk2 = FTKR (7 :$: 5 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)+      f = simplifyInline . conv2dUnpadded var+      env =+        extendEnv varName (dDnotShared (AstRaw var) (DeltaZero ftk)) emptyEnv+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f env) ftk2+  "\\u0 -> " ++ printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u0 -> \\u1 -> rfromS (ssum @125 (stranspose @[4,0,1,2,3] (sreshape @[7,5,7,7,125] (str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i64, i66] -> [i64 + i66]))) (\\[i41, i42] -> [i41 + i42])))) * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u0)))))))))"+  "\\u0 -> " ++ printArtifactPrimalPretty artifactRev+    @?= "\\u0 -> \\u1 -> let w43 = str (sreplicate @5 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) in rfromS (ssum @125 (stranspose @[4,0,1,2,3] (sreshape @[7,5,7,7,125] (w43 * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u0)))))))))"+  "\\u0 -> " ++ printArtifactPretty artifactRev+    @?= "\\u0 -> \\dret u1 -> let w45 = sreshape @[7,5,7,7,5,5,5] (stranspose @[1,2,3,4,0] (sreplicate @125 (sfromR dret))) in rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[2,5,0,1,3,4] (ssum @5 (str (sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u0))))) * w45)))) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49])))"+  "\\u0 -> " ++ printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\u0 -> \\dret u1 -> rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (sdot1In (stranspose @[3,6,0,2,4,5,1] (sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR u0))))))) (stranspose @[3,6,0,2,4,5,1] (sreshape @[7,5,7,7,5,5,5] (stranspose @[1,2,3,4,0] (sreplicate @125 (sfromR dret)))))) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49])))"++testCNNOPP1bW :: Assertion+testCNNOPP1bW = do+  resetVarCounter+  let f :: AstTensor AstMethodLet FullSpan+                     (TKProduct (TKR 4 Double) (TKR 4 Double))+        -> AstTensor AstMethodLet FullSpan+                     (TKR 4 Double)+      f v = simplifyInline $ conv2dUnpadded (tproject1 v) (tproject2 v)+      ftk = FTKProduct (FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) FTKScalar)+                       (FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) FTKScalar)+      (artifactRev, _) =+        revArtifactFromForwardPass+          UseIncomingCotangent (forwardPassByInterpretation f emptyEnv) ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (ssum @343 (stranspose @[4,0,1,2,3] (sreshape @[7,7,7,7,343] (str (sreplicate @7 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i81, i83] -> [i81 + i83]))) (\\[i41, i42] -> [i41 + i42])))) * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR (tproject1 u1))))))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w43 = str (sreplicate @7 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) in rfromS (ssum @343 (stranspose @[4,0,1,2,3] (sreshape @[7,7,7,7,343] (w43 * sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR (tproject1 u1))))))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w43 = str (sreplicate @7 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i39, i40] -> [i39 + i40]))) (\\[i41, i42] -> [i41 + i42])))) ; w45 = sreshape @[7,7,7,7,7,7,7] (stranspose @[1,2,3,4,0] (sreplicate @343 (sfromR dret))) in tpair (rfromS (ssum @7 (str (ssum @7 (str (ssum @7 (w43 * w45))))))) (rfromS (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (stranspose @[2,5,0,1,3,4] (ssum @7 (str (sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR (tproject1 u1)))))) * w45)))) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49]))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [7,7,7,7] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [7,7,7,7] FTKScalar)) ConvSX))) (STKProduct (STKS [7,7,7,7] STKScalar) (STKS [7,7,7,7] STKScalar)) (let w45 = sreshape @[7,7,7,7,7,7,7] (stranspose @[1,2,3,4,0] (sreplicate @343 (sfromR dret))) in tpair (ssum @7 (ssum @7 (sdot1In (stranspose @[2,3,0,4,5,6,1] (sreplicate @7 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (stranspose @[2,0,1] (sfromR (tproject2 u1))) (\\[i98, i100] -> [i98 + i100]))) (\\[i41, i42] -> [i41 + i42]))))) (stranspose @[2,3,1,4,5,6,0] w45)))) (stranspose @[1,2,0] (sscatter (stranspose @[2,4,1,3,0] (sscatter (sdot1In (stranspose @[3,6,0,2,4,5,1] (sreplicate @7 (str (sreplicate @7 (str (sreplicate @7 (sfromR (tproject1 u1)))))))) (stranspose @[3,6,0,2,4,5,1] w45)) (\\[i46, i47] -> [i46 + i47]))) (\\[i48, i49] -> [i48 + i49]))))"++testCNNOPP4bW :: Assertion+testCNNOPP4bW = do+  resetVarCounter+  let !artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2) (FTKR [7, 7, 7, 7] (FTKScalar @Double))+      !artSimp = simplifyArtifact artifactRev+  let ftk1 = FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)+      ftkDt = FTKR (7 :$: 7 :$: 3 :$: 3 :$: ZSR) (FTKScalar @Double)+      env = extendEnv (artVarDtRev artSimp)+                      (tconcrete ftkDt (treplTarget 7 ftkDt))+            $ extendEnv (artVarDomainRev artSimp)+                        (tconcrete ftk1 (treplTarget 42 ftk1)) emptyEnv+  interpretAstPrimal @Concrete env (artPrimalRev artifactRev)+    @?= interpretAstPrimal @Concrete env (artPrimalRev artSimp)+  interpretAstPrimal @Concrete env (artDerivativeRev artifactRev)+    @?= interpretAstPrimal @Concrete env (artDerivativeRev artSimp)+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i137, i138] -> [2 * i137 + i138]))) (\\[i45, i46] -> [2 * i45 + i46]))) in sgather w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kfromS (smaxIndex (w47 !$ [i48, i49, i50, i51]))]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) in rfromS (sgather w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kfromS (smaxIndex (w47 !$ [i48, i49, i50, i51]))]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) in rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (w47 !$ [i53, i54, i55, i56]))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"+  -- The remH comes from the indexing of reshape rule and it looks terrible,+  -- but w42 looks even worse, depending on available primitives,+  -- so the rule is probably fine.+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (sgather (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i159, i160] -> [2 * i159 + i160]))) (\\[i45, i46] -> [2 * i45 + i46])) (\\[i151] -> [remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 16) 3, remH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 4, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 1008) 7, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 144) 7, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 48) 3, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i151) 4) 4])))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"+  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)+    @?= "rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR u52) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) !$ [i53, i54, i55, i56]))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"++testCNNOPP4bD :: Assertion+testCNNOPP4bD = do+  resetVarCounter+  setTotalSharing True+  let !artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2) (FTKR [7, 7, 7, 7] (FTKScalar @Double))+      !artSimp = simplifyArtifact artifactRev+  setTotalSharing False+  let ftk1 = FTKR (7 :$: 7 :$: 7 :$: 7 :$: ZSR) (FTKScalar @Double)+      ftkDt = FTKR (7 :$: 7 :$: 3 :$: 3 :$: ZSR) (FTKScalar @Double)+      env = extendEnv (artVarDtRev artSimp)+                      (tconcrete ftkDt (treplTarget 7 ftkDt))+            $ extendEnv (artVarDomainRev artSimp)+                        (tconcrete ftk1 (treplTarget 42 ftk1)) emptyEnv+  interpretAstPrimal @Concrete env (artPrimalRev artifactRev)+    @?= interpretAstPrimal @Concrete env (artPrimalRev artSimp)+  interpretAstPrimal @Concrete env (artDerivativeRev artifactRev)+    @?= interpretAstPrimal @Concrete env (artDerivativeRev artSimp)+  printArtifactPrimalPretty artSimp+    @?= "\\u1 -> rfromS (let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i98, i99] -> [2 * i98 + i99]))) (\\[i45, i46] -> [2 * i45 + i46]))) in sgather w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kfromS (smaxIndex (w47 !$ [i48, i49, i50, i51]))]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) in rfromS (sgather w47 (\\[i48, i49, i50, i51] -> [i48, i49, i50, i51, kfromS (smaxIndex (w47 !$ [i48, i49, i50, i51]))]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w47 = sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) in rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (w47 !$ [i53, i54, i55, i56]))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"+  printArtifactPretty artSimp+    @?= "\\dret u1 -> rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR dret) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (sgather (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i81, i82] -> [2 * i81 + i82]))) (\\[i45, i46] -> [2 * i45 + i46])) (\\[i73] -> [remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 16) 3, remH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 4, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 1008) 7, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 144) 7, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 48) 3, remH (quotH ((((1008 * i53 + 144 * i54) + 48 * i55) + 16 * i56) + i73) 4) 4])))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"+  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)+    @?= "rfromS (stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (stranspose @[3,5,0,1,2,4] (sreshape @[7,7,3,3,4,4] (sscatter (sfromR u52) (\\[i53, i54, i55, i56] -> [i53, i54, i55, i56, kfromS (smaxIndex (sreshape @[7,7,3,3,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] (sfromR u1)) (\\[i43, i44] -> [2 * i43 + i44]))) (\\[i45, i46] -> [2 * i45 + i46]))) !$ [i53, i54, i55, i56]))])))) (\\[i57, i58] -> [2 * i57 + i58]))) (\\[i59, i60] -> [2 * i59 + i60])))"++testCNNOPP5aW :: Assertion+testCNNOPP5aW = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2 . conv2dC) (FTKR [7, 2, 7, 7] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (let t49 = sreshape @[2,7,16] (ssum @98 (stranspose @[4,1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i109, i111] -> [i109 + i111]))) (\\[i46, i47] -> [i46 + i47])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))))) in stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (sgather t49 (\\[i50, i51] -> [i50, i51, kfromS (smaxIndex (t49 !$ [i50, i51]))]))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w48 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i44, i45] -> [i44 + i45]))) (\\[i46, i47] -> [i46 + i47])))))) ; t49 = sreshape @[2,7,16] (ssum @98 (stranspose @[4,1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (w48 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))))) in rfromS (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (sgather t49 (\\[i50, i51] -> [i50, i51, kfromS (smaxIndex (t49 !$ [i50, i51]))]))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w48 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i44, i45] -> [i44 + i45]))) (\\[i46, i47] -> [i46 + i47])))))) ; t49 = sreshape @[2,7,16] (ssum @98 (stranspose @[4,1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (w48 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))))) ; w55 = stranspose @[4,1,2,3,0] (sreplicate @98 (sreshape @[2,7,4,4] (sscatter (ssum @1 (stranspose @[2,0,1] (ssum @1 (stranspose @[2,0,1] (sfromR dret))))) (\\[i53, i54] -> [i53, i54, kfromS (smaxIndex (t49 !$ [i53, i54]))])))) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w48 * sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) w55))))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (let w48 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i155, i157] -> [i155 + i157]))) (\\[i46, i47] -> [i46 + i47]) in ssum @2 (ssum @2 (sdot1In (stranspose @[1,6,0,4,5,3,2] (sreplicate @7 (stranspose @[3,2,1,4,5,0] w48))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[4,1,2,3,0] (sreplicate @98 (sreshape @[2,7,4,4] (sscatter (stranspose @[2,3,0,1] (sfromR dret) !$ [0, 0]) (\\[i53, i54] -> [i53, i54, kfromS (smaxIndex (ssum @98 (str (sgather (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w48)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))) (\\[i129] -> [remH ((112 * i53 + 16 * i54) + i129) 4, remH (quotH ((112 * i53 + 16 * i54) + i129) 112) 2, remH (quotH ((112 * i53 + 16 * i54) + i129) 16) 7, remH (quotH ((112 * i53 + 16 * i54) + i129) 4) 4])))))])))))))))) !$ [0]))))"+  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)+    @?= "rfromS (let w48 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i44, i45] -> [i44 + i45]))) (\\[i46, i47] -> [i46 + i47]) in ssum @2 (ssum @2 (sdot1In (stranspose @[4,2,3,0,5,6,7,1] (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w48)))) !$ [0]) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[4,1,2,3,0] (sreplicate @98 (sreshape @[2,7,4,4] (sscatter (stranspose @[2,3,0,1] (sfromR u52) !$ [0, 0]) (\\[i53, i54] -> [i53, i54, kfromS (smaxIndex (sreshape @[2,7,16] (ssum @98 (stranspose @[4,1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w48)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) (sconcrete (sreplicate [2,2,7,2,98] 0.0)))) (sconcrete (sreplicate [2,2,7,4,98] 0.0))))) !$ [i53, i54]))])))))))))) !$ [0]))))"++testCNNOPP5bW :: Assertion+testCNNOPP5bW = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2 . relu) (FTKR [7, 2, 7, 7] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (let w75 = sreshape @[7,2,3,3,16] (stranspose @[4,5,0,1,2,3] (sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i65, i67]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i66, i68])])) * stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i69, i71]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i70, i72])]))) in sgather w75 (\\[i76, i77, i78, i79] -> [i76, i77, i78, i79, kfromS (smaxIndex (w75 !$ [i76, i77, i78, i79]))]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let m59 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; m60 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; w73 = stranspose @[4,5,0,1,2,3] (sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (m59 !$ [i65, i67]), kfromS (m60 !$ [i66, i68])])) ; w74 = stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (m59 !$ [i69, i71]), kfromS (m60 !$ [i70, i72])])) ; w75 = sreshape @[7,2,3,3,16] (w73 * w74) in rfromS (sgather w75 (\\[i76, i77, i78, i79] -> [i76, i77, i78, i79, kfromS (smaxIndex (w75 !$ [i76, i77, i78, i79]))]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let m59 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; m60 = str (sreplicate @4 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @4)) ; w73 = stranspose @[4,5,0,1,2,3] (sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (m59 !$ [i65, i67]), kfromS (m60 !$ [i66, i68])])) ; w74 = stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (m59 !$ [i69, i71]), kfromS (m60 !$ [i70, i72])])) ; w75 = sreshape @[7,2,3,3,16] (w73 * w74) in rfromS (stranspose @[2,3,0,1] (sscatter (stranspose @[2,3,4,5,0,1] (w73 * sreshape @[7,2,3,3,4,4] (sscatter (sfromR dret) (\\[i81, i82, i83, i84] -> [i81, i82, i83, i84, kfromS (smaxIndex (w75 !$ [i81, i82, i83, i84]))])))) (\\[i85, i86, i87, i88] -> [kfromS (m59 !$ [i85, i87]), kfromS (m60 !$ [i86, i88])])))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (let w73 = sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i65, i67]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i66, i68])]) in stranspose @[2,3,0,1] (sscatter (w73 * stranspose @[2,3,4,5,0,1] (sreshape @[7,2,3,3,4,4] (sscatter (sfromR dret) (\\[i81, i82, i83, i84] -> [i81, i82, i83, i84, kfromS (smaxIndex (sgather (stranspose @[4,5,0,1,2,3] w73 * stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i69, i71]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i70, i72])]))) (\\[i118] -> [remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 288) 7, remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 144) 2, remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 48) 3, remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 16) 3, remH (quotH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 4) 4, remH ((((288 * i81 + 144 * i82) + 48 * i83) + 16 * i84) + i118) 4])))])))) (\\[i85, i86, i87, i88] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i85, i87]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i86, i88])])))"+  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)+    @?= "rfromS (let w73 = sgather (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i61, i62, i63, i64] -> [ifH (sscalar -0.0 <=. negate (sfromR u1 !$ [i63, i64, i61, i62])) 0 1])) (\\[i65, i66, i67, i68] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i65, i67]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i66, i68])]) in stranspose @[2,3,0,1] (sscatter (w73 * stranspose @[2,3,4,5,0,1] (sreshape @[7,2,3,3,4,4] (sscatter (sfromR u80) (\\[i81, i82, i83, i84] -> [i81, i82, i83, i84, kfromS (smaxIndex (sreshape @[7,2,3,3,16] (stranspose @[4,5,0,1,2,3] w73 * stranspose @[4,5,0,1,2,3] (sgather (stranspose @[2,3,0,1] (sfromR u1)) (\\[i69, i70, i71, i72] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i69, i71]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i70, i72])]))) !$ [i81, i82, i83, i84]))])))) (\\[i85, i86, i87, i88] -> [kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i85, i87]), kfromS (sconcrete (sfromListLinear [3,4] [0,1,2,3,2,3,4,5,4,5,6,7]) !$ [i86, i88])])))"++testCNNOPP5cW :: Assertion+testCNNOPP5cW = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (relu . conv2dC) (FTKR [7, 2, 7, 7] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (let u45 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i114, i116] -> [i114 + i116]))) (\\[i42, i43] -> [i42 + i43])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) in sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (u45 !$ [i46, i47, i48, i49])) 0 1]) * u45)"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w44 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i40, i41] -> [i40 + i41]))) (\\[i42, i43] -> [i42 + i43])))))) ; u45 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (w44 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u50 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (u45 !$ [i46, i47, i48, i49])) 0 1]) in rfromS (u50 * u45)"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w44 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i40, i41] -> [i40 + i41]))) (\\[i42, i43] -> [i42 + i43])))))) ; u45 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (w44 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u50 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (u45 !$ [i46, i47, i48, i49])) 0 1]) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w44 * sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (u50 * sfromR dret)))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (let w44 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i169, i171] -> [i169 + i171]))) (\\[i42, i43] -> [i42 + i43]) in ssum @2 (ssum @2 (sdot1In (stranspose @[1,6,0,4,5,3,2] (sreplicate @7 (stranspose @[3,2,1,4,5,0] w44))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (ssum0 (sgather (stranspose @[6,1,0,2,4,5,3] (sreplicate @7 (stranspose @[3,2,1,4,5,0] w44)) * sreplicate @2 (sreplicate @2 (str (sreplicate @2 (sfromR u1))))) (\\[i132] -> [remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 98) 2, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 196) 2, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 392) 7, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 2744) 2, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 49) 2, remH (quotH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 7) 7, remH ((((2744 * i46 + 392 * i47) + 196 * i48) + 98 * i49) + i132) 7])))) 0 1]) * sfromR dret)))) !$ [0]))))"+  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)+    @?= "rfromS (let w44 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i40, i41] -> [i40 + i41]))) (\\[i42, i43] -> [i42 + i43]) in ssum @2 (ssum @2 (sdot1In (stranspose @[4,2,3,0,5,6,7,1] (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w44)))) !$ [0]) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i46, i47, i48, i49] -> [ifH (sscalar -0.0 <=. negate (ssum0 (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w44)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))) !$ [i46, i47, i48, i49]))) 0 1]) * sfromR u51)))) !$ [0]))))"++testCNNOPP5dW :: Assertion+testCNNOPP5dW = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded 4 2 . relu . conv2dC) (FTKR [7, 2, 7, 7] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (let u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i130, i132] -> [i130 + i132]))) (\\[i64, i65] -> [i64 + i65])))))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; t74 = sreshape @[2,7,16] (stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) * stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)))) in stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (sgather t74 (\\[i75, i76] -> [i75, i76, kfromS (smaxIndex (t74 !$ [i75, i76]))]))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w66 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i62, i63] -> [i62 + i63]))) (\\[i64, i65] -> [i64 + i65])))))) ; u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (w66 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u72 = stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; u73 = stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; t74 = sreshape @[2,7,16] (u72 * u73) in rfromS (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (sgather t74 (\\[i75, i76] -> [i75, i76, kfromS (smaxIndex (t74 !$ [i75, i76]))]))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w66 = str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] (sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i62, i63] -> [i62 + i63]))) (\\[i64, i65] -> [i64 + i65])))))) ; u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (w66 * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u72 = stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; u73 = stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0))) ; t74 = sreshape @[2,7,16] (u72 * u73) ; u80 = stranspose @[3,0,1,2] (u72 * sreshape @[2,7,4,4] (sscatter (ssum @1 (stranspose @[2,0,1] (ssum @1 (stranspose @[2,0,1] (sfromR dret))))) (\\[i78, i79] -> [i78, i79, kfromS (smaxIndex (t74 !$ [i78, i79]))]))) in rfromS (ssum @1 (str (ssum @2 (str (ssum @2 (str (ssum @2 (w66 * sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) u80))))))))))))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (let w66 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i172, i174] -> [i172 + i174]))) (\\[i64, i65] -> [i64 + i65]) ; u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w66)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u72 = sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)) in ssum @2 (ssum @2 (sdot1In (stranspose @[1,6,0,4,5,3,2] (sreplicate @7 (stranspose @[3,2,1,4,5,0] w66))) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) u72))) * stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,0,1,2] (sreshape @[2,7,4,4] (sscatter (stranspose @[2,3,0,1] (sfromR dret) !$ [0, 0]) (\\[i78, i79] -> [i78, i79, kfromS (smaxIndex (sgather (stranspose @[1,2,3,0] u72 * stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)))) (\\[i146] -> [remH (quotH ((112 * i78 + 16 * i79) + i146) 112) 2, remH (quotH ((112 * i78 + 16 * i79) + i146) 16) 7, remH (quotH ((112 * i78 + 16 * i79) + i146) 4) 4, remH ((112 * i78 + 16 * i79) + i146) 4])))]))))))))))) !$ [0]))))"+  printAstPretty (simplifyInlineContractNoExpand $ artDerivativeRev artifactRev)+    @?= "rfromS (let w66 = sgather (stranspose @[4,2,0,3,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [5.0,2.0,-2.0,0.0,13.1,9.0,582934.0,2.99432,6.0,1.0,0.1,-0.2,8.0,-4.0,-335.0,26.0])) (\\[i62, i63] -> [i62 + i63]))) (\\[i64, i65] -> [i64 + i65]) ; u67 = ssum @98 (stranspose @[4,0,1,2,3] (sreshape @[2,7,2,2,98] (str (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w66)))) * sreplicate @2 (str (sreplicate @2 (str (sreplicate @2 (str (sreplicate @1 (sfromR u1)))))))))) ; u72 = sappend (stranspose @[3,1,2,0] (sappend (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i68, i69, i70, i71] -> [ifH (sscalar -0.0 <=. negate (u67 !$ [i69, i70, i68, i71])) 0 1])) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)) in ssum @2 (ssum @2 (sdot1In (stranspose @[4,2,3,0,5,6,7,1] (sreplicate @7 (stranspose @[1,2,3,0] (sreplicate @1 (stranspose @[2,3,0,4,5,1] w66)))) !$ [0]) (stranspose @[4,2,3,1,5,6,7,0] (sreshape @[2,7,2,2,1,2,7,7] (stranspose @[1,2,3,4,0] (sreplicate @98 (stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) u72))) * stranspose @[1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,1,2,0] (sslice (SNat @0) (SNat @2) (stranspose @[3,0,1,2] (sreshape @[2,7,4,4] (sscatter (stranspose @[2,3,0,1] (sfromR u77) !$ [0, 0]) (\\[i78, i79] -> [i78, i79, kfromS (smaxIndex (sreshape @[2,7,16] (stranspose @[1,2,3,0] u72 * stranspose @[1,2,3,0] (sappend (stranspose @[3,1,2,0] (sappend (stranspose @[2,0,1] u67) (sconcrete (sreplicate [2,2,7,2] 0.0)))) (sconcrete (sreplicate [2,2,7,4] 0.0)))) !$ [i78, i79]))]))))))))))) !$ [0]))))"
+ test/simplified/TestGatherSimplified.hs view
@@ -0,0 +1,1614 @@+{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Tests of the gather and scatter operations and of operations that expand+-- to gather and of fusion of all of them.+module TestGatherSimplified (testTrees) where++import Prelude++import Data.Int (Int64)+import GHC.Exts (IsList (..))+import GHC.TypeLits (Div, KnownNat, type (<=))+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape++import HordeAd+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId (resetVarCounter)+import HordeAd.Core.AstInterpret+import HordeAd.Core.CarriersAst+import HordeAd.Core.Ops++import CrossTesting+import EqEpsilon++testTrees :: [TestTree]+testTrees =+  [ testCase "gatherNested1" testGatherNested1+  , testCase "gatherNestedBuild1" testGatherNestedBuild1+  , testCase "gather1" testGather1+  , testCase "gatherBuild1" testGatherBuild1+  , testCase "gatherSimpPP1" testGatherSimpPP1+  , testCase "gatherSimp1" testGatherSimp1+  , testCase "gatherNested02" testGatherNested02+  , testCase "gatherNested2" testGatherNested2+  , testCase "gatherNestedBuild2" testGatherNestedBuild2+  , testCase "gather2" testGather2+  , testCase "gatherBuild2" testGatherBuild2+  , testCase "gatherSimpPP2" testGatherSimpPP2+  , testCase "gatherSimp2" testGatherSimp2+  , testCase "gatherNested12" testGatherNested12+  , testCase "gatherNestedBuild12" testGatherNestedBuild12+  , testCase "gather12" testGather12+  , testCase "gatherBuild12" testGatherBuild12+  , testCase "gatherSimpPP12" testGatherSimpPP12+  , testCase "gatherSimp12" testGatherSimp12+  , testCase "gatherReshape22" testGatherReshape22+  , testCase "gatherReshapeBuild22" testGatherReshapeBuild22+  , testCase "gatherSimpPP22" testGatherSimpPP22+  , testCase "gatherSimp22" testGatherSimp22+  , testCase "gatherSimpPP23" testGatherSimpPP23+  , testCase "gatherSimp23" testGatherSimp23+  , testCase "gatherTranspose33" testGatherTranspose33+  , testCase "gatherTransposeBuild33" testGatherTransposeBuild33+  , testCase "gatherTransposeBuild331" testGatherTransposeBuild331+  , testCase "gatherTransposeBuild332" testGatherTransposeBuild332+  , testCase "gatherTransposeBuild333" testGatherTransposeBuild333+  , testCase "gatherTransposeBuild334" testGatherTransposeBuild334+  , testCase "gatherTransposeBuild335" testGatherTransposeBuild335+  , testCase "gatherTransposeBuild336" testGatherTransposeBuild336+  , testCase "gatherSimpPP33" testGatherSimpPP33+  , testCase "gatherSimpPP34" testGatherSimpPP34+{- TODO: re-enable the tests once we drop GHC 9.10+   (they don't type-check with 9.10)+  , testCase "gatherCond" testGatherCond+  , testCase "gatherCondBuild" testGatherCondBuild+  , testCase "gatherCond2" testGatherCond2+  , testCase "gatherCondBuild2" testGatherCondBuild2+  , testCase "gatherSimpCond" testGatherSimpCond+  , testCase "gatherCond3" testGatherCond3+  , testCase "gatherCondBuild3" testGatherCondBuild3+  , testCase "gatherCond4" testGatherCond4+  , testCase "gatherCondBuild4" testGatherCondBuild4+  , testCase "gatherSimpCond3" testGatherSimpCond3+  , testCase "gatherCond5" testGatherCond5+  , testCase "gatherCondBuild5" testGatherCondBuild5+  , testCase "gatherCond6" testGatherCond6+  , testCase "gatherCondBuild6" testGatherCondBuild6+  , testCase "gatherSimpCond5" testGatherSimpCond5+-}++  , testCase "scatterNested1" testScatterNested1+  , testCase "scatterNestedBuild1" testScatterNestedBuild1+  , testCase "scatter1" testScatter1+  , testCase "scatterBuild1" testScatterBuild1+  , testCase "scatterSimpPP1" testScatterSimpPP1+  , testCase "scatterSimp1" testScatterSimp1+  , testCase "scatterNested2" testScatterNested2+  , testCase "scatterNestedBuild2" testScatterNestedBuild2+  , testCase "scatter2" testScatter2+  , testCase "scatterBuild2" testScatterBuild2+  , testCase "scatterSimpPP2" testScatterSimpPP2+  , testCase "scatterSimp2" testScatterSimp2+  , testCase "scatterNested12" testScatterNested12+  , testCase "scatterNestedBuild12" testScatterNestedBuild12+  , testCase "scatter12" testScatter12+  , testCase "scatterBuild12" testScatterBuild12+  , testCase "scatterSimpPP12" testScatterSimpPP12+  , testCase "scatterSimp12" testScatterSimp12++  , testCase "shmatterBarReluADVal320" testBarReluADVal320+  , testCase "shmatterReluSimpPP" testReluSimpPP++  , testCase "sminimizedCNNOPP2" testCNNOPP2+  , testCase "sminimizedCNNOPP2b" testCNNOPP2b+--  , testCase "sminimizedCNNOPP3" testCNNOPP3+--  , testCase "sminimizedCNNOPP3b" testCNNOPP3b+  , testCase "sminimizedCNNOPP4" testCNNOPP4+  , testCase "sminimizedCNNOPP4b" testCNNOPP4b+  , testCase "sminimizedCNNOPP5" testCNNOPP5+  , testCase "sminimizedCNNOPP5b" testCNNOPP5b+  , testCase "sminimizedCNNOPP6" testCNNOPP6+  , testCase "sminimizedCNNOPP6b" testCNNOPP6b+  , testCase "sminimizedCNNOPP7" testCNNOPP7+  , testCase "sminimizedCNNOPP7b" testCNNOPP7b+  , testCase "minimizedCNNOPP4bU" testCNNOPP4bU+  ]+++-- * Gathers++gatherNested1 :: forall target r. (ADReady target, GoodScalar r)+              => target (TKR 2 r) -> target (TKR 1 r)+gatherNested1 t =+  rgather @1+          (2 :$: ZSR)+          (rgather @1+                   (4 :$: 2 :$: ZSR) t+                   (\(k3 :.: ZIR) -> k3 :.: ZIR))+          (\(i2 :.: ZIR) -> i2 + i2 :.: i2 :.: ZIR)++testGatherNested1 :: Assertion+testGatherNested1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,0.0,0.0,0.0,0.0,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @1 gatherNested1 (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherNestedBuild1 :: Assertion+testGatherNestedBuild1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [3.0,1.0,1.0,1.0,1.0,3.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2+          (\t -> rbuild1 5 (\i ->+             ifH (i >. 2) (gatherNested1 t) (t ! [i])))+          (rreplicate 7 $ ringestData [2] [0, 1]))++gather1 :: forall target r. (ADReady target, GoodScalar r)+        => target (TKR 2 r) -> target (TKR 1 r)+gather1 t =+  rgather @1+          (2 :$: ZSR)+          (rslice 0 4 t)+          (\(i2 :.: ZIR) -> i2 + i2 :.: i2 :.: ZIR)++testGather1 :: Assertion+testGather1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,0.0,0.0,0.0,0.0,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @1 gather1+          (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherBuild1 :: Assertion+testGatherBuild1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [3.0,1.0,1.0,1.0,1.0,3.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2+          (\t -> rbuild1 5 (\i ->+             ifH (i >. 2) (gather1 t) (t ! [i])))+          (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherSimpPP1 :: Assertion+testGatherSimpPP1 = do+  resetVarCounter+  let !t1 = gatherNested1 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 315+  resetVarCounter+  let !t2 = gather1 $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 315+  length (show (simplifyInlineContract @(TKR 1 Float) t1))+    @?= length (show (simplifyInlineContract @(TKR 1 Float) @PrimalSpan t2))++testGatherSimp1 :: Assertion+testGatherSimp1 = do+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]+      env = extendEnv varName (ringestData [7, 2] vals) emptyEnv+  let !t1 = gatherNested1 @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = gather1 (ringestData [7, 2] vals)+  let !t1n = unAstNoSimplify $ gatherNested1 $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ gather1 $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete env t1n+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 1 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 1 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 1 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 1 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++gatherNested02 :: forall target r. (ADReady target, GoodScalar r)+               => target (TKR 1 r) -> target (TKR 1 r)+gatherNested02 t =+  rgather @1+          (1 :$: ZSR)+          (rgather @1+                   (2 :$: ZSR) t+                   (\(k3 :.: ZIR) -> k3 + k3 :.: ZIR))+          (\(i1 :.: ZIR) -> i1 + i1 + i1 :.: ZIR)++testGatherNested02 :: Assertion+testGatherNested02 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [4] [1.0,0.0,0.0,0.0])+    (rev' @Double @1 gatherNested02 (rreplicate 4 (rscalar 0.1)))++gatherNested2 :: forall target r. (ADReady target, GoodScalar r)+              => target (TKR 2 r) -> target (TKR 2 r)+gatherNested2 t =+  rgather @2+          (2 :$: 3 :$: ZSR)+          (rgather @3+                   (2 :$: 3 :$: 4 :$: 2 :$: ZSR) t+                   (\(k1 :.: k2 :.: k3 :.: ZIR) -> k1 + k2 + k3 :.: ZIR))+          (\(i1 :.: i2 :.: ZIR) -> i1 :.: i2 :.: i1 + i2 :.: i1 :.: ZIR)++testGatherNested2 :: Assertion+testGatherNested2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,0.0,0.0,0.0,1.0,1.0,0.0,0.0,1.0,1.0,0.0,0.0,0.0,1.0])+    (rev' @Double @2 gatherNested2+          (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherNestedBuild2 :: Assertion+testGatherNestedBuild2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [6.0,0.0,0.0,0.0,6.0,6.0,0.0,0.0,6.0,6.0,0.0,0.0,0.0,6.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gatherNested2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+          (rreplicate 7 $ ringestData [2] [0, 1]))++gather2 :: forall target r. (ADReady target, GoodScalar r)+        => target (TKR 2 r) -> target (TKR 2 r)+gather2 t =+  rgather @2+          (2 :$: 3 :$: ZSR)+          t+          (\(i1 :.: i2 :.: ZIR) -> i1 + i2 + i1 + i2 :.: i1 :.: ZIR)++testGather2 :: Assertion+testGather2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,0.0,0.0,0.0,1.0,1.0,0.0,0.0,1.0,1.0,0.0,0.0,0.0,1.0])+    (rev' @Double @2 gather2 (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherBuild2 :: Assertion+testGatherBuild2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [6.0,0.0,0.0,0.0,6.0,6.0,0.0,0.0,6.0,6.0,0.0,0.0,0.0,6.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gather2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+          (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherSimpPP2 :: Assertion+testGatherSimpPP2 = do+  resetVarCounter+  let !t1 = gatherNested2 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 582+  resetVarCounter+  let !t2 = gather2 $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 394+  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t1)) @?= 582+  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t2)) @?= 394++testGatherSimp2 :: Assertion+testGatherSimp2 = do+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]+      env = extendEnv varName (ringestData [7, 2] vals) emptyEnv+  let !t1 = gatherNested2 @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = gather2 (ringestData [7, 2] vals)+  let !t1n = unAstNoSimplify $ gatherNested2 $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ gather2 $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete env t1n+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 2 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++gatherNested12 :: forall target r. (ADReady target, GoodScalar r)+               => target (TKR 2 r) -> target (TKR 2 r)+gatherNested12 t =+  rgather @1+          (2 :$: 4 :$: ZSR)+          (rgather @3+                   (2 :$: 3 :$: 4 :$: ZSR) t+                   (\(k1 :.: k2 :.: k3 :.: ZIR) -> k1 + k2 + k3 :.: k1 :.: ZIR))+          (\(i1 :.: ZIR) -> i1 :.: i1 + i1 :.: ZIR)++testGatherNested12 :: Assertion+testGatherNested12 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,0.0,1.0,0.0,1.0,0.0,1.0,1.0,0.0,1.0,0.0,1.0,0.0,1.0])+    (rev' @Double @2 gatherNested12 (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherNestedBuild12 :: Assertion+testGatherNestedBuild12 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [0.0,0.0,4.0,4.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2+          (\t -> rindex (rbuild1 5 (\i ->+             ifH (i >. 2) (gatherNested12 t)+                          (rtranspose [1, 0] $ rreplicate 4 $ t ! [i]))) [1])+          (rreplicate 7 $ ringestData [2] [0, 1]))++gather12 :: forall target r. (ADReady target, GoodScalar r)+         => target (TKR 2 r) -> target (TKR 2 r)+gather12 t =+  rgather @2+          (2 :$: 4 :$: ZSR)+          t+          (\(i1 :.: k3 :.: ZIR) -> i1 + i1 + i1 + k3 :.: i1 :.: ZIR)++testGather12 :: Assertion+testGather12 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,0.0,1.0,0.0,1.0,0.0,1.0,1.0,0.0,1.0,0.0,1.0,0.0,1.0])+    (rev' @Double @2 gather12 (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherBuild12 :: Assertion+testGatherBuild12 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [0.0,0.0,4.0,4.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2+          (\t -> rindex (rbuild1 5 (\i ->+             ifH (i >. 2) (gather12 t)+                          (rtranspose [1, 0] $ rreplicate 4 $ t ! [i]))) [1])+          (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherSimpPP12 :: Assertion+testGatherSimpPP12 = do+  resetVarCounter+  let !t1 = gatherNested12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 515+  resetVarCounter+  let !t2 = gather12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 341+  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 515+  length (show (simplifyInlineContract @(TKR 2 Float) t2)) @?= 341++testGatherSimp12 :: Assertion+testGatherSimp12 = do+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]+      env = extendEnv varName (ringestData [7, 2] vals) emptyEnv+  let !t1 = gatherNested12 @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = gather12 (ringestData [7, 2] vals)+  let !t1n = unAstNoSimplify $ gatherNested12 $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ gather12 $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete env t1n+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 2 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++gatherReshape22 :: forall target r. (ADReady target, GoodScalar r)+                => target (TKR 2 r) -> target (TKR 2 r)+gatherReshape22 t =+  rreshape @6 [2, 6]+  $ rreshape [3, 1, 2, 1, 1, 2]+  $ rreshape @4 (1 :$: 12 :$: 1 :$: ZSR)+  $ rreshape @3 [3, 1, 1, 4]+  $ rreshape [2, 2, 3] t++testGatherReshape22 :: Assertion+testGatherReshape22 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [6,2]+                 [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @2 gatherReshape22+                               (rreplicate 6 $ ringestData [2] [0, 1]))++testGatherReshapeBuild22 :: Assertion+testGatherReshapeBuild22 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [6,2]+                 [6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gatherReshape22 (t * rreplicate0N [6, 2] (rfromIndex0 i))))+          (rreplicate 6 $ ringestData [2] [0, 1]))++testGatherSimpPP22 :: Assertion+testGatherSimpPP22 = do+  resetVarCounter+  let !t1 = gatherReshape22 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 159+  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 159+  resetVarCounter+  let !t2 = rreshape @2 @2 [2, 6]+            $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 159+  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t2)) @?= 159++testGatherSimp22 :: Assertion+testGatherSimp22 = do+  let varName = mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0]+      env = extendEnv varName (ringestData [6, 2] vals) emptyEnv+  let !t1 = gatherReshape22 @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = rreshape @2 @2 [2, 6] (ringestData [6, 2] vals)+  let !t1n = unAstNoSimplify $ gatherReshape22 $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ rreshape @2 @2 [2, 6] $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete env t1n+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 2 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++testGatherSimpPP23 :: Assertion+testGatherSimpPP23 = do+  resetVarCounter+  let !t1 = (\t -> rbuild1 4 (\i ->+              gatherReshape22 @(AstTensor AstMethodLet PrimalSpan)+                (t * rreplicate0N [6, 2] (rfromIndex0 i))))+            $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 450+  length (show (simplifyInlineContract @(TKR 3 Float) t1)) @?= 456+  resetVarCounter+  let !t2 = (\t -> rbuild1 4 (\i ->+              rreshape @2 @2 [2, 6]+                (t * rreplicate0N [6, 2] (rfromIndex0 i))))+            $ AstVar (mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 450+  length (show (simplifyInlineContract @(TKR 3 Float) @PrimalSpan t2)) @?= 456++testGatherSimp23 :: Assertion+testGatherSimp23 = do+  let varName = mkAstVarName (FTKR [6, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0]+      env = extendEnv varName (ringestData [6, 2] vals) emptyEnv+  let !t1 = (\t -> rbuild1 4 (\i ->+              gatherReshape22 @(AstTensor AstMethodLet PrimalSpan)+                (t * rreplicate0N [6, 2] (rfromIndex0 i)))) var+  let !t2 = (\t -> rbuild1 4 (\i ->+              rreshape @2 @2 [2, 6]+                (t * rreplicate0N [6, 2] (rfromIndex0 i)))) (ringestData [6, 2] vals)+  let !t1n = unAstNoSimplify $ (\t -> rbuild1 4 (\i ->+              gatherReshape22+                (t * rreplicate0N [6, 2] (rfromIndex0 i)))) $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ (\t -> rbuild1 4 (\i ->+              rreshape @2 @2 [2, 6]+                (t * rreplicate0N [6, 2] (rfromIndex0 i)))) $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete env t1n+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 3 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 3 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 3 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 3 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++-- Depending on if and how transpose it desugared, this may or may not result+-- in dozens of nested gathers that should vanish after simplification.+gatherTranspose33 :: forall target r. (ADReady target, GoodScalar r, RealFloat r)+                  => target (TKR 10 r) -> target (TKR 2 r)+gatherTranspose33 t =+  rmatmul2 (rreshape [6, 8] (rconcrete $ unConcrete t48))+    (rtr+     $ rreshape @4 [16, 8]+     $ rtranspose [0, 1, 2]+     $ rtranspose [2, 0, 1]+     $ rtranspose [1, 2, 0]+     $ rtranspose [1, 0, 2]+     $ rtranspose [1, 0]+     $ rtranspose [0, 1, 2, 3]+     $ rtranspose [1, 2, 3, 0]+     $ rtranspose [3, 0, 2, 1]+     $ rreshape [2, 2, 8, 4]+     $ rtranspose [0, 1, 2, 3]+     $ rtranspose [1, 2, 3, 0]+     $ rtranspose [1, 0, 3, 2]+     $ rtranspose [0, 1, 2, 3, 4, 5, 6, 7, 9, 8]+     $ rtranspose [0, 1, 2, 3, 7, 5, 6, 4]+     $ rtranspose [0, 1, 2, 3, 4, 5, 6]+     $ rtranspose [5, 0, 1, 2, 3, 4]+     $ rtranspose [0, 1, 2, 4, 3, 5, 6, 7, 9, 8]+     $ rtranspose []+     $ rtranspose [0]+     $ rtranspose [0, 1]+     $ rtranspose [1, 0]+     $ rtranspose [0, 1, 7, 4, 5, 3, 6, 2, 8]+       t)++testGatherTranspose33 :: Assertion+testGatherTranspose33 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,2,2,1,2,2,2,2,2,1] [81.3003,71.0,81.3003,71.0,81.3003,71.0,81.3003,71.0,80.0,79.0,80.0,79.0,80.0,79.0,80.0,79.0,81.3003,71.0,81.3003,71.0,81.3003,71.0,81.3003,71.0,80.0,79.0,80.0,79.0,80.0,79.0,80.0,79.0,81.3003,71.0,81.3003,71.0,81.3003,71.0,81.3003,71.0,80.0,79.0,80.0,79.0,80.0,79.0,80.0,79.0,81.3003,71.0,81.3003,71.0,81.3003,71.0,81.3003,71.0,80.0,79.0,80.0,79.0,80.0,79.0,80.0,79.0,166.8003,137.70326,166.8003,137.70326,166.8003,137.70326,166.8003,137.70326,186.1003,162.3889400002,186.1003,162.3889400002,186.1003,162.3889400002,186.1003,162.3889400002,166.8003,137.70326,166.8003,137.70326,166.8003,137.70326,166.8003,137.70326,186.1003,162.3889400002,186.1003,162.3889400002,186.1003,162.3889400002,186.1003,162.3889400002,166.8003,137.70326,166.8003,137.70326,166.8003,137.70326,166.8003,137.70326,186.1003,162.3889400002,186.1003,162.3889400002,186.1003,162.3889400002,186.1003,162.3889400002,166.8003,137.70326,166.8003,137.70326,166.8003,137.70326,166.8003,137.70326,186.1003,162.3889400002,186.1003,162.3889400002,186.1003,162.3889400002,186.1003,162.3889400002])+    (rev' @Double @2 gatherTranspose33 t128)++testGatherTransposeBuild33 :: Assertion+testGatherTransposeBuild33 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,2,2,1,2,2,2,2,2,1] [487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,480.0,474.0,480.0,474.0,480.0,474.0,480.0,474.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,480.0,474.0,480.0,474.0,480.0,474.0,480.0,474.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,480.0,474.0,480.0,474.0,480.0,474.0,480.0,474.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,487.80179999999996,426.0,480.0,474.0,480.0,474.0,480.0,474.0,480.0,474.0,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1000.8018,826.21956,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012,1116.6018,974.3336400012])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gatherTranspose33 (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (rfromIndex0 i))))+          t128)++testGatherTransposeBuild331 :: Assertion+testGatherTransposeBuild331 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 3] [1,1,1,1,1,1])+    (rev' @Double @3+          (\t -> rbuild1 2 (\i ->+             rtranspose [1, 0] (t * rreplicate0N [2, 3] (rfromIndex0 i))))+          (ringestData [2, 3] [1,2,3,4,5,6]))++testGatherTransposeBuild332 :: Assertion+testGatherTransposeBuild332 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 3] [1,1,1,1,1,1])+    (rev' @Double @3+          (\t -> rbuild1 2 (\i ->+             rtranspose [1, 0] (t * rreplicate0N [2, 3] (rfromIndex0 i))))+          (ringestData [2, 3] [1,2,3,4,5,6]))++testGatherTransposeBuild333 :: Assertion+testGatherTransposeBuild333 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2] [1,1])+    (rev' @Double @2+          (\t -> rbuild1 2 (\i ->+             t * rreplicate0N [2] (rfromIndex0 i)))+          (ringestData [2] [0,0]))++testGatherTransposeBuild334 :: Assertion+testGatherTransposeBuild334 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 1] [1,1])+    (rev' @Double @3+          (\t -> rbuild1 2 (\i ->+             t * rreplicate 2 (rreplicate 1 (rfromIndex0 i))))+         (ringestData [2, 1] [1,2]))++testGatherTransposeBuild335 :: Assertion+testGatherTransposeBuild335 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 1] [1,1])+    (rev' @Double @3+          (\t ->+             rreplicate 2 t * rtranspose [2,0,1] (rreplicate 2 (rreplicate 1 (rfromIntegral @Int64 (rconcrete $ Nested.rfromListPrimLinear (fromList [2]) [0, 1])))))+         (ringestData [2, 1] [1,2]))++testGatherTransposeBuild336 :: Assertion+testGatherTransposeBuild336 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2, 1] [1,1])+    (rev' @Double @3+          (\t ->+             rreplicate 2 t * rtranspose [2,0,1] (rreplicate 2 (rreplicate 1 (rfromList [rscalar 0, rscalar 1]))))+         (ringestData [2, 1] [1,2]))++-- These are different terms, but they should have similar lengths,+-- because they differ only by single transpose and reshape, most probably,+-- and all the rest of the element reordering should cancel out.+-- Still, probably impossible to lower the gap to zero.+testGatherSimpPP33 :: Assertion+testGatherSimpPP33 = do+  resetVarCounter+  let !t1 = gatherTranspose33 @(AstTensor AstMethodLet PrimalSpan)+            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 1117+  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 847+  resetVarCounter+  let !t2 = (\t -> rmatmul2 (rreshape [6, 8] (rconcrete $ unConcrete t48))+                            (rreshape @10 [8, 16] t))+            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 796+  length (show (simplifyInlineContract @(TKR 2 Float) @PrimalSpan t2)) @?= 526++testGatherSimpPP34 :: Assertion+testGatherSimpPP34 = do+  resetVarCounter+  let !t1 = (\t -> rbuild1 4 (\i ->+             gatherTranspose33 @(AstTensor AstMethodLet PrimalSpan) (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (rfromIndex0 i))))+            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 2504+  length (show (simplifyInlineContract @(TKR 3 Float) t1)) @?= 19838+  resetVarCounter+  let !t2 = (\t -> rbuild1 4 (\i ->+              (\t' -> rmatmul2 (rreshape [6, 8] (rconcrete $ unConcrete t48))+                               (rreshape @10 [8, 16] t'))+                (t * rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (rfromIndex0 i))))+            $ AstVar (mkAstVarName (FTKR [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 2145+  length (show (simplifyInlineContract @(TKR 3 Float) @PrimalSpan t2)) @?= 19479++{- TODO: re-enable the tests once we drop GHC 9.10+   (they don't type-check with 9.10)+gatherCond :: forall target r. (ADReady target, GoodScalar r)+           => target (TKR 2 r) -> target (TKR 2 r)+gatherCond u =+  let v = rtranspose [2, 0, 1] $ rreplicate (2 * rwidth u) u+  in rgather [rwidth u, 2] v (\(i :.: j :.: ZIR) ->+                                ifH (i ==. 3) 0 j :.: 2 * i :.: i :.: ZIR)++testGatherCond :: Assertion+testGatherCond =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,1.0,1.0,1.0,1.0,1.0,2.0,0.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @2 gatherCond (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherCondBuild :: Assertion+testGatherCondBuild =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [6.0,6.0,6.0,6.0,6.0,6.0,12.0,0.0,6.0,6.0,6.0,6.0,6.0,6.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gatherCond (t * rreplicate0N [7, 2] (rfromIndex0 i))))+          (rreplicate 7 $ ringestData [2] [0, 1]))++gatherCond2 :: forall target r. (ADReady target, GoodScalar r)+            => target (TKR 2 r) -> target (TKR 2 r)+gatherCond2 u =+  let v = rreplicate (2 * rwidth u) u+  in rtr $ rgather [2, rwidth u] v (\(j :.: i :.: ZIR) ->+                                      2 * i :.: i :.: ifH (i ==. 3) 0 j :.: ZIR)++testGatherCond2 :: Assertion+testGatherCond2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,1.0,1.0,1.0,1.0,1.0,2.0,0.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @2 gatherCond2 (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherCondBuild2 :: Assertion+testGatherCondBuild2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [6.0,6.0,6.0,6.0,6.0,6.0,12.0,0.0,6.0,6.0,6.0,6.0,6.0,6.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gatherCond2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+          (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherSimpCond :: Assertion+testGatherSimpCond = do+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]+      env = extendEnv varName (ringestData [7, 2] vals) emptyEnv+  let !t1 = gatherCond @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = gatherCond2 (ringestData [7, 2] vals)+  let !t1n = unAstNoSimplify $ gatherCond $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ gatherCond2 $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete env t1n+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 2 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++gatherCond3 :: forall target r. (ADReady target, GoodScalar r)+            => target (TKR 2 r) -> target (TKR 2 r)+gatherCond3 u =+  let v = rtranspose [2, 0, 1] $ rreplicate (2 * rwidth u) u+  in rgather [rwidth u, 2] v (\(i :.: j :.: ZIR) ->+                                2 * i :.: i :.: ifH (i ==. 3) 0 j :.: ZIR)++testGatherCond3 :: Assertion+testGatherCond3 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,0.0,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2 gatherCond3 (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherCondBuild3 :: Assertion+testGatherCondBuild3 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [6.0,0.0,6.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gatherCond3 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+          (rreplicate 7 $ ringestData [2] [0, 1]))++gatherCond4 :: forall target r. (ADReady target, GoodScalar r)+            => target (TKR 2 r) -> target (TKR 2 r)+gatherCond4 u =+  let v = rreplicate (2 * rwidth u) u+  in rtr $ rgather [2, rwidth u] v (\(j :.: i :.: ZIR) ->+                                      i :.: ifH (i ==. 3) 0 j :.: 2 * i :.: ZIR)++testGatherCond4 :: Assertion+testGatherCond4 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,0.0,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2 gatherCond4 (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherCondBuild4 :: Assertion+testGatherCondBuild4 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [6.0,0.0,6.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gatherCond4 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+          (rreplicate 7 $ ringestData [2] [0, 1]))++testGatherSimpCond3 :: Assertion+testGatherSimpCond3 = do+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]+      env = extendEnv varName (ringestData [7, 2] vals) emptyEnv+  let !t1 = gatherCond3 @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = gatherCond4 (ringestData [7, 2] vals)+  let !t1n = unAstNoSimplify $ gatherCond3 $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ gatherCond4 $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete env t1n+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 2 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++gatherCond5 :: forall target r. (ADReady target, GoodScalar r)+            => target (TKR 3 r) -> target (TKR 2 r)+gatherCond5 v =+  rgather [rwidth v, 2] v (\(i :.: j :.: ZIR) ->+                             ifH (i ==. 1) 0 j :.: 2 * i :.: i :.: ZIR)++testGatherCond5 :: Assertion+testGatherCond5 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,4,2]+                 [1.0,0.0,0.0,0.0,0.0,2.0,0.0,0.0,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2 gatherCond5 (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))++testGatherCondBuild5 :: Assertion+testGatherCondBuild5 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,4,2]+                 [6.0,0.0,0.0,0.0,0.0,12.0,0.0,0.0,6.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gatherCond5 (t * rreplicate0N [2,4,2] (rfromIndex0 i))))+          (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))++gatherCond6 :: forall target r. (ADReady target, GoodScalar r)+            => target (TKR 3 r) -> target (TKR 2 r)+gatherCond6 u =+  let v = rtranspose [2, 0, 1] u+  in rtr $ rgather [2, rwidth v] v (\(j :.: i :.: ZIR) ->+                                      i :.: ifH (i ==. 1) 0 j :.: 2 * i :.: ZIR)++testGatherCond6 :: Assertion+testGatherCond6 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,4,2]+                 [1.0,0.0,0.0,0.0,0.0,2.0,0.0,0.0,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2 gatherCond6 (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))++testGatherCondBuild6 :: Assertion+testGatherCondBuild6 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,4,2]+                 [6.0,0.0,0.0,0.0,0.0,12.0,0.0,0.0,6.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             gatherCond6 (t * rreplicate0N [2,4,2] (rfromIndex0 i))))+          (rreplicate 2 $ rreplicate 4 $ ringestData [2] [0, 1]))++testGatherSimpCond5 :: Assertion+testGatherSimpCond5 = do+  let varName = mkAstVarName (FTKR [2,4,2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1,0,2.0,2.13,0.2,11.0,-17.0,23.0,29.0,-35.0,41.0,1.4,-0.33,33.0,0.1,0.007]+      env = extendEnv varName (ringestData [2,4,2] vals) emptyEnv+  let !t1 = gatherCond5 @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = gatherCond6 (ringestData [2,4,2] vals)+  let !t1n = unAstNoSimplify $ gatherCond5 $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ gatherCond6 $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete env t1n+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 2 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n+-}+++-- * Scatters instead of gathers++scatterNested1 :: forall target r. (ADReady target, GoodScalar r)+               => target (TKR 2 r) -> target (TKR 1 r)+scatterNested1 t =+  rscatter @2+          (2 :$: ZSR)+          (rscatter @1+                   (7 :$: 2 :$: ZSR) t+                   (\(k3 :.: ZIR) -> k3 :.: ZIR))+          (\(i1 :.: i2 :.: ZIR) -> i2 `quotH` (1 + i1) :.: ZIR)++testScatterNested1 :: Assertion+testScatterNested1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @1 scatterNested1 (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterNestedBuild1 :: Assertion+testScatterNestedBuild1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [3.0,3.0,3.0,3.0,3.0,3.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0])+    (rev' @Double @2+          (\t -> rbuild1 5 (\i ->+             ifH (i >. 2) (scatterNested1 t) (t ! [i])))+          (rreplicate 7 $ ringestData [2] [0, 1]))++scatter1 :: forall target r. (ADReady target, GoodScalar r)+         => target (TKR 2 r) -> target (TKR 1 r)+scatter1 t =+  rscatter @2+          (2 :$: ZSR)+          t+          (\(i1 :.: i2 :.: ZIR) -> minH (i2 + 2 * i1) 1 :.: ZIR)++testScatter1 :: Assertion+testScatter1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @1 scatter1 (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterBuild1 :: Assertion+testScatterBuild1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [3.0,3.0,3.0,3.0,3.0,3.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0])+    (rev' @Double @2+          (\t -> rbuild1 5 (\i ->+             ifH (i >. 2) (scatter1 t) (t ! [i])))+          (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterSimpPP1 :: Assertion+testScatterSimpPP1 = do+  resetVarCounter+  let !t1 = scatterNested1 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 397+  resetVarCounter+  let !t2 = scatter1 $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 492+  length (show (simplifyInlineContract @(TKR 1 Float) @PrimalSpan t1)) @?= 397+  length (show (simplifyInlineContract @(TKR 1 Float) @PrimalSpan t2)) @?= 492++testScatterSimp1 :: Assertion+testScatterSimp1 = do+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]+      env = extendEnv varName (ringestData [7, 2] vals) emptyEnv+  let !t1 = scatterNested1 @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = scatter1 (ringestData [7, 2] vals)+  let !t1n = unAstNoSimplify $ scatterNested1 $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ scatter1 $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  -- TODO: scatter fusion isn't sound? or just incorrectly manually done here?+  -- interpretAstPrimal @Concrete env t1n+  --   @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 1 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 1 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 1 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 1 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++scatterNested2 :: forall target r. (ADReady target, GoodScalar r)+              => target (TKR 2 r) -> target (TKR 2 r)+scatterNested2 t =+  rscatter @4+          (2 :$: 3 :$: ZSR)+          (rscatter @1+                   (2 :$: 3 :$: 4 :$: 2 :$: ZSR) t+                   (\(k1 :.: ZIR) -> minH k1 1 :.: minH k1 2  :.: minH k1 3 :.: ZIR))+          (\(i1 :.: i2 :.: _i3 :.: i4 :.: ZIR) ->+            minH (i1 + i2) 1 :.: minH (i4 + i1) 2 :.: ZIR)++testScatterNested2 :: Assertion+testScatterNested2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @2 scatterNested2 (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterNestedBuild2 :: Assertion+testScatterNestedBuild2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             scatterNested2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+          (rreplicate 7 $ ringestData [2] [0, 1]))++scatter2 :: forall target r. (ADReady target, GoodScalar r)+        => target (TKR 2 r) -> target (TKR 2 r)+scatter2 t =+  rscatter @2+          (2 :$: 3 :$: ZSR)+          t+          (\(i1 :.: i2 :.: ZIR) -> minH (i1 + i2 + i1 + i2) 1 :.: minH i1 2 :.: ZIR)++testScatter2 :: Assertion+testScatter2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @2 scatter2 (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterBuild2 :: Assertion+testScatterBuild2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0])+    (rev' @Double @3+          (\t -> rbuild1 4 (\i ->+             scatter2 (t * rreplicate0N [7, 2] (rfromIndex0 i))))+          (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterSimpPP2 :: Assertion+testScatterSimpPP2 = do+  resetVarCounter+  let !t1 = scatterNested2 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 1022+  resetVarCounter+  let !t2 = scatter2 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 782+  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 1022+  length (show (simplifyInlineContract @(TKR 2 Float) t2)) @?= 782++testScatterSimp2 :: Assertion+testScatterSimp2 = do+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]+      env = extendEnv varName (ringestData [7, 2] vals) emptyEnv+  let !t1 = scatterNested2 @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = scatter2 (ringestData [7, 2] vals)+  let !t1n = unAstNoSimplify $ scatterNested2 $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ scatter2 $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  -- TODO: scatter fusion isn't sound? or just incorrectly manually done here?+  -- interpretAstPrimal @Concrete env t1n+  --  @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 2 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++scatterNested12 :: forall target r. (ADReady target, GoodScalar r)+               => target (TKR 2 r) -> target (TKR 2 r)+scatterNested12 t =+  rscatter @2+          (2 :$: 4 :$: ZSR)+          (rscatter @2+                   (2 :$: 3 :$: 4 :$: ZSR) t+                   (\(k1 :.: k2 :.: ZIR) ->+                     minH k1 1 :.: minH (k2 + k1) 2 :.: minH k1 3 :.: ZIR))+          (\(i1 :.: _i2 :.: ZIR) -> minH (i1 + i1) 1 :.: ZIR)++testScatterNested12 :: Assertion+testScatterNested12 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @2 scatterNested12 (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterNestedBuild12 :: Assertion+testScatterNestedBuild12 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [0.0,0.0,4.0,4.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2+          (\t -> rindex (rbuild1 5 (\i ->+             ifH (i >. 2) (scatterNested12 t)+                          (rtranspose [1, 0] $ rreplicate 4 $ t ! [i]))) [1])+          (rreplicate 7 $ ringestData [2] [0, 1]))++scatter12 :: forall target r. (ADReady target, GoodScalar r)+         => target (TKR 2 r) -> target (TKR 2 r)+scatter12 t =+  rscatter @2+          (2 :$: 4 :$: ZSR)+          t+          (\(i1 :.: k3 :.: ZIR) -> minH (i1 + i1 + i1 + k3) 1 :.: minH i1 3 :.: ZIR)++testScatter12 :: Assertion+testScatter12 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])+    (rev' @Double @2 scatter12 (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterBuild12 :: Assertion+testScatterBuild12 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7,2]+                 [0.0,0.0,4.0,4.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @2+          (\t -> rindex (rbuild1 5 (\i ->+             ifH (i >. 2) (scatter12 t)+                          (rtranspose [1, 0] $ rreplicate 4 $ t ! [i]))) [1])+          (rreplicate 7 $ ringestData [2] [0, 1]))++testScatterSimpPP12 :: Assertion+testScatterSimpPP12 = do+  resetVarCounter+  let !t1 = scatterNested12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 952+  resetVarCounter+  let !t2 = scatter12 @(AstTensor AstMethodLet PrimalSpan) $ AstVar (mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 625+  length (show (simplifyInlineContract @(TKR 2 Float) t1)) @?= 952+  length (show (simplifyInlineContract @(TKR 2 Float) t2)) @?= 625++testScatterSimp12 :: Assertion+testScatterSimp12 = do+  let varName = mkAstVarName (FTKR [7, 2] FTKScalar) Nothing . intToAstVarId $ 100000000+      var = AstVar varName+      vals = [-1, 0, 2.0,5.0,11.0,-17.0,23.0,29.0,-35.0,41.0,47.0,33.0, 0.1, 0.007]+      env = extendEnv varName (ringestData [7, 2] vals) emptyEnv+  let !t1 = scatterNested12 @(AstTensor AstMethodLet PrimalSpan) var+  let !t2 = scatter12 (ringestData [7, 2] vals)+  let !t1n = unAstNoSimplify $ scatterNested12 $ AstNoSimplify var+  let !t2n = unAstNoSimplify $ scatter12 $ AstNoSimplify var+  interpretAstPrimal @Concrete env t1+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete env t1n+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete emptyEnv t2+    @?= interpretAstPrimal @Concrete env t2n+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1)+    @?= interpretAstPrimal @Concrete env t1+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t1n)+    @?= interpretAstPrimal @Concrete env t1n+  interpretAstPrimal @Concrete emptyEnv+                     (simplifyInlineContract @(TKR 2 Float) t2)+    @?= interpretAstPrimal @Concrete emptyEnv t2+  interpretAstPrimal @Concrete env+                     (simplifyInlineContract @(TKR 2 Float) t2n)+    @?= interpretAstPrimal @Concrete env t2n++foo :: RealFloatH a => (a,a,a) -> a+foo (x,y,z) =+  let w = x * sin y+  in atan2H z w + z * w++bar :: forall a. RealFloatH a => (a, a) -> a+bar (x, y) =+  let w = foo (x, y, x) * sin y+  in atan2H x w + y * w++barRelu+  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )+  => target (TKR n r) -> target (TKR n r)+barRelu x = let t = rreplicate0N (rshape x) (rscalar 0.001) * x+            in relu $ bar (t, relu t)++barRelu10xSlower+  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )+  => target (TKR n r) -> target (TKR n r)+barRelu10xSlower x = let t = rmap0N (* rscalar 0.001) x+                     in relu $ bar (t, relu t)++testBarReluADVal320 :: Assertion+testBarReluADVal320 =+  assertEqualUpToEpsilon 1e-10+    (ringestData [1,2,2,1,2,2,2,2,2,1] [2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885034988157713e-4,2.885923176600045e-4,2.887454843457817e-4,2.886097295122454e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.8851415976532735e-4,2.885923176600045e-4,2.887454843457817e-4,2.8849246223035154e-4,2.884182085399516e-4,2.884075468755327e-4,2.8842176240868867e-4,2.8840399312321096e-4,0.0,2.887454843457817e-4,2.886097295122454e-4,2.887454843457817e-4,2.88599069218435e-4,2.887454843457817e-4,2.886097295122454e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.885145151321922e-4,2.8854294397024206e-4,2.8858878438222746e-4,2.885923176600045e-4,0.0,2.884007943794131e-4,0.0,2.884469945274759e-4,2.8843242392031246e-4,2.884288700806792e-4,0.0,2.885034988157713e-4,2.884110805753153e-4,0.0,2.8849283778617973e-4,2.884075468755327e-4,2.884075468755327e-4,2.884075468755327e-4,2.884075468755327e-4,0.0,0.0,0.0,0.0,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,0.0,0.0,0.0,0.0,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,2.8854294397024206e-4,2.884288700806792e-4,2.884395315486472e-4,0.0,2.8849246223035154e-4,2.8850276789489724e-4,0.0,2.8849212704517413e-4,2.8854294397024206e-4,2.884288700806792e-4,2.884395315486472e-4,0.0,2.8849246223035154e-4,2.8850276789489724e-4,0.0,2.8849212704517413e-4,2.8842922547482884e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.894378297730782e-4,2.885923176600045e-4,2.887454843457817e-4,2.88599069218435e-4,2.887454843457817e-4,2.887056688523444e-4,2.887454843457817e-4,2.887056688523444e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.884786229769816e-4,2.885923176600045e-4,2.887454843457817e-4,2.886950092188272e-4,2.887454843457817e-4,2.884818011261814e-4,2.887454843457817e-4,2.886097295122454e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885145151321922e-4,2.8854294397024206e-4,2.887167039107226e-4,2.885923176600045e-4,2.887454843457817e-4,2.8860262265516213e-4,2.887454843457817e-4,2.885884088500461e-4,2.887454843457817e-4,2.88599069218435e-4])+    (grad (kfromR . rsum0 @10 @(TKScalar Double) . barRelu10xSlower)+          (rmap0N (* rscalar 0.001) t128))++testReluSimpPP :: Assertion+testReluSimpPP = do+  resetVarCounter+  let !t1 = barRelu10xSlower @(AstTensor AstMethodLet PrimalSpan)+            $ AstVar (mkAstVarName (FTKR [1,2,2,1,2,2,2,2,2,1] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t1) @?= 22675+  length (show (simplifyInlineContract @(TKR 10 Float) t1)) @?= 22675+  resetVarCounter+  let !t2 = barRelu @(AstTensor AstMethodLet PrimalSpan)+            $ AstVar (mkAstVarName (FTKR [1,2,2,1,2,2,2,2,2,1] FTKScalar) Nothing . intToAstVarId $ 100000000)+  length (show t2) @?= 12258+  length (show (simplifyInlineContract @(TKR 10 Float) t2)) @?= 12258++testCNNOPP2 :: Assertion+testCNNOPP2 = do+  resetVarCounter+  let t = maxPool2dUnpadded2+            (rconcrete $ Nested.rreplicateScal (1 :$: 1 :$: 2 :$: 2 :$: ZSR) 1)+  printAstPretty (simplifyInlineContract t)+    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printAstPretty t+    @?= "rfromS (sreplicate @2 (sreplicate @2 (let u37 = let u42 = sgather (sgather (sreplicate @1 (let w32 = sgather (stranspose @[3,2,0,1] (sgather (sconcrete (sfromListLinear [2,3,2] [1.0,1.0,0.0,0.0,0.0,0.0,1.0,1.0,0.0,0.0,0.0,0.0])) (\\[i26, i15] -> [i26 + i15]))) (\\[i22, i16] -> [i22 + i16]) in stranspose @[1,2,3,0] (sappend (sconcrete (sreplicate [1,2,2,2,2] 0.0)) (sappend (sreplicate @1 (stranspose @[2,0,4,1,3] w32 !$ [1])) (sconcrete (sreplicate [1,2,2,2,2] 0.0)))))) (\\[i20] -> [i20, i20, i20, 0])) (\\[i45, i40, i36, i8] -> [2 * i40 + i8, i40, 2 * i45 + i36]) in str (sappend (sconcrete (sreplicate [1,2,2,2] 0.0)) (sreplicate @1 (str u42 !$ [1]))) in stranspose @[2,3,0,1] u37 !$ [0, 0])))"++testCNNOPP2b :: Assertion+testCNNOPP2b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded2 (FTKR [1, 1, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (sconcrete (sreplicate [1,1,2,2] 0.0))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (sconcrete (sreplicate [1,1,2,2] 0.0))"++maxPool2dUnpadded2+  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded2 a =+  rbuild [2, 2, 2, 2] $ \case+    [_, _, iBh, iBw] ->+      let arrt = slicez2 (conv2dUnpadded2 a) [iBw, 1, 2 * iBh, 2 * iBw]+      in rmaximum2 arrt+    _ -> error "maxPool2dUnpadded2: impossible pattern needlessly required"++conv2dUnpadded2+  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded2 a =+  rbuild [3, 3, 2, 2] $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez2 a [iImg, 0, iBh, iBw]+      in rindex0 arrAt [0, iBw, iBw, 0]+    _ -> error "conv2dUnpadded2: impossible pattern needlessly required"++slicez2+  :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double, n ~ 4)+  => target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez2 d ixBase =+  rbuild [1, 1, 2, 2] $ \ixResult -> indexz02 d (ixrZipWith (+) ixBase ixResult)++indexz02+  :: forall target r n.+     (target ~ AstTensor AstMethodLet FullSpan, r ~ Double, n ~ 4)+  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)+indexz02 d ix = ifH (1 ==. (toList ix !! 0)) (d ! ix) (rscalar 0)++rmaximum2 :: (target ~ AstTensor AstMethodLet FullSpan, r ~ Double)+         => target (TKR 4 r) -> target (TKR 0 r)+rmaximum2 t0 = tlet t0 $ \t -> rindex0 t [0, 0, 0, 0]++{- TODO: divergent result; bring back when GHC 9.10 dropped:+testCNNOPP3 :: Assertion+testCNNOPP3 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @3) knownSTK+                   $ AstReplicate (SNat @3) knownSTK+                   $ AstReplicate (SNat @3) knownSTK+                   $ AstReplicate (SNat @3) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded3 $ conv2dUnpadded3 blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (sconcrete (sfromListLinear [2,2,2,2] [14.0,0.0,14.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printAstPretty afcnn2T+    @?= "rfromS (stranspose @[1,2,0] (sreplicate @2 (let w32 = stranspose @[1,2,3,0] (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @2 (stranspose @[1,2,3,0] (sreplicate @2 (let t44 = sreplicate @2 (sgather (sappend (stranspose @[4,1,2,3,5,0] (sappend (sreplicate @1 (stranspose @[0,1,2,4,3] (sgather (str (sreplicate @1 (sconcrete (sfromListLinear [2] [7.0,0.0])))) (\\[i31, i25, i20, i17] -> [ifH (notB (notB (0 <=. negate i31 + negate i20) &&* notB (1 <=. i31 + i20 &&* (-1) <=. negate i31 + negate i20)) &&* notB (notB (0 <=. negate i25 + negate i17) &&* notB (1 <=. i25 + i17 &&* (-1) <=. negate i25 + negate i17))) 0 1])))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @1 (sreplicate @2 (sscalar 0.0))))))))) (stranspose @[4,1,2,3,5,0] (sappend (sreplicate @1 (stranspose @[0,1,2,4,3] (sgather (str (sreplicate @1 (sconcrete (sfromListLinear [2] [7.0,0.0])))) (\\[i30, i24, i19, i17] -> [ifH (notB (notB (0 <=. negate i30 + negate i19) &&* notB (1 <=. i30 + i19 &&* (-1) <=. negate i30 + negate i19)) &&* notB (notB (0 <=. negate i24 + negate i17) &&* notB (1 <=. i24 + i17 &&* (-1) <=. negate i24 + negate i17))) 0 1])))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sreplicate @2 (sreplicate @1 (sreplicate @2 (sscalar 0.0)))))))))) (\\[i39, i4] -> [2 * i4, 0, 2 * i4, 0, 0, 2 * i39]) + sgather (sappend (sreplicate @1 (sgather (stranspose @[1,2,3,4,5,0] (sappend (sreplicate @1 (stranspose @[0,2,4,3,1] (sgather (str (sreplicate @1 (sconcrete (sfromListLinear [2] [7.0,0.0])))) (\\[i28, i22, i19, i17] -> [ifH (notB (notB (0 <=. negate i28 + negate i19) &&* notB (1 <=. i28 + i19 &&* (-1) <=. negate i28 + negate i19)) &&* notB (notB (0 <=. negate i22 + negate i17) &&* notB (1 <=. i22 + i17 &&* (-1) <=. negate i22 + negate i17))) 0 1])))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sreplicate @1 (sreplicate @2 (sreplicate @2 (sscalar 0.0))))))))) (\\[i27] -> [i27, i27, 0, 1, 0]))) (str (sreplicate @2 (str (sreplicate @2 (sconcrete (sfromListLinear [1] [0.0]))))))) (\\[i38, i4] -> [2 * i4, 0, 2 * i38])) in sappend (str (sappend (stranspose @[1,2,0] (sappend (sgather t44 (\\[i34, i37, i42] -> [i42, i37, i34])) (sreplicate @1 (sreplicate @1 (sreplicate @1 (sscalar 0.0)))))) (sreplicate @1 (sreplicate @1 (sreplicate @2 (sscalar 0.0)))))) (sreplicate @1 (sreplicate @2 (sreplicate @2 (sscalar 0.0)))))))))))) in stranspose @[3,4,5,6,0,1,2] w32 !$ [0, 0, 0, 0])))"+-}++-- TODO: OOMs+_testCNNOPP3b :: Assertion+_testCNNOPP3b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3 . conv2dUnpadded3) (FTKR [3, 3, 3, 3] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 1] + sfromR u1 !$ [0, 1, 1, 1])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 1] + sfromR u1 !$ [0, 1, 1, 1])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0)))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let t80 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) in rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t80)))))) [0, 0, 0, 1] + soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t80)))))) [0, 1, 1, 1])"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (let t80 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) in soneHot (t80 !$ [0, 0, 0]) [0, 0, 0, 1] + soneHot (t80 !$ [0, 0, 0]) [0, 1, 1, 1])"++maxPool2dUnpadded3+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded3 arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez3 [2, 2, 2, 2] arr [iBh `quotH` 2, aa, bb, iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded3: impossible pattern needlessly required"++conv2dUnpadded3+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded3 arrA =+  let shB = [2, 2, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez4 shB arrA [iImg, 0, iBw, 1]+      in rindex0 arrAt [0, iBw, iImg, iBh] + rindex0 arrAt [iImg, 1, iBw + 1, iBh]+    _ -> error "conv2dUnpadded3: impossible pattern needlessly required"++slicez3+  :: (ADReady target, GoodScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez3 shOut d ixBase =+  rbuild shOut $ \_ -> indexz03 d (ixrZipWith (+) ixBase ixBase)++indexz03+  :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)+  => target (TKR n r) -> IxROf target n -> target (TKR 0 r)+indexz03 d ix = ifH (within3 @target (rshape @target d) ix) (d ! ix) (rscalar 0)++within3+  :: forall target n. (ADReady target, KnownNat n)+  => IShR n -> IxROf target n -> BoolOf target+within3 sh ix =+  let within :: IntOf target -> IntOf target -> BoolOf target+      within i dim = 0 ==. i ||* dim - 2 ==. i+  in foldr (&&*) true+     $ zipWith within (toList ix) (map fromIntegral $ toList sh)++rmaximum3 :: (BaseTensor target, LetTensor target, KnownNat n, GoodScalar r)+         => target (TKR n r) -> target (TKR 0 r)+rmaximum3 t0 = tlet t0 $ \t -> rindex0 t [0, 0, 0, 0]++testCNNOPP4 :: Assertion+testCNNOPP4 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @3) knownSTK+                   $ AstReplicate (SNat @3) knownSTK+                   $ AstReplicate (SNat @3) knownSTK+                   $ AstReplicate (SNat @3) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded4 blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (let w77 = stranspose @[5,6,7,8,0,1,2,3,4] (sfromVector (fromList [stranspose @[2,3,4,0,5,6,7,1] (sgather (stranspose @[6,0,3,1,4,5,2] (sgather (stranspose @[4,1,3,0,2] (sgather (stranspose @[3,0,4,1,2] (sgather (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i52, i55] -> [i52 + i55]))) (\\[i58, i60] -> [2 + (negate i60 + i58), i60]))) (\\[i62, i64, i67] -> [i62 * i64 + i67]))) (\\[i22, i8] -> [2 * i22 + i8])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0, 0] in stranspose @[1,2,3,0] (sappend (sreplicate @1 (str (sappend (sreplicate @1 (stranspose @[0,2,4,1,3] w77 !$ [1, 0, 0])) (sreplicate @1 (stranspose @[0,2,4,1,3] w77 !$ [1, 1, 0]))))) (sreplicate @1 (stranspose @[0,4,1,2,3] w77 !$ [1, 1]))))"+      -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printAstPretty afcnn2T+    @?= "rfromS (let w19 = sgather (sfromVector (fromList [stranspose @[2,3,4,0,5,6,7,1] (sgather (stranspose @[6,0,3,1,4,5,2] (sgather (stranspose @[3,0,2,1] (sgather (stranspose @[0,2,1] (sgather (sconcrete (sreplicate [2,3,3,3] 7.0)) (\\[i32, i5] -> [i32 + i5]))) (\\[i31, i6] -> [i31, 2 + (negate i31 + i6)]))) (\\[i36, i26, i7] -> [i36 * i26 + i7]))) (\\[i22, i8] -> [2 * i22 + i8])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) (\\[i34, i28, i24, i21, i15, i12, i10, i9] -> [ifH (0 <=. negate i28 + negate i15 &&* (1 <=. i28 + negate i12 &&* (notB (notB (0 <=. negate i34 * i24 + negate i10) &&* notB (1 <=. i34 * i24 + i10 &&* (-1) <=. negate i34 * i24 + negate i10)) &&* notB (notB (0 <=. (-2) * i21 + negate i9) &&* notB (1 <=. 2 * i21 + i9 &&* (-1) <=. (-2) * i21 + negate i9))))) 0 1, i34, i28, i24, i21, i15, i12, i10, i9]) in stranspose @[4,5,6,7,0,1,2,3] w19 !$ [0, 0, 0, 0])"++-- In this test primal is trivial but gradient is not, so we know how much+-- scatters should be able to simplify in the future.+testCNNOPP4b :: Assertion+testCNNOPP4b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent maxPool2dUnpadded4 (FTKR [3, 3, 3, 3] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (let w120 = stranspose @[5,6,7,8,0,1,2,3,4] (sfromVector (fromList [stranspose @[2,3,4,0,5,6,7,1] (sgather (stranspose @[4,7,0,3,1,5,6,2] (sgather (stranspose @[3,4,7,1,5,6,0,2] (sgather (stranspose @[6,0,7,4,3,2,1,5] (sgather (sslice (SNat @1) (SNat @2) (stranspose @[0,2,1] (sfromR u1))) (\\[i241, i243, i244, i245, i246] -> [i241 + i243]))) (\\[i252, i255] -> [2 + (negate i255 + i252), i255]))) (\\[i258, i260, i263] -> [i258, i260, i258 * i260 + i263]))) (\\[i118, i119] -> [i118, 2 * i118 + i119])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0, 0] in stranspose @[1,2,3,0] (sappend (sreplicate @1 (str (sappend (sreplicate @1 (stranspose @[0,2,4,1,3] w120 !$ [1, 0, 0])) (sreplicate @1 (stranspose @[0,2,4,1,3] w120 !$ [1, 1, 0]))))) (sreplicate @1 (stranspose @[0,4,1,2,3] w120 !$ [1, 1]))))"+      -- TODO: was once "\\u1 -> rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w120 = stranspose @[5,6,7,8,0,1,2,3,4] (sfromVector (fromList [stranspose @[2,3,4,0,5,6,7,1] (sgather (stranspose @[4,7,0,3,1,5,6,2] (sgather (stranspose @[3,4,7,1,5,6,0,2] (sgather (stranspose @[6,0,7,4,3,2,1,5] (sgather (sslice (SNat @1) (SNat @2) (stranspose @[0,2,1] (sfromR u1))) (\\[i108, i109, i110, i111, i112] -> [i108 + i109]))) (\\[i113, i114] -> [2 + (negate i114 + i113), i114]))) (\\[i115, i116, i117] -> [i115, i116, i115 * i116 + i117]))) (\\[i118, i119] -> [i118, 2 * i118 + i119])), sconcrete (sreplicate [2,2,2,2,2,2,2,2] 0.0)])) !$ [0, 0, 0, 0] in rfromS (stranspose @[1,2,3,0] (sappend (sreplicate @1 (str (sappend (sreplicate @1 (stranspose @[0,2,4,1,3] w120 !$ [1, 0, 0])) (sreplicate @1 (stranspose @[0,2,4,1,3] w120 !$ [1, 1, 0]))))) (sreplicate @1 (stranspose @[0,4,1,2,3] w120 !$ [1, 1]))))"+      -- TODO: was once "\\u1 -> rfromS (sconcrete (sfromListLinear [2,2,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w122 = stranspose @[4,5,6,7,8,0,1,2,3] (soneHot (stranspose @[0,3,1,4,2] (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (str (ssum @1 (sslice (SNat @0) (SNat @1) (stranspose @[3,0,1,2] (sfromR dret))))))) [1, 0, 0]) + (stranspose @[0,3,1,4,2] (soneHot (ssum @1 (sslice (SNat @1) (SNat @1) (str (ssum @1 (sslice (SNat @0) (SNat @1) (stranspose @[3,0,1,2] (sfromR dret))))))) [1, 1, 0]) + stranspose @[0,2,3,4,1] (soneHot (ssum @1 (sslice (SNat @1) (SNat @1) (stranspose @[3,0,1,2] (sfromR dret)))) [1, 1]))) [0, 0, 0, 0]) in rfromS (stranspose @[0,2,1] (sappend (sconcrete (sreplicate [1,3,3,3] 0.0)) (sappend (sscatter (stranspose @[1,6,5,4,3,7,0,2] (sscatter (stranspose @[6,3,7,0,1,4,5,2] (sscatter (stranspose @[2,4,7,3,0,5,6,1] (sscatter (stranspose @[3,7,0,1,2,4,5,6] (w122 !$ [0])) (\\[i123, i124] -> [i123, 2 * i123 + i124]))) (\\[i125, i126, i127] -> [i125, i126, i125 * i126 + i127]))) (\\[i128, i129] -> [2 + (negate i129 + i128), i129]))) (\\[i130, i131, i132, i133, i134] -> [i130 + i131])) (sconcrete (sfromListLinear [0,3,3,3] [])))))"+      -- TODO: was once "\\dret u1 -> rfromS (sconcrete (sreplicate [3,3,3,3] 0.0))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (sappend (sconcrete (sreplicate [1,3,3,3] 0.0)) (stranspose @[0,2,1] (sscatter (stranspose @[1,6,5,4,3,7,0,2] (sscatter (stranspose @[6,3,7,0,1,4,5,2] (sscatter (stranspose @[2,4,7,3,0,5,6,1] (sscatter (stranspose @[4,8,3,5,6,7,0,1,2] (soneHot (stranspose @[0,3,1,4,2] (soneHot (stranspose @[1,3,0,2] (sfromR dret) !$ [0, 0]) [1, 0, 0]) + (stranspose @[0,3,1,4,2] (soneHot (stranspose @[1,3,0,2] (sfromR dret) !$ [1, 0]) [1, 1, 0]) + stranspose @[0,2,3,4,1] (soneHot (stranspose @[3,0,1,2] (sfromR dret) !$ [1]) [1, 1]))) [0, 0, 0, 0]) !$ [0]) (\\[i123, i124] -> [i123, 2 * i123 + i124]))) (\\[i125, i126, i127] -> [i125, i126, i125 * i126 + i127]))) (\\[i128, i129] -> [2 + (negate i129 + i128), i129]))) (\\[i130, i131, i132, i133, i134] -> [i130 + i131]))))"+      -- TODO: was once "\\dret u1 -> rfromS (sconcrete (sreplicate [3,3,3,3] 0.0))"++testCNNOPP5 :: Assertion+testCNNOPP5 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @6) knownSTK+                   $ AstReplicate (SNat @6) knownSTK+                   $ AstReplicate (SNat @6) knownSTK+                   $ AstReplicate (SNat @6) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = conv2dUnpadded4 blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (sconcrete (sfromListLinear [1,1,2,2] [7.0,0.0,0.0,0.0]))"+  printAstPretty afcnn2T+    @?= "rfromS (sconcrete (sfromListLinear [1,1,2,2] [7.0,0.0,0.0,0.0]))"++-- In this test primal is trivial but gradient is not, so we know how much+-- scatters should be able to simplify in the future.+testCNNOPP5b :: Assertion+testCNNOPP5b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent conv2dUnpadded4 (FTKR [5, 5, 5, 5] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (sreplicate @1 (sreplicate @1 (str (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (sreplicate @1 (sreplicate @1 (str (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (str (ssum @1 (ssum @1 (sfromR dret)))))))) [0, 0, 0, 0])"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (sfromR dret !$ [0, 0, 0, 0]) [0, 0, 0, 0])"++maxPool2dUnpadded4+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded4 arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez4 [2, 2, 2, 2] arr [bb + 1, 2 - bb, aa * iBh, 2 * iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded4: impossible pattern needlessly required"++conv2dUnpadded4+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded4 arrA =+  let shB = [1, 1, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez4 shB arrA [iImg, 0, iBh, iBw]+      in rindex0 arrAt [0, 0, 0, 0]+    _ -> error "conv2dUnpadded4: impossible pattern needlessly required"++slicez4+  :: (ADReady target, GoodScalar r, KnownNat n)+  => IShR n -> target (TKR n r) -> IxROf target n -> target (TKR n r)+slicez4 shOut d ixBase =+  rbuild shOut $ \ixResult -> indexz03 d (ixrZipWith (+) ixBase ixResult)++testCNNOPP6 :: Assertion+testCNNOPP6 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded3 $ conv2dUnpadded3z blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (let t45 = str (sfromVector (fromList [stranspose @[1,2,0] (sgather (stranspose @[2,1,0] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sfromListLinear [2,2,2] [7.0,7.0,0.0,0.0,7.0,7.0,0.0,0.0])) (\\[i48] -> [2 * i48]))) (\\[i51] -> [2 * i51]))) (\\[i4] -> [2 * i4])), sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0] in sappend (sreplicate @1 (sappend (sreplicate @1 (t45 !$ [0, 0, 0])) (sreplicate @1 (t45 !$ [1, 0, 1])))) (sreplicate @1 (t45 !$ [1, 1])))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+      -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printAstPretty afcnn2T+    @?= "rfromS (stranspose @[1,2,0] (sreplicate @2 (sgather (sfromVector (fromList [stranspose @[1,2,0] (sgather (stranspose @[2,1,0] (sgather (str (sgather (sreplicate @2 (str (sreplicate @2 (sgather (sfromVector (fromList [str (sgather (str (sgather (sconcrete (sreplicate [2,2,2,2] 7.0)) (\\[i9] -> [2 * i9, 2 * i9, 2 * i9]))) (\\[i12] -> [2 * i12])), sconcrete (sreplicate [2,2] 0.0)])) (\\[i19, i17] -> [let i18 = ifH (notB (notB (0 <=. negate i19) &&* notB (0 <=. negate i19)) &&* (notB (notB (0 <=. negate i19) &&* notB (0 <=. negate i19)) &&* notB (notB (0 <=. negate i17) &&* notB (0 <=. negate i17)))) 0 1 in ifH (0 <=. negate i19) i18 (ifH (0 <=. negate i19) i18 1), i19, i17]))))) (\\[i1] -> [2 * i1, 0]))) (\\[i2] -> [2 * i2]))) (\\[i4] -> [2 * i4])), sconcrete (sreplicate [2,2,2] 0.0)])) (\\[i26, i24, i22] -> [let i23 = ifH (notB (notB (0 <=. negate i24) &&* notB (0 <=. negate i24)) &&* notB (notB (0 <=. negate i22) &&* notB (0 <=. negate i22))) 0 1 in ifH (0 <=. negate i26) i23 (ifH (0 <=. negate i26) i23 1), i26, i24, i22]))))"++testCNNOPP6b :: Assertion+testCNNOPP6b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3 . conv2dUnpadded3z) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (let t58 = str (sfromVector (fromList [stranspose @[1,2,0] (sgather (stranspose @[2,1,0] (sgather (stranspose @[2,0,1] (sgather (stranspose @[1,2,0] (sreplicate @2 (sreplicate @2 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))))) (\\[i93] -> [2 * i93]))) (\\[i96] -> [2 * i96]))) (\\[i57] -> [2 * i57])), sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0] in stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (t58 !$ [0, 0, 0])) (sreplicate @1 (t58 !$ [1, 0, 1])))) (sreplicate @1 (t58 !$ [1, 1])))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let t58 = str (sfromVector (fromList [stranspose @[1,2,0] (sgather (stranspose @[2,1,0] (sgather (stranspose @[2,0,1] (sgather (stranspose @[1,2,0] (sreplicate @2 (sreplicate @2 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))))) (\\[i55] -> [2 * i55]))) (\\[i56] -> [2 * i56]))) (\\[i57] -> [2 * i57])), sconcrete (sreplicate [2,2,2] 0.0)])) !$ [0] in rfromS (stranspose @[1,2,0] (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sappend (sreplicate @1 (t58 !$ [0, 0, 0])) (sreplicate @1 (t58 !$ [1, 0, 1])))) (sreplicate @1 (t58 !$ [1, 1])))) (sconcrete (sreplicate [1,2,2] 0.0)))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let t60 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) ; u61 = str (soneHot (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t60)))))) [0, 0, 0] + (soneHot (ssum @1 (sslice (SNat @1) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t60)))))) [1, 0, 1] + soneHot (ssum @1 (sslice (SNat @1) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) t60)))) [1, 1])) [0]) ; v65 = ssum @2 (ssum @2 (stranspose @[2,0,1] (sscatter (stranspose @[1,2,0] (sscatter (stranspose @[2,1,0] (sscatter (stranspose @[2,0,1] (u61 !$ [0])) (\\[i62] -> [2 * i62]))) (\\[i63] -> [2 * i63]))) (\\[i64] -> [2 * i64])))) in rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) v65)) [0, 0, 0, 0])"++  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (let t60 = ssum @2 (stranspose @[2,0,1] (sfromR dret)) in soneHot (ssum0 (stranspose @[1,2,0] (sscatter (stranspose @[1,2,0] (sscatter (stranspose @[2,1,0] (sscatter (stranspose @[1,3,0,2] (soneHot (soneHot (t60 !$ [0, 0, 0]) [0, 0, 0] + (soneHot (t60 !$ [0, 0, 1]) [1, 0, 1] + soneHot (t60 !$ [0, 1]) [1, 1])) [0]) !$ [0]) (\\[i62] -> [2 * i62]))) (\\[i63] -> [2 * i63]))) (\\[i64] -> [2 * i64])) !$ [0])) [0, 0, 0, 0])"+      -- TODO: was once "\\dret u1 -> rfromS (soneHot (ssum0 (stranspose @[0,1,3,2] (sfromR dret) !$ [0, 0, 0])) [0, 0, 0, 0])"++conv2dUnpadded3z+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded3z arrA =+  let shB = [2, 2, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBw]+      in rindex0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dUnpadded3z: impossible pattern needlessly required"++testCNNOPP7 :: Assertion+testCNNOPP7 = do+  resetVarCounter+  let blackGlyph :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      blackGlyph = AstFromPrimal $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: AstTensor AstMethodLet FullSpan (TKR 4 Double)+      afcnn2T = maxPool2dUnpadded3y $ conv2dUnpadded3y blackGlyph+  printAstPretty (simplifyInlineContract afcnn2T)+    @?= "rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sgather (stranspose @[3,2,1,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sgather (stranspose @[3,2,0,1] (sgather (stranspose @[2,3,0,1] (sgather (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,0.0,0.0,7.0,0.0,0.0,0.0,7.0,0.0,0.0,0.0,7.0,0.0,0.0,0.0])) (\\[i49] -> [2 * i49]))) (\\[i52, i53] -> [2 * i53, 2 * i52]))) (\\[i4] -> [2 * i4])), sconcrete (sreplicate [2,2,2,2] 0.0)])) !$ [0]) (\\[i29, i26, i22] -> [i26, i29, ifH (0 <=. negate i29) (ifH (notB (notB (0 <=. negate i26) &&* notB (0 <=. negate i26)) &&* notB (notB (0 <=. negate i22) &&* notB (0 <=. negate i22))) 0 1) (ifH (0 <=. negate i29) (ifH (notB (notB (0 <=. negate i26) &&* notB (0 <=. negate i26)) &&* notB (notB (0 <=. negate i22) &&* notB (0 <=. negate i22))) 0 1) 1), i22]))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"+      -- TODO: was once "rfromS (sconcrete (sfromListLinear [2,2,2,2] [7.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]))"+  printAstPretty afcnn2T+    @?= "rfromS (sgather (sfromVector (fromList [stranspose @[1,2,3,0] (sgather (stranspose @[3,2,0,1] (sgather (stranspose @[1,2,0] (sgather (sreplicate @2 (stranspose @[1,2,0] (sreplicate @2 (sgather (sfromVector (fromList [str (sgather (str (sgather (sconcrete (sreplicate [2,2,2,2] 7.0)) (\\[i9] -> [2 * i9, 2 * i9, 2 * i9]))) (\\[i11] -> [2 * i11])), sconcrete (sreplicate [2,2] 0.0)])) (\\[i19, i17] -> [let i18 = ifH (notB (notB (0 <=. negate i19) &&* notB (0 <=. negate i19)) &&* (notB (notB (0 <=. negate i19) &&* notB (0 <=. negate i19)) &&* notB (notB (0 <=. negate i17) &&* notB (0 <=. negate i17)))) 0 1 in ifH (0 <=. negate i19) i18 (ifH (0 <=. negate i19) i18 1), i19, i17]))))) (\\[i1] -> [2 * i1]))) (\\[i28, i3] -> [2 * i3, 2 * i28]))) (\\[i4] -> [2 * i4])), sconcrete (sreplicate [2,2,2,2] 0.0)])) (\\[i29, i26, i24, i22] -> [let i23 = ifH (notB (notB (0 <=. negate i29) &&* notB (0 <=. negate i29)) &&* (notB (notB (0 <=. negate i26) &&* notB (0 <=. negate i26)) &&* notB (notB (0 <=. negate i22) &&* notB (0 <=. negate i22)))) 0 1 in ifH (0 <=. negate i24) i23 (ifH (0 <=. negate i24) i23 1), i29, i26, i24, i22]))"++testCNNOPP7b :: Assertion+testCNNOPP7b = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpadded3y . conv2dUnpadded3y) (FTKR [2, 2, 2, 2] (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sgather (stranspose @[3,2,1,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sgather (stranspose @[3,2,0,1] (sgather (stranspose @[2,3,0,1] (sgather (sreplicate @2 (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0))))) (\\[i119] -> [2 * i119]))) (\\[i122, i123] -> [2 * i123, 2 * i122]))) (\\[i62] -> [2 * i62])), sconcrete (sreplicate [2,2,2,2] 0.0)])) !$ [0]) (\\[i63, i64, i65] -> [i64, i63, ifH (0 <=. negate i63) (ifH (notB (notB (0 <=. negate i64) &&* notB (0 <=. negate i64)) &&* notB (notB (0 <=. negate i65) &&* notB (0 <=. negate i65))) 0 1) (ifH (0 <=. negate i63) (ifH (notB (notB (0 <=. negate i64) &&* notB (0 <=. negate i64)) &&* notB (notB (0 <=. negate i65) &&* notB (0 <=. negate i65))) 0 1) 1), i65]))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> rfromS (stranspose @[1,2,0] (sappend (sreplicate @1 (sgather (stranspose @[3,2,1,0] (sfromVector (fromList [stranspose @[1,2,3,0] (sgather (stranspose @[3,2,0,1] (sgather (stranspose @[2,3,0,1] (sgather (sreplicate @2 (sreplicate @2 (sappend (sreplicate @1 (sappend (sreplicate @1 (sfromR u1 !$ [0, 0, 0, 0])) (sconcrete (sfromListLinear [1] [0.0])))) (sconcrete (sreplicate [1,2] 0.0))))) (\\[i59] -> [2 * i59]))) (\\[i60, i61] -> [2 * i61, 2 * i60]))) (\\[i62] -> [2 * i62])), sconcrete (sreplicate [2,2,2,2] 0.0)])) !$ [0]) (\\[i63, i64, i65] -> [i64, i63, ifH (0 <=. negate i63) (ifH (notB (notB (0 <=. negate i64) &&* notB (0 <=. negate i64)) &&* notB (notB (0 <=. negate i65) &&* notB (0 <=. negate i65))) 0 1) (ifH (0 <=. negate i63) (ifH (notB (notB (0 <=. negate i64) &&* notB (0 <=. negate i64)) &&* notB (notB (0 <=. negate i65) &&* notB (0 <=. negate i65))) 0 1) 1), i65]))) (sconcrete (sreplicate [1,2,2,2] 0.0))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w70 = stranspose @[3,2,1,0] (soneHot (sscatter (ssum @1 (sslice (SNat @0) (SNat @1) (stranspose @[2,0,1] (sfromR dret)))) (\\[i67, i68, i69] -> [i68, i67, ifH (0 <=. negate i67) (ifH (notB (notB (0 <=. negate i68) &&* notB (0 <=. negate i68)) &&* notB (notB (0 <=. negate i69) &&* notB (0 <=. negate i69))) 0 1) (ifH (0 <=. negate i67) (ifH (notB (notB (0 <=. negate i68) &&* notB (0 <=. negate i68)) &&* notB (notB (0 <=. negate i69) &&* notB (0 <=. negate i69))) 0 1) 1), i69])) [0]) ; m75 = ssum @2 (ssum @2 (sscatter (stranspose @[2,3,0,1] (sscatter (stranspose @[2,3,1,0] (sscatter (stranspose @[3,0,1,2] (w70 !$ [0])) (\\[i71] -> [2 * i71]))) (\\[i72, i73] -> [2 * i73, 2 * i72]))) (\\[i74] -> [2 * i74]))) in rfromS (soneHot (ssum @1 (sslice (SNat @0) (SNat @1) (ssum @1 (sslice (SNat @0) (SNat @1) m75)))) [0, 0, 0, 0])"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret u1 -> rfromS (soneHot (ssum0 (stranspose @[2,3,0,1] (sscatter (stranspose @[2,3,0,1] (sscatter (stranspose @[2,3,1,0] (sscatter (stranspose @[3,4,2,1,0] (soneHot (sscatter (stranspose @[2,0,1] (sfromR dret) !$ [0]) (\\[i67, i68, i69] -> [i68, i67, ifH (0 <=. negate i67) (ifH (notB (notB (0 <=. negate i68) &&* notB (0 <=. negate i68)) &&* notB (notB (0 <=. negate i69) &&* notB (0 <=. negate i69))) 0 1) (ifH (0 <=. negate i67) (ifH (notB (notB (0 <=. negate i68) &&* notB (0 <=. negate i68)) &&* notB (notB (0 <=. negate i69) &&* notB (0 <=. negate i69))) 0 1) 1), i69])) [0]) !$ [0]) (\\[i71] -> [2 * i71]))) (\\[i72, i73] -> [2 * i73, 2 * i72]))) (\\[i74] -> [2 * i74])) !$ [0, 0])) [0, 0, 0, 0])"+      -- TODO: was once "\\dret u1 -> rfromS (soneHot (sfromR dret !$ [0, 0, 0, 0]) [0, 0, 0, 0])"++maxPool2dUnpadded3y+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+maxPool2dUnpadded3y arr =+  rbuild [2, 2, 2, 2] $ \case+    [aa, bb, iBh, iBw] ->+      let arrt = slicez3 [2, 2, 2, 2] arr [iBh, aa, bb, iBw]+      in rmaximum3 arrt+    _ -> error "maxPool2dUnpadded3y: impossible pattern needlessly required"++conv2dUnpadded3y+  :: (ADReady target, GoodScalar r)+  => target (TKR 4 r) -> target (TKR 4 r)+conv2dUnpadded3y arrA =+  let shB = [2, 2, 2, 2]+  in rbuild shB $ \case+    [iImg, _, iBh, iBw] ->+      let arrAt = slicez3 shB arrA [iImg, iImg, iImg, iBh]+      in rindex0 arrAt [iBh, iBw, iImg, iBh]+    _ -> error "conv2dUnpadded3y: impossible pattern needlessly required"++-- This test uses a disastrous version of smaximum, but shows how+-- smaxIndex gets non-trivially vectorized, preserving sharing, too.+testCNNOPP4bU :: Assertion+testCNNOPP4bU = do+  resetVarCounter+  let artifactRev = revArtifactAdapt UseIncomingCotangent (maxPool2dUnpaddedS4 @4 @2) (FTKS (SNat @31 :$$ SNat @31 :$$ SNat @31 :$$ SNat @31 :$$ ZSS) (FTKScalar @Double))+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> let w52 = sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] u1) (\\[i99, i100] -> [2 * i99 + i100]))) (\\[i50, i51] -> [2 * i50 + i51]) ; u53 = smaxIndex (sreshape @[31,31,15,15,16] (stranspose @[2,3,4,0,5,1] w52)) in sgather w52 (\\[i54, i55, i56, i57] -> [i57, remH (kfromS (u53 !$ [i54, i55, i56, i57])) 4, i54, i55, i56, remH (quotH (kfromS (u53 !$ [i54, i55, i56, i57])) 4) 4])"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w52 = sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] u1) (\\[i48, i49] -> [2 * i48 + i49]))) (\\[i50, i51] -> [2 * i50 + i51]) ; u53 = smaxIndex (sreshape @[31,31,15,15,16] (stranspose @[2,3,4,0,5,1] w52)) in sgather w52 (\\[i54, i55, i56, i57] -> [i57, remH (kfromS (u53 !$ [i54, i55, i56, i57])) 4, i54, i55, i56, remH (quotH (kfromS (u53 !$ [i54, i55, i56, i57])) 4) 4])"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w52 = sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] u1) (\\[i48, i49] -> [2 * i48 + i49]))) (\\[i50, i51] -> [2 * i50 + i51]) ; u53 = smaxIndex (sreshape @[31,31,15,15,16] (stranspose @[2,3,4,0,5,1] w52)) in stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (sscatter dret (\\[i59, i60, i61, i62] -> [i62, remH (kfromS (u53 !$ [i59, i60, i61, i62])) 4, i59, i60, i61, remH (quotH (kfromS (u53 !$ [i59, i60, i61, i62])) 4) 4])) (\\[i63, i64] -> [2 * i63 + i64]))) (\\[i65, i66] -> [2 * i65 + i66]))"+-- TODO: different test result with GHC 9.10:   printArtifactPretty (simplifyArtifact artifactRev)+--    @?= "\\dret u1 -> let u53 = smaxIndex (sreshape @[31,31,15,15,16] (stranspose @[2,3,4,0,5,1] (sgather (stranspose @[4,2,3,0,1] (sgather (stranspose @[2,0,1] u1) (\\[i82, i83] -> [2 * i82 + i83]))) (\\[i50, i51] -> [2 * i50 + i51])))) in stranspose @[1,2,0] (sscatter (stranspose @[3,4,1,2,0] (sscatter (sscatter dret (\\[i59, i60, i61, i62] -> [i62, remH (kfromS (u53 !$ [i59, i60, i61, i62])) 4, i59, i60, i61, remH (quotH (kfromS (u53 !$ [i59, i60, i61, i62])) 4) 4])) (\\[i63, i64] -> [2 * i63 + i64]))) (\\[i65, i66] -> [2 * i65 + i66]))"++smaximum4 :: forall r sh target. (ADReady target, GoodScalar r, KnownShS sh)+          => target (TKS sh r) -> target (TKS '[] r)+smaximum4 t0 =+  tlet t0 $ \t ->+  ttletPrimal (tprimalPart $ kfromS $ smaxIndex (sflatten t)) $ \maxIndex ->+    sindex0 t+    $ fromLinearIdxS (tprimalPart @target . kconcrete . fromIntegral)+                     (sshape t)+                     maxIndex++maxPool2dUnpaddedS4+  :: forall ksize stride batch_size channels h w target r shOut shK1.+     ( KnownNat ksize, KnownNat stride, KnownNat batch_size, KnownNat channels+     , KnownNat h, KnownNat w+     , 1 <= stride  -- wrongly reported as redundant due to plugins+     , ADReady target, GoodScalar r+     , shOut ~ '[batch_size, channels, h `Div` stride, w `Div` stride]+     , shK1 ~ '[1, 1, ksize, ksize]+     )+  => target (TKS '[batch_size, channels, h, w] r)+  -> target (TKS shOut r)+maxPool2dUnpaddedS4 arr =+  let stride = valueOf @stride :: Int+  in sbuild @(Rank shOut) $ \case+    [iImg, iChan, iBh, iBw] ->+      smaximum4 $ slicezS @shK1 arr [ iImg, iChan+                                    , fromIntegral stride * iBh+                                    , fromIntegral stride * iBw ]+    _ -> error "maxPool2dUnpaddedS4: impossible pattern needlessly required"
+ test/simplified/TestHighRankSimplified.hs view
@@ -0,0 +1,762 @@+{-# LANGUAGE OverloadedLists #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Assorted mostly high rank tensor tests.+module TestHighRankSimplified (testTrees) where++import Prelude++import Data.Int (Int64)+import GHC.Exts (IsList (..))+import GHC.TypeLits (KnownNat, type (+), type (-), type (<=))+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Mixed.Shape+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape++import HordeAd+import HordeAd.Core.Ops (tD, tfromPrimal)++import CrossTesting+import EqEpsilon++testTrees :: [TestTree]+testTrees =+  [ testCase "3foo" testFoo+  , testCase "3bar" testBar+--  , testCase "3barS" testBarS+  , testCase "3fooD T Double [1.1, 2.2, 3.3]" testFooD+  , testCase "3fooBuild0" testFooBuild0+  , testCase "3fooBuildOut" testFooBuildOut+  , testCase "3fooBuild91" testFooBuild91+  , testCase "3fooBuild92" testFooBuild92+  , testCase "3fooBuild21" testFooBuild21+  , testCase "3fooBuild25" testFooBuild25+  , testCase "3fooBuild21S" testFooBuild21S+  , testCase "3fooBuild25S" testFooBuild25S+  , testCase "3fooBuildNest21S" testFooBuildNest21S+  , testCase "3fooBuildNest25S" testFooBuildNest25S+  , testCase "3fooBuild3" testFooBuild3+  , testCase "3fooBuildDt" testFooBuildDt+  , testCase "3fooBuildDt2" testFooBuildDt2+  , testCase "3fooBuild5" testFooBuild5+  , testCase "3fooBuild1" testFooBuild1+  , testCase "3fooMap" testFooMap+  , testCase "3fooMap1" testFooMap1+  , testCase "3fooNoGo" testFooNoGo+  , testCase "3fooNoGo10" testFooNoGo10+  , testCase "3nestedBuildMap1" testNestedBuildMap1+  , testCase "3nestedBuildMap10" testNestedBuildMap10+  , testCase "3nestedBuildMap11" testNestedBuildMap11+--  , testCase "3nestedBuildMap7" testNestedBuildMap7+  , testCase "3nestedSumBuild1" testNestedSumBuild1+--  , testCase "3nestedSumBuild5" testNestedSumBuild5+  , testCase "3nestedSumBuildB" testNestedSumBuildB+  , testCase "3nestedBuildIndex" testNestedBuildIndex+  , testCase "3barReluADValDt" testBarReluADValDt+  , testCase "3barReluADValDt2" testBarReluADValDt2+  , testCase "3barReluADVal" testBarReluADVal+  , testCase "3barReluADVal3" testBarReluADVal3+  , testCase "3braidedBuilds" testBraidedBuilds+  , testCase "3braidedBuilds1" testBraidedBuilds1+  , testCase "3recycled" testRecycled+-- takes too long (can't be helped)  , testCase "3recycled1" testRecycled1+  , testCase "3concatBuild0" testConcatBuild0+  , testCase "3concatBuild1" testConcatBuild1+  , testCase "3concatBuild0m" testConcatBuild0m+  , testCase "3concatBuild1m" testConcatBuild1m+  , testCase "3concatBuild2" testConcatBuild2+  , testCase "3concatBuild22" testConcatBuild22+  , testCase "3concatBuild3" testConcatBuild3+  , testCase "3logistic0" testLogistic0+  , testCase "3logistic5" testLogistic5+  , testCase "3logistic52" testLogistic52+  , testCase "3logistic0Old" testLogistic0Old+  , testCase "3logistic5Old" testLogistic5Old+  , testCase "3logistic52Old" testLogistic52Old+  , testCase "3logisticA0" testLogisticA0+  , testCase "3logisticB0" testLogisticB0+  , testCase "3logisticC0" testLogisticC0+  ]++foo :: RealFloatH a => (a,a,a) -> a+foo (x,y,z) =+  let w = x * sin y+  in atan2H z w + z * w++_fooF :: RealFloatH a => (a,a,a) -> a+_fooF (x,y,z) =+  let w = x * sin y+  in atan2H z w + z * w++testFoo :: Assertion+testFoo =+  assertEqualUpToEpsilon 1e-3+    (ringestData [2,2,1, 2,2] [-4.6947093,1.5697206,-1.6332961,0.34882763,1.5697206,-1.0,-0.9784988,-0.9158946,6.6326222,3.6699238,7.85237,-2.9069107,17.976654,0.3914159,32.98194,19.807974], ringestData [2,2,1, 2,2] [6.943779,-1.436789,33.67549,0.22397964,-1.436789,-1.0,-0.975235,-0.90365005,147.06645,-73.022705,-9.238474,-10.042692,-980.2843,-7.900571,-14.451739,436.9084], ringestData [2,2,1, 2,2] [-4.8945336,2.067469,-1.7196897,1.3341143,2.067469,1.0,0.99846554,0.99536234,6.6943173,3.7482092,7.977362,-3.1475093,18.000969,0.48736274,33.01224,19.845064])+    (grad (kfromR . rsum0 @5 @(TKScalar Float) . foo) (t16, t16, t16))++bar :: forall a. RealFloatH a => (a, a) -> a+bar (x, y) =+  let w = foo (x, y, x) * sin y+  in atan2H x w + y * w++_barF :: forall a. RealFloatH a => (a, a) -> a+_barF (x, y) =+  let w = _fooF (x, y, x) * sin y+  in atan2H x w + y * w++testBar :: Assertion+testBar =+  assertEqualUpToEpsilon 1e-5+    (Concrete $ Nested.rfromListLinear [3,1,2,2,1,2,2] [304.13867,914.9335,823.0187,1464.4688,5264.3306,1790.0055,1535.4309,3541.6572,304.13867,914.9335,823.0187,1464.4688,6632.4355,6047.113,1535.4309,1346.6815,45.92141,6.4903135,5.5406737,1.4242969,6.4903135,1.1458766,4.6446533,2.3550234,88.783676,27.467598,125.27507,18.177452,647.1915,0.3878851,2177.6152,786.1792,6.4903135,6.4903135,6.4903135,6.4903135,2.3550234,2.3550234,2.3550234,2.3550234,21.783596,2.3550234,2.3550234,2.3550234,21.783596,21.783596,21.783596,21.783596],Concrete $ Nested.rfromListLinear [3,1,2,2,1,2,2] [-5728.7617,24965.113,32825.07,-63505.953,-42592.203,145994.88,-500082.5,-202480.06,-5728.7617,24965.113,32825.07,-63505.953,49494.473,-2446.7632,-500082.5,-125885.58,-43.092484,-1.9601002,-98.97709,2.1931143,-1.9601002,1.8243169,-4.0434446,-1.5266153,2020.9731,-538.0603,-84.28137,62.963814,-34986.996,-9.917454,135.30023,17741.998,-1.9601002,-1.9601002,-1.9601002,-1.9601002,-1.5266153,-1.5266153,-1.5266153,-1.5266153,-4029.1775,-1.5266153,-1.5266153,-1.5266153,-4029.1775,-4029.1775,-4029.1775,-4029.1775])+    (cgrad (kfromR . rsum0 . bar @(ADVal Concrete (TKR 7 Float))) (t48, t48))++{- TODO: divergent result; bring back when GHC 9.10 dropped:+testBarS :: Assertion+testBarS =+  assertEqualUpToEpsilon 1e-5+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[3, 1, 2, 2, 1, 2, 2] knownShS [304.13867,914.9335,823.0187,1464.4688,5264.3306,1790.0055,1535.4309,3541.6572,304.13867,914.9335,823.0187,1464.4688,6632.4355,6047.113,1535.4309,1346.6815,45.92141,6.4903135,5.5406737,1.4242969,6.4903135,1.1458766,4.6446533,2.3550234,88.783676,27.467598,125.27507,18.177452,647.1915,0.3878851,2177.6152,786.1792,6.4903135,6.4903135,6.4903135,6.4903135,2.3550234,2.3550234,2.3550234,2.3550234,21.783596,2.3550234,2.3550234,2.3550234,21.783596,21.783596,21.783596,21.783596], sconcrete $ Nested.sfromListPrimLinear @_ @'[3, 1, 2, 2, 1, 2, 2] knownShS [-5728.761,24965.113,32825.074,-63505.957,-42592.203,145994.89,-500082.5,-202480.05,-5728.761,24965.113,32825.074,-63505.957,49494.473,-2446.7632,-500082.5,-125885.58,-43.092484,-1.9601007,-98.97708,2.1931143,-1.9601007,1.8243167,-4.0434446,-1.5266151,2020.9731,-538.06036,-84.28139,62.963818,-34986.992,-9.917454,135.3003,17741.996,-1.9601007,-1.9601007,-1.9601007,-1.9601007,-1.5266151,-1.5266151,-1.5266151,-1.5266151,-4029.1775,-1.5266151,-1.5266151,-1.5266151,-4029.1775,-4029.1775,-4029.1775,-4029.1775])+    (cgrad (kfromS . ssum0 . barF @(ADVal Concrete (TKS '[3, 1, 2, 2, 1, 2, 2] Float))) (sfromR t48, sfromR t48))+-}++-- A dual-number and list-based version of a function that goes+-- from `R^3` to `R`.+fooD :: forall r n. (RealFloatH (ADVal Concrete (TKR n r)))+     => ListR 3 (ADVal Concrete (TKR n r)) -> ADVal Concrete (TKR n r)+fooD (x ::: y ::: z ::: ZR) =+  let w = x * sin y+  in atan2H z w + z * w++testFooD :: Assertion+testFooD =+  assertEqualUpToEpsilon 1e-10+    (fromList [ringestData [1,2,2,1,2,2,2,2,2,1] [18.73108960474591,20.665204824764675,25.821775835995922,18.666613887422585,34.775664100213014,62.54884873632415,37.93303229694526,11.635186977032971,18.73108960474591,20.665204824764675,25.821775835995922,20.600738734367262,34.775664100213014,62.54884873632415,16.663997008808924,3.1300339898598155,1.060799258653783,3.78942741815228,0.1889454555944933,-1.060799258653783,62.54884873632415,37.93303229694526,62.54884873632415,35.99996432769119,62.54884873632415,37.93303229694526,11.635186977032971,18.73108960474591,20.665204824764675,25.821775835995922,20.665204824764675,20.665204824764675,25.821775835995922,34.134947381491145,34.775664100213014,45527.22315787758,-4.488300547708207,2.1475176207684497,8.404498097344806,5.747373381623309,5.096832468946128,-2.4630526910399646,18.666613887422585,1.7769486222994448,-215.8115662030395,16.73214939773215,1.060799258653783,1.060799258653783,1.060799258653783,1.060799258653783,2.1475176207684497,2.1475176207684497,2.1475176207684497,2.1475176207684497,16.08742477551077,16.08742477551077,16.08742477551077,16.08742477551077,2.1475176207684497,2.1475176207684497,2.1475176207684497,2.1475176207684497,16.08742477551077,16.08742477551077,16.08742477551077,16.08742477551077,25.821775835995922,5.096832468946128,7.045006174919766,-1.7808956511653404,16.663997008744435,18.533999054066836,-25.177267779903083,16.60317012020362,25.821775835995922,5.096832468946128,7.045006174919766,-1.7808956511653404,16.663997008744435,18.533999054066836,-12.280721583745471,16.60317012020362,5.161956818274285,18.73108960474591,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,188.11000552192755,34.775664100213014,62.54884873632415,35.99996432769119,62.54884873632415,55.32933980086011,62.54884873632415,55.32933980086011,11.635186977032971,18.73108960474591,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,14.152094926881784,34.775664100213014,62.54884873632415,53.39649491503442,62.54884873632415,14.72904006548922,62.54884873632415,37.93303229694526,11.635186977032971,18.73108960474591,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,20.665204824764675,25.821775835995922,57.33025874582143,34.775664100213014,62.54884873632415,36.64432517917614,62.54884873632415,34.06684929392724,62.54884873632415,35.99996432769119], ringestData [1,2,2,1,2,2,2,2,2,1] [647.1354943759653,787.5605199613974,1229.333367336918,642.6917612678424,2229.2701397674327,7210.705208776531,2652.3459120285806,250.02943073785886,647.1354943759653,787.5605199613974,1229.333367336918,782.6578815409038,2229.2701397674327,7210.705208776531,512.2982591657892,18.580536443699742,2.518850510725482,26.993800503829114,0.2243239488720164,2.518850510725482,7210.705208776531,2652.3459120285806,7210.705208776531,2388.9603285490866,7210.705208776531,2652.3459120285806,250.02943073785886,647.1354943759653,787.5605199613974,1229.333367336918,787.5605199613974,787.5605199613974,1229.333367336918,2147.9011858437157,2229.2701397674327,-0.5405182383359878,-0.5328698165396271,-0.5099245509210925,130.7140495214786,61.4116989316311,48.40938174779479,11.696956758139343,642.6917612678424,6.317020301049852,85833.87394976329,516.4928003659018,2.518850510725482,2.518850510725482,2.518850510725482,2.518850510725482,-0.5099245509210925,-0.5099245509210925,-0.5099245509210925,-0.5099245509210925,477.4973215160379,477.4973215160379,477.4973215160379,477.4973215160379,-0.5099245509210925,-0.5099245509210925,-0.5099245509210925,-0.5099245509210925,477.4973215160379,477.4973215160379,477.4973215160379,477.4973215160379,1229.333367336918,48.40938174779479,92.00538642301063,6.3430614471479245,512.2982591618282,633.5999783697488,1168.7578661039847,508.56903530563443,1229.333367336918,48.40938174779479,92.00538642301063,6.3430614471479245,512.2982591618282,633.5999783697488,278.48156010484087,508.56903530563443,49.64077766932281,647.1354943759653,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,65212.963738386214,2229.2701397674327,7210.705208776531,2388.9603285490866,7210.705208776531,5642.338335044463,7210.705208776531,5642.338335044463,250.02943073785886,647.1354943759653,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,369.6431004072799,2229.2701397674327,7210.705208776531,5255.048317224881,7210.705208776531,400.3514287686239,7210.705208776531,2652.3459120285806,250.02943073785886,647.1354943759653,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,787.5605199613974,1229.333367336918,6057.774447242021,2229.2701397674327,7210.705208776531,2475.225838667682,7210.705208776531,2139.3419044407133,7210.705208776531,2388.9603285490866], ringestData [1,2,2,1,2,2,2,2,2,1] [18.76237979248771,20.69357069589509,25.8444826804669,18.698011972363496,34.7925278085306,62.558226125235436,37.948492946856575,11.685493300971446,18.76237979248771,20.69357069589509,25.8444826804669,20.629193248844963,34.7925278085306,62.558226125235436,16.699160877305292,3.3121428825170947,1.516071490296981,3.9411848287000124,1.0994899188808887,-1.516071490296981,62.558226125235436,37.948492946856575,62.558226125235436,36.01625479268449,62.558226125235436,37.948492946856575,11.685493300971446,18.76237979248771,20.69357069589509,25.8444826804669,20.69357069589509,20.69357069589509,25.8444826804669,34.1521274657041,34.7925278085306,-45527.22317076194,4.617144085155745,-2.4052046956635262,8.474005308282699,5.84854498865513,5.210650526856928,-2.6906888068615635,18.698011972363496,2.0810391881996813,-215.8142842462135,16.767170338627782,1.516071490296981,1.516071490296981,1.516071490296981,1.516071490296981,-2.4052046956635262,-2.4052046956635262,-2.4052046956635262,-2.4052046956635262,16.123846116986126,16.123846116986126,16.123846116986126,16.123846116986126,-2.4052046956635262,-2.4052046956635262,-2.4052046956635262,-2.4052046956635262,16.123846116986126,16.123846116986126,16.123846116986126,16.123846116986126,25.8444826804669,5.210650526856928,7.127782944309438,-2.0844104722608057,16.69916087724094,18.565621417897145,-25.200555362084323,16.638462541261234,25.8444826804669,5.210650526856928,7.127782944309438,-2.0844104722608057,16.69916087724094,18.565621417897145,-12.328394068734287,16.638462541261234,5.2743697149763085,18.76237979248771,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,188.113123824884,34.7925278085306,62.558226125235436,36.01625479268449,62.558226125235436,55.33994055377702,62.558226125235436,55.33994055377702,11.685493300971446,18.76237979248771,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,14.193483311576621,34.7925278085306,62.558226125235436,53.40747931617656,62.558226125235436,14.768811697198851,62.558226125235436,37.948492946856575,11.685493300971446,18.76237979248771,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,20.69357069589509,25.8444826804669,57.34048958248757,34.7925278085306,62.558226125235436,36.660329315674915,62.558226125235436,34.08406370302229,62.558226125235436,36.01625479268449]])+    (cgrad (kfromR . rsum0 . fooD) (fromList [ t128+               , rreplicate0N [1, 2, 2, 1, 2, 2, 2, 2, 2, 1] (rscalar (0.7 :: Double))+               , t128 ]))++fooBuild0 :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)+          => target (TKR (1 + n) r) -> target (TKR (1 + n) r)+fooBuild0 v =+  let r = rsum v+  in rbuild1 2 $ const r++testFooBuild0 :: Assertion+testFooBuild0 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,2,1,2,2] [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0])+    (rev' @Double @5 fooBuild0 t16)++fooBuildOut+  :: forall target r n. (ADReady target, GoodScalar r, KnownNat n)+  => target (TKR (1 + n) r) -> target (TKR (1 + n) r)+fooBuildOut v =+  rbuild1 2 $ \ix -> ifH (ix ==. 0)+                         (rindex v [ix + 1])  -- index out of bounds; guarded+                         (rsum v)++testFooBuildOut :: Assertion+testFooBuildOut =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,2,1,2,2] [1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0])+    (rev' @Double @5 fooBuildOut t16)++fooBuild2+  :: forall target r n.+     (ADReady target, GoodScalar r, KnownNat n, Floating (target (TKR n r)), RealFloat r)+  => target (TKR (1 + n) r) -> target (TKR (1 + n) r)+fooBuild2 v =+  rbuild1 2 $ \ix' -> let ix :: PrimalOf target (TKS '[] Int64)+                          ix = sfromR $ rfromK ix' in+    ifH (ix - (sprimalPart . sfloor . sfromR) (rsum0  @5 @(TKScalar r)+                      $ rreplicate0N [5,12,11,9,4] (rsum0 v)) - sscalar 10001 >=. sscalar 0+         &&* ix - (sprimalPart . sfloor . sfromR) (rsum0 @5 @(TKScalar r) @target+                          $ rreplicate0N [5,12,11,9,4] (rsum0 v)) - sscalar 10001 <=. sscalar 1)+        (rindex v [kfromR $ rfromS $ ix - (sprimalPart . sfloor . sfromR) (rsum0  @5 @(TKScalar r) @target+                                $ rreplicate0N [5,12,11,9,4] (rsum0 v)) - sscalar 10001])+           -- index out of bounds; also fine+        (sqrt $ abs $ rindex v [kfromS+                                $ let rr = (ix - (sfromR . rprimalPart . rfloor) (rsum0 v) - sscalar 10001) `remH` sscalar 2+                                  in ifH (signum rr ==. negate (signum $ sscalar 2))+                                     (rr + sscalar 2)+                                     rr])++fooBuild2L+  :: forall k target r n.+     (ADReady target, GoodScalar r, KnownNat n, Floating (target (TKR n r)), RealFloat r)+  => ListR k (target (TKR (1 + n) r)) -> target (TKR (1 + n) r)+fooBuild2L = foldr1 (+) . fmap fooBuild2++testFooBuild91 :: Assertion+testFooBuild91 =+  assertEqualUpToEpsilon 1e-8+    (fromList $ map (ringestData [2]) [[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299]])+    (cgrad (kfromR . rsum0 @1 . fooBuild2L @50 @(ADVal Concrete) @Double @0)+       (fromList $ map (ringestData [2]) [[0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9]]))++testFooBuild92 :: Assertion+testFooBuild92 =+  assertEqualUpToEpsilon 1e-8+    (fromList $ map (ringestData [2]) [[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299],[1.5811388300841895,1.118033988749895],[1.118033988749895,0.9128709291752769],[0.9128709291752769,0.7905694150420948],[0.7905694150420948,0.7071067811865475],[0.7071067811865475,0.6454972243679028],[0.6454972243679028,0.5976143046671968],[0.5976143046671968,0.5590169943749475],[0.5590169943749475,0.5270462766947299],[0.5270462766947299,1.5811388300841895],[1.5811388300841895,0.5270462766947299]])+    (grad+       (kfromR . rsum0 @1 . fooBuild2L @50 @(AstTensor AstMethodLet FullSpan) @Double @0)+       (fromList $ map (ringestData [2]) [[0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9], [0.1, 0.2], [0.2, 0.3], [0.3, 0.4], [0.4, 0.5], [0.5, 0.6], [0.6, 0.7], [0.7, 0.8], [0.8, 0.9], [0.9, 0.1], [0.1, 0.9]]))++testFooBuild21 :: Assertion+testFooBuild21 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2] [0.2886751345948129,0.35355339059327373])+    (rev' @Double @1 fooBuild2 (ringestData [2] [3.0,2.0]))++testFooBuild25 :: Assertion+testFooBuild25 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,2,1,2,2] [0.22360679774997896,0.35355339059327373,0.20412414523193154,0.5,-0.35355339059327373,500.0,1.5811388300841895,-1.118033988749895,0.1381447409988844,0.16666666666666666,0.17677669529663687,-0.25,8.574929257125441e-2,0.288948802391873,-8.703882797784893e-2,9.805806756909202e-2])+    (rev' @Double @5 fooBuild2 t16)++fooBuild2S+  :: forall k sh target r.+     (ADReady target, GoodScalar r, KnownNat k, Floating (target (TKS sh r)), RealFloat r, KnownShS sh)+  => target (TKS (k : sh) r) -> target (TKR (1 + Rank sh) r)+fooBuild2S v = rfromS $+  sbuild1 @2 $ \ix' -> let ix :: PrimalOf target (TKS '[] Int64)+                           ix = sfromR $ rfromK ix' in+    ifH (ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+             $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001 >=. srepl 0+         &&* ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+             $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001 <=. srepl 1)+        (sindex v ((kfromS $ ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+             $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001) :.$ ZIS ))+           -- index out of bounds; also fine+        (sqrt $ abs $ sindex v ((kfromR $ rfromS $ let rr = (ix - (sprimalPart . sfloor) (ssum0 v) - srepl 10001) `remH` srepl 2+                                in ifH (signum rr ==. negate (signum $ srepl 2))+                                   (rr + srepl 2)+                                   rr) :.$ ZIS))++testFooBuild21S :: Assertion+testFooBuild21S =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2] [0.2886751345948129,0.35355339059327373])+    (rev' @Double @1 (fooBuild2S @2 @'[] . sfromR) (ringestData [2] [3.0,2.0]))++testFooBuild25S :: Assertion+testFooBuild25S =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,2,1,2,2] [0.22360679774997896,0.35355339059327373,0.20412414523193154,0.5,-0.35355339059327373,500.0,1.5811388300841895,-1.118033988749895,0.1381447409988844,0.16666666666666666,0.17677669529663687,-0.25,8.574929257125441e-2,0.288948802391873,-8.703882797784893e-2,9.805806756909202e-2])+    (rev' @Double @5 (fooBuild2S @2 @[2, 1, 2, 2] . sfromR) t16)++fooBuildNest2S+  :: forall k sh target r.+     (ADReady target, GoodScalar r, KnownNat k, Floating (target (TKS sh r)), RealFloat r, KnownShS sh)+  => target (TKS (k : sh) r) -> target (TKR (1 + Rank sh) r)+fooBuildNest2S v = rfromS $+  sbuild1 @2 $ \ix' -> let ix :: PrimalOf target (TKS '[] Int64)+                           ix = sfromR $ rfromK ix' in+    ifH (ix - (sunNest @_ @'[] @'[] . sprimalPart . snest knownShS . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+             $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001 >=. srepl 0+         &&* ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r)+             $ sreplicate0N @[5,12,11,9,4] (ssum0 v)) - srepl 10001 <=. srepl 1)+-- TODO:        (sindex v (ShapedList.singletonIndex (ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @r  $ sunNest $ treplicate (SNat @5) knownSTK $ snest (knownShS @[12,11])+        (sindex v ((kfromR $ rfromS $ ix - (sprimalPart . sfloor) (ssum0 @[5,12,11,9,4] @(TKScalar r) @target $ sunNest $ tproject2 $ tfromPrimal knownSTK $ tpair tunit (sprimalPart $ snest (knownShS @[5,12,11])+             $ sreplicate0N @[5,12,11,9,4] (ssum0 v))) - srepl 10001) :.$ ZIS))+           -- index out of bounds; also fine+-- TODO:        (sunNest @_ @'[] @sh $ tlet (snest (knownShS @'[]) $ (sfromPrimal ix - sfloor (ssum0 v) - srepl 10001) `remH` srepl 2) $ \rr -> snest (knownShS @'[]) $ sqrt $ abs $ sindex v (ShapedList.singletonIndex (ifH (signum (sprimalPart (sunNest rr)) ==. negate (signum $ srepl 2)) (sprimalPart (sunNest rr) + srepl 2) (sprimalPart (sunNest rr)))))+        (sunNest @_ @'[] @sh $ tlet ((sfromPrimal ix - sfloor (ssum0 v) - srepl 10001) `remH` srepl 2) $ \rr -> snest (knownShS @'[]) $ sqrt $ abs $ sindex v ((kfromS $ ifH (signum (sprimalPart rr) ==. negate (signum $ srepl 2)) (sprimalPart rr + srepl 2) (sprimalPart rr)) :.$ ZIS))++testFooBuildNest21S :: Assertion+testFooBuildNest21S =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2] [0.2886751345948129,0.35355339059327373])+    (rev' @Double @1 (fooBuildNest2S @2 @'[] . sfromR) (ringestData [2] [3.0,2.0]))++testFooBuildNest25S :: Assertion+testFooBuildNest25S =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,2,1,2,2] [0.22360679774997896,0.35355339059327373,0.20412414523193154,0.5,-0.35355339059327373,500.0,1.5811388300841895,-1.118033988749895,0.1381447409988844,0.16666666666666666,0.17677669529663687,-0.25,8.574929257125441e-2,0.288948802391873,-8.703882797784893e-2,9.805806756909202e-2])+    (rev' @Double @5 (fooBuildNest2S @2 @[2, 1, 2, 2] . sfromR) t16)++fooBuild3 :: forall target r n.+             ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+          => target (TKR (1 + n) r) -> target (TKR (1 + n) r)+fooBuild3 v =+  rbuild1 22 $ \ix ->+    bar ( rreplicate0N (shrTail $ rshape v) (rscalar 1)+        , rindex v [minH 1 (ix + 1)] )  -- index not out of bounds++testFooBuild3 :: Assertion+testFooBuild3 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,2,1,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,423.72976235076516,-260.41676627885636,-17.60047532855961,151.18955028869385,-1059.9668424433578,-65.00898015327623,-21.49245448729951,743.7622427949768])+    (rev' @Double @5 fooBuild3 t16)++fooBuild5 :: forall target r n.+             ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+          => target (TKR (1 + n) r) -> target (TKR (1 + n) r)+fooBuild5 v =+  let r = rsum v+      v' = rreplicate0N (shrTail $ rshape v) $ rminimum $ rflatten v+  in rbuild1 2 $ \ix ->+       r * foo ( rreplicate0N (shrTail $ rshape v) (rscalar 3)+               , rrepl (rshape r) 5 * r+               , r * v')+       + bar (r, rindex v [minH 1 (ix + 1)])  -- index not out of bounds++testFooBuildDt :: Assertion+testFooBuildDt =+  assertEqualUpToEpsilon 1e-5+    (rconcrete $ Nested.rfromListPrimLinear [2,2,1,2,2] [1.1033568028244503e7,74274.22833989389,-5323238.2765011545,253074.03394016018,4.14744804041263e7,242643.98750578283,-1.922371592087736e7,2.730274503834733e7,1.135709425204681e7,6924.195066252549,-5345004.080027547,255679.51406100337,3.8870981856703006e7,241810.92121468345,-1.9380955730171032e7,2.877024321777493e7])+    (vjp @_ @(TKR 5 Double)+           fooBuild5 t16 (rreplicate0N [2, 2, 1, 2, 2] (rscalar 42)))++testFooBuildDt2 :: Assertion+testFooBuildDt2 =+  assertEqualUpToEpsilon 1e-5+    (rconcrete $ Nested.rfromListPrimLinear [2,2,1,2,2] [2.206713605648901e7,148548.45667978778,-1.0646476553002307e7,506148.0678803204,8.294896080825263e7,485287.9750115657,-3.844743184175473e7,5.460549007669466e7,2.271418850409362e7,13848.390132505112,-1.0690008160055092e7,511359.0281220066,7.774196371340603e7,483621.8424293669,-3.876191146034207e7,5.754048643554987e7])+    (vjp @_ @(TKProduct (TKR 5 Double) (TKR 5 Double))+           (\x -> let y = fooBuild5 x in tpair y y) t16 (let dt = rreplicate0N [2, 2, 1, 2, 2] (rscalar 42) in tpair dt dt))++testFooBuild5 :: Assertion+testFooBuild5 =+  assertEqualUpToEpsilon' 1e-5+    (ringestData [3,1,2,2,1,2,2] [-613291.6547530327,571164.2201603781,-1338602.6247083102,528876.2566682736,1699442.2143691683,2874891.369778316,-3456754.605470273,3239487.8744244366,554916.1344235454,-775449.1803684114,3072.200583200206,1165767.8436804386,-1.0686356667942494e7,-6606976.194539241,-6457671.748790982,4791868.42112978,-615556.7946425928,569660.3506343022,-1348678.1169100606,534886.9366492515,1696036.143341285,2883992.9672165257,-3456212.5353846983,3240296.690514803,629047.8398075115,-794389.5797803313,-1143.8025173051583,1177448.8083517442,-1.15145721735623e7,-6618648.839812404,-6462386.031613377,5358224.852822481,-613291.6547530327,571164.2201603781,-1338602.6247083102,528876.2566682736,1699442.2143691683,2874891.369778316,-3456754.605470273,3239487.8744244366,554916.1344235454,-775449.1803684114,3072.200583200206,1165767.8436804386,-1.0686356667942494e7,-6606976.194539241,-6457671.748790982,4791868.42112978])+    (rev' @Double @7 fooBuild5 t48)++fooBuild1 :: forall target r n.+             ( ADReady target, GoodScalar r, KnownNat n, RealFloatH (target (TKR n r)) )+          => target (TKR (1 + n) r) -> target (TKR (1 + n) r)+fooBuild1 v =+  let r = rsum v+      tk = rreplicate0N (shrTail $ rshape v)+      v' = tk $ rminimum $ rflatten v+  in rbuild1 3 $ \ix ->+       r * foo ( tk (rscalar 3)+               , tk (rscalar 5) * r+               , r * v')+       + bar (r, rindex v [minH 1 (ix + 1)])++testFooBuild1 :: Assertion+testFooBuild1 =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,2,1,2,2] [394056.00100873224,2652.651012139068,-190115.65273218407,9038.358355005721,1481231.4430045108,8665.8566966351,-686561.2828884773,975098.0370838332,405610.50900167174,247.29268093759174,-190893.00285812665,9131.411216464405,1388249.3520251075,8636.104329095837,-692176.9903632513,1027508.6863491047])+    (rev' @Double @5 fooBuild1 t16)++fooMap1 :: (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+        => IShR (1 + n) -> target (TKR 0 r) -> target (TKR (1 + n) r)+fooMap1 sh r =+  let v = fooBuild1 $ rreplicate0N sh (r * r)+  in rmap0N (\x -> x * r + rscalar 5) v++testFooMap :: Assertion+testFooMap =+  assertEqualUpToEpsilon' 1e-3+    (rscalar 2.7518227)+    (rev' @Float @1 (fooMap1 [130]) (rscalar 0.1))++-- Reduced test, because this takes forever with Ast but without vectorization.+testFooMap1 :: Assertion+testFooMap1 =+  assertEqualUpToEpsilon 1e-6+    (rscalar 3901.312463734578)+    (grad (kfromR @_ @Double . rsum0 @7 . fooMap1 [4, 3, 2, 3, 4, 5, 3]) (rscalar 0.1))++fooNoGo :: forall target r n.+           ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )+        => target (TKR (1 + n) r) -> target (TKR (1 + n) r)+fooNoGo v =+  let r = rsum v+      r0 = rsum0 v+      shTail = shrTail (rshape v)+  in rbuild1 3 (\ix ->+       bar ( rreplicate0N shTail (rscalar 3.14)+           , bar ( rrepl shTail 3.14+                 , rindex v [ix]) )+       + ifH (rindex v (ix * 2 :.: ZIR) <=. rreplicate0N shTail (rscalar 0) &&* 6 >. abs ix)+               r (rreplicate0N shTail (rscalar 5) * r))+     / rslice 1 3 (rmap0N (\x -> ifH (x >. r0) r0 x) v)+     * rbuild1 3 (const $ rrepl shTail 1)++testFooNoGo :: Assertion+testFooNoGo =+  assertEqualUpToEpsilon' 1e-6+   (ringestData [5] [344.3405885672822,-396.1811403813819,7.735358041386672,-0.8403418295960372,5.037878787878787])+   (rev' @Double @1 fooNoGo+         (ringestData [5] [1.1 :: Double, 2.2, 3.3, 4, 5]))++testFooNoGo10 :: Assertion+testFooNoGo10 =+  assertEqualUpToEpsilon 1e-10+    (ringestData [5, 3, 1, 2, 2, 1, 2, 2] [8.096867407436072e-8,9.973025492756426e-8,9.976696178938985e-8,5.614458707681111e-8,-1.8338500573636686e-7,-2.144970334428336e-7,7.354143606421902e-7,-1.8140041785503643e-7,8.096867407436072e-8,9.973025492756426e-8,9.976696178938985e-8,5.614458707681111e-8,-2.01381292700262e-7,-2.221588091014473e-7,7.354143606421902e-7,-1.9951065225263367e-7,1.7230532848112822e-7,4.5426218104870796e-7,1.430886696893587e-7,9.354993295163118e-7,-5.225515010723883e-7,1.019433073376504e-6,9.64067025472343e-6,-4.872227980305747e-6,8.089200625992941e-8,9.924319994964371e-8,1.092480101004153e-7,-2.8478802468285825e-7,9.641049518625974e-8,2.9624147815716037e-7,-1.950868158558337e-7,9.547754822865364e-8,4.5426218104870796e-7,4.5426218104870796e-7,4.5426218104870796e-7,4.5426218104870796e-7,-4.872227980305747e-6,-4.872227980305747e-6,-4.872227980305747e-6,-4.872227980305747e-6,9.361277121832246e-8,-4.872227980305747e-6,-4.872227980305747e-6,-4.872227980305747e-6,9.361277121832246e-8,9.361277121832246e-8,9.361277121832246e-8,9.361277121832246e-8,-5.488572216677945e-7,-1.8496203182958057e-7,-1.4603644180845103e-7,-1.2145268106051633e-7,-2.817402689957553e-7,-2.9913537180597976e-7,6.272804203945257e-7,-2.3697344464172694e-7,-5.488572216677945e-7,-1.8496203182958057e-7,-1.4603644180845103e-7,-1.2145268106051633e-7,-2.613973017956691e-7,-3.0013408634207794e-7,6.272804203945257e-7,-2.916736028401805e-7,-7.0114505846358575e-6,-4.303381366239431e-5,-4.897282418246382e-6,-1.710952247892854e-4,-4.2040039667393255e-5,-2.0204742564752248e-4,-1.7017980671040968e-2,-4.247008401789142e-3,-1.056090348050961e-6,-2.210187184450231e-6,-2.7842041329045203e-6,-1.0402806498987974e-5,-1.2967382896879757e-7,-1.9315601705070884e-5,-2.40087090725031e-7,-2.4419692405172046e-7,-4.303381366239431e-5,-4.303381366239431e-5,-4.303381366239431e-5,-4.303381366239431e-5,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-2.683138631810477e-7,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-2.683138631810477e-7,-2.683138631810477e-7,-2.683138631810477e-7,-2.683138631810477e-7,-5.488572216677945e-7,-1.8496203182958057e-7,-1.4603644180845103e-7,-1.2145268106051633e-7,-2.817402689957553e-7,-2.9913537180597976e-7,6.272804203945257e-7,-2.3697344464172694e-7,-5.488572216677945e-7,-1.8496203182958057e-7,-1.4603644180845103e-7,-1.2145268106051633e-7,-2.613973017956691e-7,-3.0013408634207794e-7,6.272804203945257e-7,-2.916736028401805e-7,-7.0114505846358575e-6,-4.303381366239431e-5,-4.897282418246382e-6,-1.710952247892854e-4,-4.2040039667393255e-5,-2.0204742564752248e-4,-1.7017980671040968e-2,-4.247008401789142e-3,-1.056090348050961e-6,-2.210187184450231e-6,-2.7842041329045203e-6,-1.0402806498987974e-5,-1.2967382896879757e-7,-1.9315601705070884e-5,-2.40087090725031e-7,-2.4419692405172046e-7,-4.303381366239431e-5,-4.303381366239431e-5,-4.303381366239431e-5,-4.303381366239431e-5,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-2.683138631810477e-7,-4.247008401789142e-3,-4.247008401789142e-3,-4.247008401789142e-3,-2.683138631810477e-7,-2.683138631810477e-7,-2.683138631810477e-7,-2.683138631810477e-7,-5.469529675653596e-7,-2.331458950045675e-7,-1.9907443163522408e-7,-1.4019078434680374e-7,-6.95091094132346e-8,-5.685763846730528e-8,-9.268594848659335e-8,-3.010367762029461e-8,-5.469529675653596e-7,-2.331458950045675e-7,-1.9907443163522408e-7,-1.4019078434680374e-7,-3.415394012988984e-8,-5.069973314807702e-8,-9.268594848659335e-8,-6.380451815099858e-8,-6.883755913116986e-6,-4.273807584344302e-5,-4.79037108793574e-6,-1.705307241188017e-4,-4.2267488166320864e-5,-2.0143642393829028e-4,-1.701262134129569e-2,-4.2496361738088365e-3,-1.0224785375169973e-6,-2.1427637177332083e-6,-2.705952143004936e-6,-1.0493018474305117e-5,-1.819666770962338e-7,-1.911089472080586e-5,-9.045482032374276e-8,-2.819821645880664e-7,-4.273807584344302e-5,-4.273807584344302e-5,-4.273807584344302e-5,-4.273807584344302e-5,-4.2496361738088365e-3,-4.2496361738088365e-3,-4.2496361738088365e-3,-4.2496361738088365e-3,-3.019273543907303e-7,-4.2496361738088365e-3,-4.2496361738088365e-3,-4.2496361738088365e-3,-3.019273543907303e-7,-3.019273543907303e-7,-3.019273543907303e-7,-3.019273543907303e-7,8.287292817679557e-8,5.154639175257732e-8,4.672897196261682e-8,3.740648379052369e-8,2.884615384615385e-8,2.7780699895840894e-8,1.5447991761071065e-8,2.546934916639589e-8,8.287292817679557e-8,5.154639175257732e-8,4.672897196261682e-8,3.740648379052369e-8,2.5862068965517245e-8,2.7275544092553562e-8,1.5447991761071065e-8,2.8358432436548274e-8,3.0000000000000004e-7,7.500000000000001e-7,2.5000000000000004e-7,1.5000000000000002e-6,-7.500000000000001e-7,1.6304347826086957e-6,1.5000000000000002e-5,-7.500000000000001e-6,1.1450381679389314e-7,1.6666666666666668e-7,1.8750000000000003e-7,-3.7500000000000006e-7,4.411764705882353e-8,5.00948462422186e-7,-4.545454545454546e-8,5.76923076923077e-8,7.500000000000001e-7,7.500000000000001e-7,7.500000000000001e-7,7.500000000000001e-7,-7.500000000000001e-6,-7.500000000000001e-6,-7.500000000000001e-6,-7.500000000000001e-6,5.999928000863991e-8,-7.500000000000001e-6,-7.500000000000001e-6,-7.500000000000001e-6,5.999928000863991e-8,5.999928000863991e-8,5.999928000863991e-8,5.999928000863991e-8])+   (grad (kfromR @_ @Double . rsum0 @8 . rmap0N (* rscalar 0.000000001) . fooNoGo) (rmap0N (* rscalar 0.01) $ rreplicate 5 t48))++nestedBuildMap :: forall target n r.+                  (ADReady target, GoodScalar r, n <= 6, KnownNat n, Differentiable r)+               => target (TKR 0 r) -> target (TKR (1 + n) r)+nestedBuildMap r =+  let w x = rreplicate0N [4] x :: target (TKR 1 r)+      v' = rreplicate0N (177 :$: ZSR) r+      nestedMap x = rmap0N (x /) (w x)+      variableLengthBuild iy = rbuild1 7 (\ix ->+        rindex v' (ix + iy :.: ZIR))+      doublyBuild =+        rbuild1 3 (rreplicate0N (shrTake @n @(6 - n)+                             $ 2 :$: 4 :$: 2 :$: 1 :$: 3 :$: 2 :$: ZSR)+                   . rminimum . variableLengthBuild)+  in rmap0N (\x -> x * rsum0+                         (rbuild1 3 (\ix -> bar (x, rindex v' [ix]))+                          + fooBuild1 (nestedMap x)+                          / fooMap1 [3] x)+            ) doublyBuild++testNestedBuildMap1 :: Assertion+testNestedBuildMap1 =+  assertEqualUpToEpsilon' 1e-8+    (rscalar 22.673212907588812)+    (rev' @Double @1 nestedBuildMap (rscalar 0.6))++testNestedBuildMap10 :: Assertion+testNestedBuildMap10 =+  assertEqualUpToEpsilon 1e-8+    (map rscalar [109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257, 109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257, 109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257])+    (map (cgrad (kfromR . rsum0 @1 @(TKScalar Double) . nestedBuildMap))+         (map (Concrete . Nested.rscalar) $ [0.1, 0.2 .. 1] ++ [0.1, 0.2 .. 1] ++ [0.1, 0.2 .. 1]))++testNestedBuildMap11 :: Assertion+testNestedBuildMap11 =+  assertEqualUpToEpsilon 1e-8+    (map rscalar [109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257, 109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257, 109.62086996459126,106.70290239773645,103.05843225947055,98.11825678264942,67.8014491889543,22.67321290758882,-163.40832575807545,376.4240286600336,-1996.9068313949347,249.28292226561257])+    (map (grad (kfromR . rsum0 @1 @(TKScalar Double) . nestedBuildMap))+         (map (Concrete . Nested.rscalar) $ [0.1, 0.2 .. 1] ++ [0.1, 0.2 .. 1] ++ [0.1, 0.2 .. 1]))++{-+testNestedBuildMap7 :: Assertion+testNestedBuildMap7 =+  assertEqualUpToEpsilon' 1e-8+    (rscalar 2176.628439128524)+    (rev' @Double @7 nestedBuildMap (rscalar 0.6))+-}++-- The n <= 4 is necessary despite what GHC claims. Applying @(2 + n)+-- to nestedBuildMap doesn't help.+nestedSumBuild+  :: forall target n r.+     (ADReady target, GoodScalar r, n <= 4, KnownNat n, Differentiable r)+  => target (TKR n r) -> target (TKR (2 + n) r)+nestedSumBuild v =+  rbuild1 13 $ \ix1 -> rbuild1 4 $ \ix2 ->+    ifH (ix2 >. ix1)+        (rmap0N ((* rscalar (-0.00000003)) . sqrt . abs)+         $ nestedBuildMap (rsum0 v)+           `rindex` (ix2 `remH` 3 :.: minH 1 ix1 :.: minH ix1 3 :.: ZIR))+        (nestedBuildMap (rscalar 0.00042)+         `rindex` (ix2 `remH` 3 :.: minH 1 ix1 :.: minH ix1 3 :.: ZIR))++testNestedSumBuild1 :: Assertion+testNestedSumBuild1 =+  assertEqualUpToEpsilon 1e-6+    (ringestData [5] [5.738943380972744e-6,5.738943380972744e-6,5.738943380972744e-6,5.738943380972744e-6,5.738943380972744e-6])+    (grad (kfromR . rsum0 @3 @(TKScalar Double) . nestedSumBuild) (ringestData [5] [1.1, 2.2, 3.3, 4, -5.22]))++{-+testNestedSumBuild5 :: Assertion+testNestedSumBuild5 =+  assertEqualUpToEpsilon' 1e-6+    (ringestData [1,2,2] [3.5330436757054903e-3,3.5330436757054903e-3,3.5330436757054903e-3,3.5330436757054903e-3])+    (rev' @Double @5 nestedSumBuild (rsum (rsum t16)))+-}++nestedSumBuildB :: forall target n r. (ADReady target, GoodScalar r, KnownNat n)+                => target (TKR (1 + n) r) -> target (TKR 3 r)+nestedSumBuildB v =+  rbuild @2 [13, 4, 2] $ \case+    [ix, ix2] ->+      flip rindex [ix2]+        (rfromList+             [ rbuild1 2 rfromIndex0+             , rsum $ rbuild [9, 2] $ const $ rfromIndex0 ix+             , rindex v (fromList+                         $ replicate (rlength v - 1)+                             (maxH 0 $ minH 1 $ ix2 `quotH` 2 + ix `quotH` 4 - 1))+             , rbuild1 2 (\_ -> rsum0 v)+             , rsum (rbuild1 7 (\ix7 ->+                 rreplicate 2 (rfromIndex0 ix7)))+             ])+    _ -> error "nestedSumBuildB: impossible pattern needlessly required"++testNestedSumBuildB :: Assertion+testNestedSumBuildB =+  assertEqualUpToEpsilon' 1e-8+    (ringestData [2,3,2,2,2] [30.0,30.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,35.0,35.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0,26.0])+    (rev' @Double @3 nestedSumBuildB (rsum $ rsum $ rtranspose [1, 4, 2, 0, 3] t48))++nestedBuildIndex :: forall target r. (ADReady target, GoodScalar r)+                 => target (TKR 5 r) -> target (TKR 3 r)+nestedBuildIndex v =+  rbuild1 2 $ \ix2 -> rindex (rbuild1 3 $ \ix3 -> rindex (rbuild1 3 $ \ix4 -> rindex v (ix4 `remH` 2 :.: ix2 :.: 0 :.: ZIR)) [ix3]) (ix2 :.: ZIR)++testNestedBuildIndex :: Assertion+testNestedBuildIndex =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,2,1,2,2] [1.0,1.0,1.0,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0])+    (rev' @Double @3 nestedBuildIndex t16)++barRelu+  :: ( ADReady target, GoodScalar r, KnownNat n, Differentiable r )+  => target (TKR n r) -> target (TKR n r)+barRelu x = let t = rreplicate0N (rshape x) (rscalar 0.001) * x+            in relu $ bar (t, relu t)++testBarReluADValDt :: Assertion+testBarReluADValDt =+  assertEqualUpToEpsilon 1e-6+    (rconcrete $ Nested.rfromListPrimLinear [2,2,1,2,2] [1.2916050471365906e-2,1.2469757606504572e-2,1.3064120086501589e-2,1.2320300700062944e-2,0.0,1.217049789428711e-2,1.2185494267265312e-2,0.0,1.4105363649830907e-2,1.3506236503127638e-2,1.3359213691150671e-2,0.0,1.7066665416485535e-2,1.2618022646204737e-2,0.0,1.595161947206668e-2])+    (vjp @_ @(TKR 5 Double)+           barRelu t16 (rreplicate0N [ 2 , 2 , 1 , 2 , 2 ] (rscalar 42.2)))++testBarReluADValDt2 :: Assertion+testBarReluADValDt2 =+  assertEqualUpToEpsilon 1e-6+    (rconcrete $ Nested.rfromListPrimLinear [2,2,1,2,2] [84.42583210117625,84.42493951543845,84.4261282404092,84.42464060162287,84.4,84.42434099465609,84.4243709887547,84.4,84.42821072755468,84.42701247325044,84.42671842762383,84.4,84.43413333114152,84.42523604552053,84.4,84.43190323923253])+    (vjp @_ @(TKProduct+                  (TKR 4 Double)+                  (TKProduct (TKR 5 Float)+                             (TKS [2,2,1,2,2] Double)))+           (\x -> tpair (rsum x) (tpair (rcast $ barRelu x) (sfromR $ barRelu x))) t16+                  (let dt = rreplicate0N [ 2 , 2 , 1 , 2 , 2 ] (rscalar 42.2)+                   in tpair (rsum dt) (tpair (rcast dt) (sfromR dt))))++testBarReluADVal :: Assertion+testBarReluADVal =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3,1,2,2,1,2,2] [3.513740871835189e-4,3.8830416352632824e-4,3.981974371104471e-4,4.2420226755643853e-4,4.6186212581292275e-4,4.6805323209889415e-4,5.933633926875981e-4,4.8311739820100107e-4,3.513740871835189e-4,3.8830416352632824e-4,3.981974371104471e-4,4.2420226755643853e-4,4.803836032226148e-4,4.7114455958615145e-4,5.933633926875981e-4,4.6464270870595213e-4,3.060675467148428e-4,2.954918864100193e-4,3.095763053673437e-4,2.9195025355591045e-4,0.0,2.9166656928452994e-4,2.887557883241243e-4,0.0,3.342503234557057e-4,3.2005299770444394e-4,3.165690448140097e-4,0.0,4.0442335110155446e-4,2.990052759764126e-4,0.0,3.780004614233832e-4,2.954918864100193e-4,2.954918864100193e-4,2.954918864100193e-4,2.954918864100193e-4,0.0,0.0,0.0,0.0,3.7466025157760897e-4,0.0,0.0,0.0,3.7466025157760897e-4,3.7466025157760897e-4,3.7466025157760897e-4,3.7466025157760897e-4])+    (rev' @Double @7 barRelu t48)++testBarReluADVal3 :: Assertion+testBarReluADVal3 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3,1,2,2,1,2,2] [2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.885852309100301e-4,2.885923176600045e-4,2.887454843457817e-4,2.886097295122454e-4,2.8846476339094805e-4,2.885038541771792e-4,2.885145151321922e-4,2.8854294397024206e-4,2.8860655161315664e-4,2.88595871110374e-4,2.887454843457817e-4,2.885884088500461e-4,2.884182085399516e-4,2.884075468755327e-4,2.8842176240868867e-4,2.8840399312321096e-4,0.0,2.8840370860416445e-4,2.884007943794131e-4,0.0,2.884469945274759e-4,2.8843242392031246e-4,2.884288700806792e-4,0.0,2.885212670262263e-4,2.884110805753153e-4,0.0,2.8849283778617973e-4,2.884075468755327e-4,2.884075468755327e-4,2.884075468755327e-4,2.884075468755327e-4,0.0,0.0,0.0,0.0,2.884892851579934e-4,0.0,0.0,0.0,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4,2.884892851579934e-4])+    (rev' @Double @7 barRelu+         (rmap0N (* rscalar 0.001) t48))++braidedBuilds :: forall target n r. (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+              => target (TKR (1 + n) r) -> target (TKR 2 r)+braidedBuilds r =+  rbuild1 3 (\ix1 ->+    rbuild1 4 (\ix2 -> rindex (rfromList+      [rfromIndex0 ix2, rscalar 7, rsum0 (rslice 1 1 r), rscalar (-0.2)]) (ix1 :.: ZIR)))++testBraidedBuilds :: Assertion+testBraidedBuilds =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [4] [0.0,4.0,0.0,0.0])+    (rev' @Double @2 (braidedBuilds @_ @0) (rreplicate0N [4] (rscalar 3.4)))++testBraidedBuilds1 :: Assertion+testBraidedBuilds1 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [2,2,1,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0])+    (rev' @Double @2 braidedBuilds t16)++recycled :: (ADReady target, GoodScalar r, KnownNat n)+         => target (TKR n r) -> target (TKR 7 r)+recycled r =+  rbuild1 2 $ \_ -> rbuild1 4 $ \_ -> rbuild1 2 $ \_ -> rbuild1 3 $ \_ ->+    nestedSumBuildB (rreplicate 4 r)++testRecycled :: Assertion+testRecycled =+  assertEqualUpToEpsilon' 1e-6+    (rrepl [2] 5616)+    (rev' @Double @7 (recycled @_ @_ @1) (rreplicate0N [2] (rscalar 1.0001)))++{-+testRecycled1 :: Assertion+testRecycled1 =+  assertEqualUpToEpsilon' 1e-6+    (ringestData [5, 4, 2] [5184.0,5184.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,5424.0,5424.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0,4992.0])+    (rev' @Double @7 (recycled @_ @_ @3) (rreplicate0N [5, 4, 2] (rscalar 0.0002)))+-}++concatBuild :: forall target r n.+               (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+            => target (TKR (1 + n) r) -> target (TKR (3 + n) r)+concatBuild r =+  rbuild1 7 (\i ->+    rconcat [ rbuild1 5 (const r)+            , rbuild1 1 (\j -> rmap0N (* rfromIndex0 (j - i)) r)+            , rbuild1 11 (\j ->+                rmap0N (* (rfromIndex0+                  (kfromR (rprimalPart @target (rscalar 125)) * (j `remH` (abs (signum i + abs i) + 1))+                   + maxH j (i `quotH` (j + 1)) * (kfromR . rprimalPart . rfloor) (rsum0 r)+                   - ifH (r <=. r &&* i <. j)+                         (kfromR $ rprimalPart $ rminIndex (rflatten r))+                         ((kfromR . rprimalPart . rfloor) $ rsum0 $ r ! ((i * j) `remH` 7 :.: ZIR))))) r)+            , rbuild1 13 (\_k ->+                rsum $ rtr $ rreplicate (rwidth r) (rslice 0 1 r)) ])++testConcatBuild0 :: Assertion+testConcatBuild0 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7] [16917.0,16280.0,16280.0,16280.0,16280.0,16280.0,16280.0])+    (rev' @Double @3 concatBuild+       (ringestData [7] [0.651,0.14,0.3414,-0.14,0.0014,0.0020014,0.9999]))++testConcatBuild1 :: Assertion+testConcatBuild1 =+  assertEqualUpToEpsilon 1e-10+    (ringestData [3,1,2,2,1,2,2] [1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4816999999999999e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3,1.4544e-3])+    (grad (kfromR . rsum0 @9 @(TKScalar Double) . concatBuild . rmap0N (* rscalar 1e-7)) t48)++concatBuildm :: forall target r n.+                (ADReady target, GoodScalar r, KnownNat n, Differentiable r)+             => target (TKR (1 + n) r) -> target (TKR (2 + n) r)+concatBuildm r =+  rbuild1 7 (\i ->+    rmap0N (* (rfromIndex0+      ((kfromR . rprimalPart . rfloor) $ rsum0 $ r ! (i :.: ZIR)))) r)++testConcatBuild0m :: Assertion+testConcatBuild0m =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [7] [-1.0,-1.0,-1.0,-1.0,-1.0,-1.0,-1.0])+    (rev' @Double @2 concatBuildm+       (ringestData [7] [0.651,0.14,0.3414,-0.14,0.0014,0.0020014,0.9999]))++testConcatBuild1m :: Assertion+testConcatBuild1m =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3,1,2,2,1,2,2] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])+    (rev' @Double @8 (concatBuildm . rmap0N (* rscalar 1e-7)) t48)++concatBuild2 :: (ADReady target, GoodScalar r, KnownNat n)+             => target (TKR (1 + n) r) -> target (TKR (3 + n) r)+concatBuild2 r =+  rbuild1 5 (\i ->+    rbuild1 2 (\j -> rmap0N (* rfromIndex0 (maxH j (i `quotH` (j + 1)))) r))++testConcatBuild2 :: Assertion+testConcatBuild2 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3] [16.0,16.0,16.0])+    (rev' @Double @3 concatBuild2 (ringestData [3] [0.651,0.14,0.3414]))++testConcatBuild22 :: Assertion+testConcatBuild22 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3,1,2,2,1,2,2] [16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0,16.0])+    (rev' @Double @9 concatBuild2 t48)++concatBuild3 :: (ADReady target, GoodScalar r)+             => target (TKR 1 r) -> target (TKR 2 r)+concatBuild3 _r =+  rbuild1 5 (\i ->+    rbuild1 2 (\j -> rfromIndex0 (maxH j (i `quotH` (j + 1)))))++testConcatBuild3 :: Assertion+testConcatBuild3 =+  assertEqualUpToEpsilon' 1e-10+    (ringestData [0] [])+    (rev' @Double @2 concatBuild3 (ringestData [0] []))++testLogistic0 :: Assertion+testLogistic0 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 4.5176659730912e-2)+    (rev' @Double @0 logistic (rscalar 3))++testLogistic5 :: Assertion+testLogistic5 =+  assertEqualUpToEpsilon' 1e-10+    (rfromListLinear [2,2,1,2,2] [6.648056670790033e-3,0.10499358540350662,2.466509291359931e-3,0.19661193324148185,0.1049935854035065,0.2499999999999375,0.24937604019289197,0.24751657271185995,2.0452222584760427e-6,1.2337934976493025e-4,3.3523767075636815e-4,1.7662706213291118e-2,1.7763568394002473e-15,4.540945566439111e-2,4.6588861451033536e-15,5.109024314693943e-12])+    (rev' @Double @5 logistic t16)++testLogistic52 :: Assertion+testLogistic52 =+  assertEqualUpToEpsilon' 1e-10+    (rfromListLinear [2,2,1,2,2] [1.3111246391159124e-3,2.1750075272657612e-2,4.8549901151740267e-4,4.312916016242333e-2,2.6155373652699744e-2,5.8750924453083504e-2,5.8238333583278255e-2,5.8847120842749026e-2,4.0211548220008027e-7,2.425923569564766e-5,6.592194028602285e-5,4.415319450597324e-3,3.4925295232121135e-16,9.12284655835676e-3,1.1647215362758384e-15,1.0044951474920845e-12])+    (rev' @Double @5 (logistic . logistic) t16)++logisticOld :: forall target r n.+            ( BaseTensor target, LetTensor target+            , BaseTensor (PrimalOf target), KnownNat n, GoodScalar r+            , Floating (PrimalOf target (TKR n r)) )+         => target (TKR n r) -> target (TKR n r)+logisticOld d0 = tlet d0 $ \d ->  -- used in rprimalPart and in tdualPart+  let sh = rshape d+      y0 = recip (rrepl sh 1 + exp (- rprimalPart @target d))+  in tlet (rfromPrimal @target y0)+     $ \y1 -> let y = rprimalPart @target y1+              in tD knownSTK y (rScale @target (y * (rrepl sh 1 - y))+                                $ rdualPart @target d)++testLogistic0Old :: Assertion+testLogistic0Old =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 4.5176659730912e-2)+    (rev' @Double @0 logisticOld (rscalar 3))++testLogistic5Old :: Assertion+testLogistic5Old =+  assertEqualUpToEpsilon' 1e-10+    (rfromListLinear [2,2,1,2,2] [6.648056670790033e-3,0.10499358540350662,2.466509291359931e-3,0.19661193324148185,0.1049935854035065,0.2499999999999375,0.24937604019289197,0.24751657271185995,2.0452222584760427e-6,1.2337934976493025e-4,3.3523767075636815e-4,1.7662706213291118e-2,1.7763568394002473e-15,4.540945566439111e-2,4.6588861451033536e-15,5.109024314693943e-12])+    (rev' @Double @5 logisticOld t16)++testLogistic52Old :: Assertion+testLogistic52Old =+  assertEqualUpToEpsilon' 1e-10+    (rfromListLinear [2,2,1,2,2] [1.3111246391159124e-3,2.1750075272657612e-2,4.8549901151740267e-4,4.312916016242333e-2,2.6155373652699744e-2,5.8750924453083504e-2,5.8238333583278255e-2,5.8847120842749026e-2,4.0211548220008027e-7,2.425923569564766e-5,6.592194028602285e-5,4.415319450597324e-3,3.4925295232121135e-16,9.12284655835676e-3,1.1647215362758384e-15,1.0044951474920845e-12])+    (rev' @Double @5 (logisticOld . logistic) t16)++logisticA :: forall target r n.+            ( BaseTensor target, LetTensor target+            , BaseTensor (PrimalOf target), KnownNat n, GoodScalar r+            , Floating (PrimalOf target (TKR n r)) )+         => target (TKR n r) -> target (TKR n r)+logisticA d0 = tlet d0 $ \d ->  -- used in rprimalPart and in tdualPart+  let sh = rshape d+      y0 = recip (rrepl sh 1 + exp (- rprimalPart @target d))+  in tlet (rfromPrimal @target y0)+     $ \y1 -> let y = rprimalPart @target y1+              in rfromPrimal y+                 + rfromDual (rScale @target (y * (rrepl sh 1 - y))+                              $ rdualPart @target d)++testLogisticA0 :: Assertion+testLogisticA0 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 4.5176659730912e-2)+    (rev' @Double @0 logisticA (rscalar 3))++logisticB :: forall target r n.+            ( BaseTensor target, LetTensor target+            , BaseTensor (PrimalOf target), KnownNat n, GoodScalar r+            , Floating (PrimalOf target (TKR n r)) )+         => target (TKR n r) -> target (TKR n r)+logisticB d0 = tlet d0 $ \d ->  -- used in rprimalPart and in tdualPart+  let sh = rshape d+      y0 = recip (rrepl sh 1 + exp (- rprimalPart @target d))+  in tlet (rfromPrimal @target y0)+     $ \y1 -> let y = rprimalPart @target y1+              in rfromPrimal y + rfromDual (rdualPart @target d)++testLogisticB0 :: Assertion+testLogisticB0 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1)+    (rev' @Double @0 logisticB (rscalar 3))++logisticC :: forall target r n.+            ( BaseTensor target, LetTensor target+            , KnownNat n, GoodScalar r )+         => target (TKR n r) -> target (TKR n r)+logisticC d0 = tlet d0 $ \d ->  -- used in rprimalPart and in tdualPart+  let y0 = rprimalPart @target d+  in rfromPrimal @target y0++testLogisticC0 :: Assertion+testLogisticC0 =+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0)+    (rev' @Double @0 logisticC (rscalar 3))
+ test/simplified/TestMnistCNNR.hs view
@@ -0,0 +1,406 @@+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+-- | Tests of "MnistCnnRanked2" convolutional neural network+-- using a few different optimization pipelines.+--+-- With the current CPU backend it's slow enough that it's hard to see+-- if it trains.+module TestMnistCNNR+  ( testTrees+  ) where++import Prelude++import Control.Monad (foldM, unless)+import System.IO (hPutStrLn, stderr)+import System.Random+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)+import Text.Printf++import Data.Array.Nested.Ranked.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstInterpret++import EqEpsilon++import MnistCnnRanked2 qualified+import MnistData++-- TODO: optimize enough that it can run for one full epoch in reasonable time+-- and then verify it trains down to ~20% validation error in a short enough+-- time to include such a training run in tests.++testTrees :: [TestTree]+testTrees = [ tensorADValMnistTestsCNNRA+            , tensorADValMnistTestsCNNRI+            , tensorADValMnistTestsCNNRO+            ]++type XParams r = X (MnistCnnRanked2.ADCnnMnistParameters Concrete r)++-- POPL differentiation, straight via the ADVal instance of RankedTensor,+-- which side-steps vectorization.+mnistTestCaseCNNRA+  :: forall r.+     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+  => String+  -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> r+  -> TestTree+mnistTestCaseCNNRA prefix epochs maxBatches khInt kwInt c_outInt n_hiddenInt+                   miniBatchSize totalBatchSize expected =+  withSNat khInt $ \(_khSNat :: SNat kh) ->+  withSNat kwInt $ \(_kwSNat :: SNat kw) ->+  withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->+  withSNat n_hiddenInt $ \(_n_hiddenSNat :: SNat n_hidden) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue+            @(Concrete (X (MnistCnnRanked2.ADCnnMnistParametersShaped+                             Concrete SizeMnistHeight SizeMnistWidth+                             kh kw c_out n_hidden r)))+            0.4 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show khInt, show kwInt+                        , show c_outInt, show n_hiddenInt+                        , show miniBatchSize+                        , show $ widthSTK $ knownSTK @(XParams r)+                        , show (tsize knownSTK targetInit) ]+      ftest :: Int -> MnistDataBatchR r -> Concrete (XParams r) -> r+      ftest batch_size mnistData pars =+        MnistCnnRanked2.convMnistTestR+          batch_size mnistData (fromTarget @Concrete pars)+  in testCase name $ do+       hPutStrLn stderr $+         printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+                prefix epochs maxBatches+       trainData <- map mkMnistDataR+                    <$> loadMnistData trainGlyphsPath trainLabelsPath+       testData <- map mkMnistDataR . take (totalBatchSize * maxBatches)+                   <$> loadMnistData testGlyphsPath testLabelsPath+       let testDataR = mkMnistDataBatchR testData+           f :: MnistDataBatchR r -> ADVal Concrete (XParams r)+             -> ADVal Concrete (TKScalar r)+           f (glyphR, labelR) adinputs =+             MnistCnnRanked2.convMnistLossFusedR+               miniBatchSize (rconcrete glyphR, rconcrete labelR)+               (fromTarget adinputs)+           runBatch :: (Concrete (XParams r), StateAdam (XParams r))+                    -> (Int, [MnistDataR r])+                    -> IO (Concrete (XParams r), StateAdam (XParams r))+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkR = map mkMnistDataBatchR+                          $ filter (\ch -> length ch == miniBatchSize)+                          $ chunksOf miniBatchSize chunk+                 res@(parameters2, _) =+                   sgdAdam f chunkR parameters stateAdam+                 trainScore =+                   ftest (length chunk) (mkMnistDataBatchR chunk) parameters2+                 testScore =+                   ftest (totalBatchSize * maxBatches) testDataR parameters2+                 lenChunk = length chunk+             unless (n_hiddenInt < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> (Concrete (XParams r), StateAdam (XParams r))+                    -> IO (Concrete (XParams r))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (n_hiddenInt < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+           ftk = tftk @Concrete (knownSTK @(XParams r)) targetInit+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal =+             1 - ftest (totalBatchSize * maxBatches) testDataR res+       testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCaseCNNRA+  :: String+  -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Double+  -> TestTree #-}++tensorADValMnistTestsCNNRA :: TestTree+tensorADValMnistTestsCNNRA = testGroup "CNNR ADVal MNIST tests"+  [ mnistTestCaseCNNRA "CNNRA 1 epoch, 1 batch" 1 1 4 4 8 16 1 1+                       (1 :: Double)+  , mnistTestCaseCNNRA "CNNRA artificial 1 2 3 4 5" 1 1 2 3 4 5 1 10+                       (1 :: Float)+  , mnistTestCaseCNNRA "CNNRA artificial 5 4 3 2 1" 5 4 3 2 1 1 1 1+                       (1 :: Double)+  , mnistTestCaseCNNRA "CNNRA 1 epoch, 0 batch" 1 0 4 4 16 64 16 50+                       (1.0 :: Float)+  ]++-- POPL differentiation, with Ast term defined and vectorized only once,+-- but differentiated anew in each gradient descent iteration.+mnistTestCaseCNNRI+  :: forall r.+     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+  => String+  -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> r+  -> TestTree+mnistTestCaseCNNRI prefix epochs maxBatches khInt kwInt c_outInt n_hiddenInt+                   miniBatchSize totalBatchSize expected =+  withSNat khInt $ \(_khSNat :: SNat kh) ->+  withSNat kwInt $ \(_kwSNat :: SNat kw) ->+  withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->+  withSNat n_hiddenInt $ \(_n_hiddenSNat :: SNat n_hidden) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue+            @(Concrete (X (MnistCnnRanked2.ADCnnMnistParametersShaped+                             Concrete SizeMnistHeight SizeMnistWidth+                             kh kw c_out n_hidden r)))+            0.4 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show khInt, show kwInt+                        , show c_outInt, show n_hiddenInt+                        , show miniBatchSize+                        , show $ widthSTK $ knownSTK @(XParams r)+                        , show (tsize knownSTK targetInit) ]+      ftest :: Int -> MnistDataBatchR r -> Concrete (XParams r) -> r+      ftest batch_size mnistData pars =+        MnistCnnRanked2.convMnistTestR+          batch_size mnistData (fromTarget @Concrete pars)+  in testCase name $ do+       hPutStrLn stderr $+         printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+                prefix epochs maxBatches+       trainData <- map mkMnistDataR+                    <$> loadMnistData trainGlyphsPath trainLabelsPath+       testData <- map mkMnistDataR . take (totalBatchSize * maxBatches)+                   <$> loadMnistData testGlyphsPath testLabelsPath+       let testDataR = mkMnistDataBatchR testData+           ftk = tftk @Concrete (knownSTK @(XParams r)) targetInit+       (_, _, var, varAst2) <- funToAstRevIO ftk+       (varGlyph, astGlyph) <-+         funToAstIO (FTKR (miniBatchSize+                           :$: sizeMnistHeightInt+                           :$: sizeMnistWidthInt+                           :$: ZSR) FTKScalar) id+       (varLabel, astLabel) <-+         funToAstIO (FTKR (miniBatchSize+                           :$: sizeMnistLabelInt+                           :$: ZSR) FTKScalar) id+       let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)+           ast = simplifyInline+                 $ MnistCnnRanked2.convMnistLossFusedR+                     miniBatchSize (astGlyph, astLabel)+                     (fromTarget varAst2)+           f :: MnistDataBatchR r -> ADVal Concrete (XParams r)+             -> ADVal Concrete (TKScalar r)+           f (glyph, label) varInputs =+             let env = extendEnv var varInputs emptyEnv+                 envMnist = extendEnv varGlyph (rconcrete glyph)+                            $ extendEnv varLabel (rconcrete label) env+             in interpretAstFull envMnist ast+           runBatch :: (Concrete (XParams r), StateAdam (XParams r))+                    -> (Int, [MnistDataR r])+                    -> IO (Concrete (XParams r), StateAdam (XParams r))+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkR = map mkMnistDataBatchR+                          $ filter (\ch -> length ch == miniBatchSize)+                          $ chunksOf miniBatchSize chunk+                 res@(parameters2, _) =+                   sgdAdam f chunkR parameters stateAdam+                 !trainScore =+                   ftest (length chunk) (mkMnistDataBatchR chunk) parameters2+                 !testScore =+                   ftest (totalBatchSize * maxBatches) testDataR parameters2+                 !lenChunk = length chunk+             unless (n_hiddenInt < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> (Concrete (XParams r), StateAdam (XParams r))+                    -> IO (Concrete (XParams r))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (n_hiddenInt < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal =+             1 - ftest (totalBatchSize * maxBatches) testDataR res+       testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCaseCNNRI+  :: String+  -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Double+  -> TestTree #-}++tensorADValMnistTestsCNNRI :: TestTree+tensorADValMnistTestsCNNRI = testGroup "CNNR Intermediate MNIST tests"+  [ mnistTestCaseCNNRI "CNNRI 1 epoch, 1 batch" 1 1 4 4 8 16 1 1+                       (1 :: Double)+  , mnistTestCaseCNNRI "CNNRI artificial 1 2 3 4 5" 1 1 2 3 4 5 1 10+                       (1 :: Float)+  , mnistTestCaseCNNRI "CNNRI artificial 5 4 3 2 1" 5 4 3 2 1 1 1 1+                       (1 :: Double)+  , mnistTestCaseCNNRI "CNNRI 1 epoch, 0 batch" 1 0 4 4 16 64 16 50+                       (1.0 :: Float)+  ]++-- JAX differentiation, Ast term built and differentiated only once+-- and the result interpreted with different inputs in each gradient+-- descent iteration.+mnistTestCaseCNNRO+  :: forall r.+     ( Differentiable r, GoodScalar r+     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+  => String+  -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> r+  -> TestTree+mnistTestCaseCNNRO prefix epochs maxBatches khInt kwInt c_outInt n_hiddenInt+                   miniBatchSize totalBatchSize expected =+  withSNat khInt $ \(_khSNat :: SNat kh) ->+  withSNat kwInt $ \(_kwSNat :: SNat kw) ->+  withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->+  withSNat n_hiddenInt $ \(_n_hiddenSNat :: SNat n_hidden) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue+            @(Concrete (X (MnistCnnRanked2.ADCnnMnistParametersShaped+                             Concrete SizeMnistHeight SizeMnistWidth+                             kh kw c_out n_hidden r)))+            0.4 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show khInt, show kwInt+                        , show c_outInt, show n_hiddenInt+                        , show miniBatchSize+                        , show $ widthSTK $ knownSTK @(XParams r)+                        , show (tsize knownSTK targetInit) ]+      ftest :: Int -> MnistDataBatchR r -> Concrete (XParams r) -> r+      ftest batch_size mnistData pars =+        MnistCnnRanked2.convMnistTestR+          batch_size mnistData (fromTarget @Concrete pars)+  in testCase name $ do+       hPutStrLn stderr $+         printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+                prefix epochs maxBatches+       trainData <- map mkMnistDataR+                    <$> loadMnistData trainGlyphsPath trainLabelsPath+       testData <- map mkMnistDataR . take (totalBatchSize * maxBatches)+                   <$> loadMnistData testGlyphsPath testLabelsPath+       let testDataR = mkMnistDataBatchR testData+           dataInit = case chunksOf miniBatchSize testData of+             d : _ -> let (dglyph, dlabel) = mkMnistDataBatchR d+                      in (rconcrete dglyph, rconcrete dlabel)+             [] -> error "empty test data"+           f :: ( MnistCnnRanked2.ADCnnMnistParameters+                    (AstTensor AstMethodLet FullSpan) r+                , ( AstTensor AstMethodLet FullSpan (TKR 3 r)+                  , AstTensor AstMethodLet FullSpan (TKR 2 r) ) )+             -> AstTensor AstMethodLet FullSpan (TKScalar r)+           f = \ (pars, (glyphR, labelR)) ->+             MnistCnnRanked2.convMnistLossFusedR+               miniBatchSize (rprimalPart glyphR, rprimalPart labelR) pars+           artRaw = gradArtifact f (fromTarget targetInit, dataInit)+           art = simplifyArtifactGradient artRaw+           go :: [MnistDataBatchR r]+              -> (Concrete (XParams r), StateAdam (XParams r))+              -> (Concrete (XParams r), StateAdam (XParams r))+           go [] (parameters, stateAdam) = (parameters, stateAdam)+           go ((glyph, label) : rest) (!parameters, !stateAdam) =+             let parametersAndInput =+                   tpair parameters (tpair (rconcrete glyph) (rconcrete label))+                 gradient =+                   tproject1 $ fst+                   $ revInterpretArtifact art parametersAndInput Nothing+             in go rest (updateWithGradientAdam+                           @(XParams r)+                           defaultArgsAdam stateAdam knownSTK parameters+                           gradient)+           runBatch :: (Concrete (XParams r), StateAdam (XParams r))+                    -> (Int, [MnistDataR r])+                    -> IO (Concrete (XParams r), StateAdam (XParams r))+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkR = map mkMnistDataBatchR+                          $ filter (\ch -> length ch == miniBatchSize)+                          $ chunksOf miniBatchSize chunk+                 res@(parameters2, _) = go chunkR (parameters, stateAdam)+                 trainScore =+                   ftest (length chunk) (mkMnistDataBatchR chunk) parameters2+                 testScore =+                   ftest (totalBatchSize * maxBatches) testDataR parameters2+                 lenChunk = length chunk+             unless (n_hiddenInt < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> (Concrete (XParams r), StateAdam (XParams r))+                    -> IO (Concrete (XParams r))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (n_hiddenInt < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+           ftk = tftk @Concrete (knownSTK @(XParams r)) targetInit+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal =+             1 - ftest (totalBatchSize * maxBatches) testDataR res+       assertEqualUpToEpsilon 1e-1 expected testErrorFinal++{-# SPECIALIZE mnistTestCaseCNNRO+  :: String+  -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Double+  -> TestTree #-}++tensorADValMnistTestsCNNRO :: TestTree+tensorADValMnistTestsCNNRO = testGroup "CNNR Once MNIST tests"+  [ mnistTestCaseCNNRO "CNNRO 1 epoch, 1 batch" 1 1 4 4 8 16 1 1+                       (1 :: Double)+  , mnistTestCaseCNNRO "CNNRO artificial 1 2 3 4 5" 1 1 2 3 4 5 1 10+                       (1 :: Float)+  , mnistTestCaseCNNRO "CNNRO artificial 5 4 3 2 1" 5 4 3 2 1 1 1 1+                       (1 :: Double)+  , mnistTestCaseCNNRO "CNNRO 1 epoch, 0 batch" 1 0 4 4 16 64 16 50+                       (1.0 :: Float)+  ]
+ test/simplified/TestMnistCNNS.hs view
@@ -0,0 +1,422 @@+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Tests of "MnistCnnShaped2" convolutional neural network+-- using a few different optimization pipelines.+--+-- With the current CPU backend it's slow enough that it's hard to see+-- if it trains.+module TestMnistCNNS+  ( testTrees+  ) where++import Prelude++import Control.Monad (foldM, unless)+import GHC.TypeLits (KnownNat, type (<=))+import System.IO (hPutStrLn, stderr)+import System.Random+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)+import Text.Printf++import Data.Array.Nested.Shaped.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstInterpret++import EqEpsilon++import MnistCnnShaped2 qualified+import MnistData++-- TODO: optimize enough that it can run for one full epoch in reasonable time+-- and then verify it trains down to ~20% validation error in a short enough+-- time to include such a training run in tests.++testTrees :: [TestTree]+testTrees = [ tensorADValMnistTestsCNNSA+            , tensorADValMnistTestsCNNSI+            , tensorADValMnistTestsCNNSO+            ]++type XParams kh kw c_out n_hidden r =+  X (MnistCnnShaped2.ADCnnMnistParametersShaped+       Concrete SizeMnistHeight SizeMnistWidth kh kw c_out n_hidden r)++-- POPL differentiation, straight via the ADVal instance of RankedTensor,+-- which side-steps vectorization.+mnistTestCaseCNNSA+  :: forall kh kw r.+     ( 1 <= kh, 1 <= kw+     , Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r )+  => String+  -> Int -> Int -> SNat kh -> SNat kw -> Int -> Int -> Int -> Int -> r+  -> TestTree+mnistTestCaseCNNSA prefix epochs maxBatches kh@SNat kw@SNat c_outInt n_hiddenInt+                   miniBatchSizeInt totalBatchSize expected =+  withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->+  withSNat n_hiddenInt $ \(_n_hiddenSNat :: SNat n_hidden) ->+  withSNat miniBatchSizeInt $ \(miniBatchSize :: SNat miniBatchSize) ->+  let targetInit =+        fst $ randomValue+                @(Concrete (X (MnistCnnShaped2.ADCnnMnistParametersShaped+                                 Concrete SizeMnistHeight SizeMnistWidth+                                 kh kw c_out n_hidden r)))+                0.4 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show (sNatValue kh), show (sNatValue kw)+                        , show c_outInt, show n_hiddenInt+                        , show miniBatchSizeInt+                        , show $ widthSTK $ knownSTK @(XParams kh kw c_out n_hidden r)+                        , show (tsize knownSTK targetInit) ]+      ftest :: KnownNat batch_size+            => MnistDataBatchS batch_size r+            -> Concrete (XParams kh kw c_out n_hidden r) -> r+      ftest @batch_size mnistData pars =+        MnistCnnShaped2.convMnistTestS kh kw (SNat @c_out) (SNat @n_hidden)+          (SNat @batch_size) mnistData (fromTarget @Concrete pars)+  in testCase name $ do+      hPutStrLn stderr $+        printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+               prefix epochs maxBatches+      trainData <- map mkMnistDataS+                   <$> loadMnistData trainGlyphsPath trainLabelsPath+      testData <- map mkMnistDataS . take (totalBatchSize * maxBatches)+                  <$> loadMnistData testGlyphsPath testLabelsPath+      withSNat (totalBatchSize * maxBatches) $ \(SNat @lenTestData) -> do+       let testDataS = mkMnistDataBatchS testData+           f :: MnistDataBatchS miniBatchSize r+             -> ADVal Concrete (XParams kh kw c_out n_hidden r)+             -> ADVal Concrete (TKScalar r)+           f (glyphR, labelR) adinputs =+             MnistCnnShaped2.convMnistLossFusedS+               kh kw (SNat @c_out) (SNat @n_hidden)+               miniBatchSize (sconcrete glyphR, sconcrete labelR)+               (fromTarget adinputs)+           runBatch :: (Concrete (XParams kh kw c_out n_hidden r), StateAdam (XParams kh kw c_out n_hidden r))+                    -> (Int, [MnistDataS r])+                    -> IO (Concrete (XParams kh kw c_out n_hidden r), StateAdam (XParams kh kw c_out n_hidden r))+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkS = map mkMnistDataBatchS+                          $ filter (\ch -> length ch == miniBatchSizeInt)+                          $ chunksOf miniBatchSizeInt chunk+                 res@(parameters2, _) =+                   sgdAdam f chunkS parameters stateAdam+                 trainScore = withSNat (length chunk) $ \(SNat @len) ->+                   ftest @len (mkMnistDataBatchS chunk) parameters2+                 testScore = ftest @lenTestData testDataS parameters2+                 lenChunk = length chunk+             unless (n_hiddenInt < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> (Concrete (XParams kh kw c_out n_hidden r), StateAdam (XParams kh kw c_out n_hidden r))+                    -> IO (Concrete (XParams kh kw c_out n_hidden r))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (n_hiddenInt < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+           ftk = tftk @Concrete (knownSTK @(XParams kh kw c_out n_hidden r)) targetInit+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal =+             1 - ftest testDataS res+       testErrorFinal @?~ expected++tensorADValMnistTestsCNNSA :: TestTree+tensorADValMnistTestsCNNSA = testGroup "CNNS ADVal MNIST tests"+  [ mnistTestCaseCNNSA "CNNSA 1 epoch, 1 batch"+                       1 1 (SNat @4) (SNat @4) 8 16 1 1+                       (1 :: Double)+  , mnistTestCaseCNNSA "CNNSA artificial 1 2 3 4 5"+                       1 1 (SNat @2) (SNat @3) 4 5 1 10+                       (1 :: Float)+  , mnistTestCaseCNNSA "CNNSA artificial 5 4 3 2 1"+                       5 4 (SNat @3) (SNat @2) 1 1 1 1+                       (1 :: Double)+  , mnistTestCaseCNNSA "CNNSA 1 epoch, 0 batch"+                       1 0 (SNat @4) (SNat @4) 16 64 16 50+                       (1.0 :: Float)+  ]++-- POPL differentiation, with Ast term defined and vectorized only once,+-- but differentiated anew in each gradient descent iteration.+mnistTestCaseCNNSI+  :: forall kh kw r.+     ( 1 <= kh, 1 <= kw+     , Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r )+  => String+  -> Int -> Int -> SNat kh -> SNat kw -> Int -> Int -> Int -> Int -> r+  -> TestTree+mnistTestCaseCNNSI prefix epochs maxBatches kh@SNat kw@SNat c_outInt n_hiddenInt+                   miniBatchSizeInt totalBatchSize expected =+  withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->+  withSNat n_hiddenInt $ \(_n_hiddenSNat :: SNat n_hidden) ->+  withSNat miniBatchSizeInt $ \(miniBatchSize :: SNat miniBatchSize) ->+  let targetInit =+        fst $ randomValue+                @(Concrete (X (MnistCnnShaped2.ADCnnMnistParametersShaped+                                 Concrete SizeMnistHeight SizeMnistWidth+                                 kh kw c_out n_hidden r)))+                0.4 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show (sNatValue kh), show (sNatValue kw)+                        , show c_outInt, show n_hiddenInt+                        , show miniBatchSizeInt+                        , show $ widthSTK $ knownSTK @(XParams kh kw c_out n_hidden r)+                        , show (tsize knownSTK targetInit) ]+      ftest :: KnownNat batch_size+            => MnistDataBatchS batch_size r+            -> Concrete (XParams kh kw c_out n_hidden r) -> r+      ftest @batch_size mnistData pars =+        MnistCnnShaped2.convMnistTestS kh kw (SNat @c_out) (SNat @n_hidden)+          (SNat @batch_size) mnistData (fromTarget @Concrete pars)+  in testCase name $ do+      hPutStrLn stderr $+        printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+               prefix epochs maxBatches+      trainData <- map mkMnistDataS+                   <$> loadMnistData trainGlyphsPath trainLabelsPath+      testData <- map mkMnistDataS . take (totalBatchSize * maxBatches)+                  <$> loadMnistData testGlyphsPath testLabelsPath+      withSNat (totalBatchSize * maxBatches) $ \(SNat @lenTestData) -> do+       let testDataS = mkMnistDataBatchS testData+           ftk = tftk @Concrete (knownSTK @(XParams kh kw c_out n_hidden r)) targetInit+       (_, _, var, varAst2) <- funToAstRevIO ftk+       (varGlyph, astGlyph) <-+         funToAstIO (FTKS (miniBatchSize+                           :$$ sizeMnistHeight+                           :$$ sizeMnistWidth+                           :$$ ZSS) FTKScalar) id+       (varLabel, astLabel) <-+         funToAstIO (FTKS (miniBatchSize+                           :$$ sizeMnistLabel+                           :$$ ZSS) FTKScalar) id+       let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)+           ast = simplifyInline+                 $ MnistCnnShaped2.convMnistLossFusedS+                     kh kw (SNat @c_out) (SNat @n_hidden)+                     miniBatchSize (astGlyph, astLabel)+                     (fromTarget varAst2)+           f :: MnistDataBatchS miniBatchSize r+             -> ADVal Concrete (XParams kh kw c_out n_hidden r)+             -> ADVal Concrete (TKScalar r)+           f (glyph, label) varInputs =+             let env = extendEnv var varInputs emptyEnv+                 envMnist = extendEnv varGlyph (sconcrete glyph)+                            $ extendEnv varLabel (sconcrete label) env+             in interpretAstFull envMnist ast+           runBatch :: (Concrete (XParams kh kw c_out n_hidden r), StateAdam (XParams kh kw c_out n_hidden r))+                    -> (Int, [MnistDataS r])+                    -> IO (Concrete (XParams kh kw c_out n_hidden r), StateAdam (XParams kh kw c_out n_hidden r))+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkS = map mkMnistDataBatchS+                          $ filter (\ch -> length ch == miniBatchSizeInt)+                          $ chunksOf miniBatchSizeInt chunk+                 res@(parameters2, _) =+                   sgdAdam f chunkS parameters stateAdam+                 !trainScore = withSNat (length chunk) $ \(SNat @len) ->+                   ftest @len (mkMnistDataBatchS chunk) parameters2+                 !testScore = ftest @lenTestData testDataS parameters2+                 !lenChunk = length chunk+             unless (n_hiddenInt < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> (Concrete (XParams kh kw c_out n_hidden r), StateAdam (XParams kh kw c_out n_hidden r))+                    -> IO (Concrete (XParams kh kw c_out n_hidden r))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (n_hiddenInt < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal =+             1 - ftest testDataS res+       testErrorFinal @?~ expected++tensorADValMnistTestsCNNSI :: TestTree+tensorADValMnistTestsCNNSI = testGroup "CNNS Intermediate MNIST tests"+  [ mnistTestCaseCNNSI "CNNSI 1 epoch, 1 batch"+                       1 1 (SNat @4) (SNat @4) 8 16 1 1+                       (1 :: Double)+  , mnistTestCaseCNNSI "CNNSI artificial 1 2 3 4 5"+                       1 1 (SNat @2) (SNat @3) 4 5 1 10+                       (1 :: Float)+  , mnistTestCaseCNNSI "CNNSI artificial 5 4 3 2 1"+                       5 4 (SNat @3) (SNat @2) 1 1 1 1+                       (1 :: Double)+  , mnistTestCaseCNNSI "CNNSI 1 epoch, 0 batch"+                       1 0 (SNat @4) (SNat @4) 16 64 16 50+                       (1.0 :: Float)+  ]++-- JAX differentiation, Ast term built and differentiated only once+-- and the result interpreted with different inputs in each gradient+-- descent iteration.+mnistTestCaseCNNSO+  :: forall kh kw r.+     ( 1 <= kh, 1 <= kw+     , Differentiable r, GoodScalar r+     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+  => String+  -> Int -> Int -> SNat kh -> SNat kw -> Int -> Int -> Int -> Int -> r+  -> TestTree+mnistTestCaseCNNSO prefix epochs maxBatches kh@SNat kw@SNat c_outInt n_hiddenInt+                   miniBatchSizeInt totalBatchSize expected =+  withSNat c_outInt $ \(_c_outSNat :: SNat c_out) ->+  withSNat n_hiddenInt $ \(_n_hiddenSNat :: SNat n_hidden) ->+  withSNat miniBatchSizeInt $ \(miniBatchSize :: SNat miniBatchSize) ->+  let targetInit =+        fst $ randomValue+                @(Concrete (X (MnistCnnShaped2.ADCnnMnistParametersShaped+                                 Concrete SizeMnistHeight SizeMnistWidth+                                 kh kw c_out n_hidden r)))+                0.4 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show (sNatValue kh), show (sNatValue kw)+                        , show c_outInt, show n_hiddenInt+                        , show miniBatchSizeInt+                        , show $ widthSTK+                          $ knownSTK @(XParams kh kw c_out n_hidden r)+                        , show (tsize knownSTK targetInit) ]+      ftest :: KnownNat batch_size+            => MnistDataBatchS batch_size r+            -> Concrete (XParams kh kw c_out n_hidden r) -> r+      ftest @batch_size mnistData pars =+        MnistCnnShaped2.convMnistTestS kh kw (SNat @c_out) (SNat @n_hidden)+          (SNat @batch_size) mnistData (fromTarget @Concrete pars)+  in testCase name $ do+      hPutStrLn stderr $+        printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+               prefix epochs maxBatches+      trainData <- map mkMnistDataS+                   <$> loadMnistData trainGlyphsPath trainLabelsPath+      testData <- map mkMnistDataS . take (totalBatchSize * maxBatches)+                  <$> loadMnistData testGlyphsPath testLabelsPath+      withSNat (totalBatchSize * maxBatches) $ \(SNat @lenTestData) -> do+       let testDataS = mkMnistDataBatchS testData+           dataInit = case chunksOf miniBatchSizeInt testData of+             d : _ -> let (dglyph, dlabel) = mkMnistDataBatchS d+                      in (sconcrete dglyph, sconcrete dlabel)+             [] -> error "empty test data"+           f :: ( MnistCnnShaped2.ADCnnMnistParametersShaped+                    (AstTensor AstMethodLet FullSpan)+                    SizeMnistHeight SizeMnistWidth+                    kh kw c_out n_hidden r+                , ( AstTensor AstMethodLet FullSpan (TKS '[miniBatchSize, SizeMnistHeight, SizeMnistWidth] r)+                  , AstTensor AstMethodLet FullSpan (TKS '[miniBatchSize, SizeMnistLabel] r) ) )+             -> AstTensor AstMethodLet FullSpan (TKScalar r)+           f = \ (pars, (glyphR, labelR)) ->+             MnistCnnShaped2.convMnistLossFusedS+               kh kw (SNat @c_out) (SNat @n_hidden)+               miniBatchSize (sprimalPart glyphR, sprimalPart labelR) pars+           artRaw = gradArtifact f (fromTarget targetInit, dataInit)+           art = simplifyArtifactGradient artRaw+           go :: [MnistDataBatchS miniBatchSize r]+              -> ( Concrete (XParams kh kw c_out n_hidden r)+                 , StateAdam (XParams kh kw c_out n_hidden r) )+              -> ( Concrete (XParams kh kw c_out n_hidden r)+                 , StateAdam (XParams kh kw c_out n_hidden r) )+           go [] (parameters, stateAdam) = (parameters, stateAdam)+           go ((glyph, label) : rest) (!parameters, !stateAdam) =+             let parametersAndInput =+                   tpair parameters (tpair (sconcrete glyph) (sconcrete label))+                 gradient =+                   tproject1 $ fst+                   $ revInterpretArtifact art parametersAndInput Nothing+             in go rest (updateWithGradientAdam+                           @(XParams kh kw c_out n_hidden r)+                           defaultArgsAdam stateAdam knownSTK parameters+                           gradient)+           runBatch :: ( Concrete (XParams kh kw c_out n_hidden r)+                       , StateAdam (XParams kh kw c_out n_hidden r) )+                    -> (Int, [MnistDataS r])+                    -> IO ( Concrete (XParams kh kw c_out n_hidden r)+                          , StateAdam (XParams kh kw c_out n_hidden r) )+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkS = map mkMnistDataBatchS+                          $ filter (\ch -> length ch == miniBatchSizeInt)+                          $ chunksOf miniBatchSizeInt chunk+                 res@(parameters2, _) = go chunkS (parameters, stateAdam)+                 trainScore = withSNat (length chunk) $ \(SNat @len) ->+                   ftest @len (mkMnistDataBatchS chunk) parameters2+                 testScore = ftest @lenTestData testDataS parameters2+                 lenChunk = length chunk+             unless (n_hiddenInt < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> (Concrete (XParams kh kw c_out n_hidden r), StateAdam (XParams kh kw c_out n_hidden r))+                    -> IO (Concrete (XParams kh kw c_out n_hidden r))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (n_hiddenInt < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+           ftk = tftk @Concrete (knownSTK @(XParams kh kw c_out n_hidden r)) targetInit+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal =+             1 - ftest testDataS res+       assertEqualUpToEpsilon 1e-1 expected testErrorFinal++tensorADValMnistTestsCNNSO :: TestTree+tensorADValMnistTestsCNNSO = testGroup "CNNS Once MNIST tests"+  [ mnistTestCaseCNNSO "CNNSO 1 epoch, 1 batch"+                       1 1 (SNat @4) (SNat @4) 8 16 1 1+                       (1 :: Double)+  , mnistTestCaseCNNSO "CNNSO artificial 1 2 3 4 5"+                       1 1 (SNat @2) (SNat @3) 4 5 1 10+                       (1 :: Float)+  , mnistTestCaseCNNSO "CNNSO artificial 5 4 3 2 1"+                       5 4 (SNat @3) (SNat @2) 1 1 1 1+                       (1 :: Double)+  , mnistTestCaseCNNSO "CNNSO 1 epoch, 0 batch"+                       1 0 (SNat @4) (SNat @4) 16 64 16 50+                       (1.0 :: Float)+  ]
+ test/simplified/TestMnistFCNNR.hs view
@@ -0,0 +1,783 @@+{-# LANGUAGE OverloadedLists #-}+-- | Tests of "MnistFcnnRanked1" and "MnistFcnnRanked2" dense neural networks+-- using a few different optimization pipelines.+module TestMnistFCNNR+  ( testTrees+  ) where++import Prelude++import Control.Arrow ((***))+import Control.Monad (foldM, unless)+import Data.Bifunctor (first)+import Data.Proxy (Proxy (Proxy))+import System.IO (hPutStrLn, stderr)+import System.Random+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)+import Test.Tasty.QuickCheck hiding (label, shuffle)+import Text.Printf++import Data.Array.Nested.Ranked.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstInterpret+import HordeAd.Core.Ops (tconcrete)++import CrossTesting+import EqEpsilon++import MnistData+import MnistFcnnRanked1 qualified+import MnistFcnnRanked2 (XParams2)+import MnistFcnnRanked2 qualified++testTrees :: [TestTree]+testTrees = [ testCase "2VTOrev" mnistTestCase2VTOrev+            , tensorADValMnistTests+            , tensorIntermediateMnistTests+            , tensorADOnceMnistTests+            , tensorADValMnistTests2+            , tensorIntermediateMnistTests2+            , tensorADOnceMnistTests2+            ]+++-- * Running rev' on the gradient of afcnnMnistLoss2++mnistTestCase2VTOrev :: Assertion+mnistTestCase2VTOrev =+  let (!targetInit, !art) =+        MnistFcnnRanked2.mnistTrainBench2VTOGradientX+          @Double (Proxy @Float) IgnoreIncomingCotangent+          1 (mkStdGen 44) 1500 500+      blackGlyph = rreplicate sizeMnistGlyphInt $ rscalar 7+      ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))+                 targetInit+      f :: forall target r. (ADReady target, r ~ Double)+        => target (TKR 1 r) -> target (TKR 1 r)+      f label =+        let val = tpair (tconcrete ftk targetInit)+                        (tpair (rconcrete $ unConcrete blackGlyph) label)+            env = extendEnv (artVarDomainRev art) val emptyEnv+        in tproject1 $ tproject2+           $ interpretAst @target env (artDerivativeRev art)+  in assertEqualUpToEpsilon' 1e-10+       (ringestData [10] [6.922657834114052e-2,-3.2210167235305924e-5,0.12334696753032606,-4.892729845753193e-3,3.010762414514606e-2,2.0344986964700877e-2,-3.78339785604896e-2,5.77360835535866e-2,0.10761507003315526,-7.909016076299641e-2])+       (rev' f (rreplicate sizeMnistLabelInt $ rscalar 8))+++-- * Using lists of vectors, which is rank 1++type XParams widthHidden widthHidden2 r =+  X (MnistFcnnRanked1.ADFcnnMnist1Parameters+       Concrete widthHidden widthHidden2 r)++-- POPL differentiation, straight via the ADVal instance of RankedTensor,+-- which side-steps vectorization.+mnistTestCase1VTA+  :: forall r.+     ( Differentiable r, GoodScalar r+     , PrintfArg r, AssertEqualUpToEpsilon r )+  => String+  -> Int -> Int -> Int -> Int -> Double -> Int -> r+  -> TestTree+mnistTestCase1VTA prefix epochs maxBatches widthHiddenInt widthHidden2Int+                  gamma batchSize expected =+  withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->+  withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters+                    Concrete widthHidden widthHidden2 r+      valsInit = fst $ randomValue 1 (mkStdGen 44)+      targetInit :: Concrete (XParams widthHidden widthHidden2 r)+      targetInit = toTarget @Concrete valsInit+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show widthHiddenInt, show widthHidden2Int+                        , show $ widthSTK+                          $ knownSTK @(XParams widthHidden widthHidden2 r)+                        , show (tsize knownSTK targetInit)+                        , show gamma ]+      ftest :: [MnistDataLinearR r]+            -> MnistFcnnRanked1.ADFcnnMnist1Parameters+                 Concrete widthHidden widthHidden2 r+            -> r+      ftest = MnistFcnnRanked1.afcnnMnistTest1 widthHiddenSNat widthHidden2SNat+  in testCase name $ do+    hPutStrLn stderr $+      printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+             prefix epochs maxBatches+    trainData <- loadMnistData trainGlyphsPath trainLabelsPath+    testData <- map mkMnistDataLinearR . take (batchSize * maxBatches)+                <$> loadMnistData testGlyphsPath testLabelsPath+    let f :: MnistDataLinearR r+          -> ADVal Concrete (XParams widthHidden widthHidden2 r)+          -> ADVal Concrete (TKScalar r)+        f (glyph, label) adinputs =+          MnistFcnnRanked1.afcnnMnistLoss1+            widthHiddenSNat widthHidden2SNat+            (rconcrete glyph, rconcrete label) (fromTarget adinputs)+    -- Mimic how backprop tests and display it, even though tests+    -- should not print, in principle.+    let runBatch :: Concrete (XParams widthHidden widthHidden2 r)+                 -> (Int, [MnistDataLinearR r])+                 -> IO (Concrete (XParams widthHidden widthHidden2 r))+        runBatch !params (k, chunk) = do+          let res = fst $ sgd gamma f chunk params+              trainScore = ftest chunk (fromTarget res)+              testScore = ftest testData (fromTarget res)+              lenChunk = length chunk+          unless (widthHiddenInt < 10) $ do+            hPutStrLn stderr $+              printf "\n%s: (Batch %d with %d points)"+                     prefix k lenChunk+            hPutStrLn stderr $+              printf "%s: Training error:   %.2f%%"+                     prefix ((1 - trainScore) * 100)+            hPutStrLn stderr $+              printf "%s: Validation error: %.2f%%"+                     prefix ((1 - testScore ) * 100)+          return res+    let runEpoch :: Int+                 -> Concrete (XParams widthHidden widthHidden2 r)+                 -> IO (Concrete (XParams widthHidden widthHidden2 r))+        runEpoch n params | n > epochs = return params+        runEpoch n !params = do+          unless (widthHiddenInt < 10) $+            hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+          let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+              chunks = take maxBatches+                       $ zip [1 ..] $ chunksOf batchSize+                       $ map mkMnistDataLinearR trainDataShuffled+          res <- foldM runBatch params chunks+          runEpoch (succ n) res+    res <- runEpoch 1 targetInit+    let testErrorFinal = 1 - ftest testData (fromTarget res)+    testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCase1VTA+  :: String+  -> Int -> Int -> Int -> Int -> Double -> Int -> Double+  -> TestTree #-}++tensorADValMnistTests :: TestTree+tensorADValMnistTests = testGroup "Ranked ADVal MNIST tests"+  [ mnistTestCase1VTA "VTA1 1 epoch, 1 batch" 1 1 300 100 0.02 5000+                      (0.2146 :: Double)+  , mnistTestCase1VTA "VTA1 artificial 1 2 3 4 5" 1 2 3 4 5 5000+                      (0.8972 :: Float)+  , mnistTestCase1VTA "VTA1 1 epoch, 0 batch" 1 0 300 100 0.02 5000+                      (1 :: Float)+  ]++-- POPL differentiation, with Ast term defined and vectorized only once,+-- but differentiated anew in each gradient descent iteration.+mnistTestCase1VTI+  :: forall r.+     ( Differentiable r, GoodScalar r+     , PrintfArg r, AssertEqualUpToEpsilon r )+  => String+  -> Int -> Int -> Int -> Int -> Double -> Int -> r+  -> TestTree+mnistTestCase1VTI prefix epochs maxBatches widthHiddenInt widthHidden2Int+                  gamma batchSize expected =+  withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->+  withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters+                    Concrete widthHidden widthHidden2 r+      valsInit = fst $ randomValue 1 (mkStdGen 44)+      targetInit :: Concrete (XParams widthHidden widthHidden2 r)+      targetInit = toTarget @Concrete valsInit+      ftk = tftk @Concrete (knownSTK @(XParams widthHidden widthHidden2 r))+                 targetInit+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show widthHiddenInt, show widthHidden2Int+                        , show $ widthSTK+                          $ knownSTK @(XParams widthHidden widthHidden2 r)+                        , show (tsize knownSTK targetInit)+                        , show gamma ]+      ftest :: [MnistDataLinearR r]+            -> MnistFcnnRanked1.ADFcnnMnist1Parameters+                 Concrete widthHidden widthHidden2 r+            -> r+      ftest = MnistFcnnRanked1.afcnnMnistTest1 widthHiddenSNat widthHidden2SNat+  in testCase name $ do+    hPutStrLn stderr $+      printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+             prefix epochs maxBatches+    trainData <- loadMnistData trainGlyphsPath trainLabelsPath+    testData <- map mkMnistDataLinearR . take (batchSize * maxBatches)+                <$> loadMnistData testGlyphsPath testLabelsPath+    (_, _, var, varAst) <- funToAstRevIO ftk+    (varGlyph, astGlyph) <-+      funToAstIO (FTKR (sizeMnistGlyphInt :$: ZSR) FTKScalar) id+    (varLabel, astLabel) <-+      funToAstIO (FTKR (sizeMnistLabelInt :$: ZSR) FTKScalar) id+    let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)+        ast = simplifyInline+              $ MnistFcnnRanked1.afcnnMnistLoss1+                  widthHiddenSNat widthHidden2SNat (astGlyph, astLabel)+                  (fromTarget varAst)+        f :: MnistDataLinearR r+          -> ADVal Concrete (XParams widthHidden widthHidden2 r)+          -> ADVal Concrete (TKScalar r)+        f (glyph, label) varInputs =+          let env = extendEnv var varInputs emptyEnv+              envMnist = extendEnv varGlyph (rconcrete glyph)+                         $ extendEnv varLabel (rconcrete label) env+          in interpretAstFull envMnist ast+    let runBatch :: Concrete (XParams widthHidden widthHidden2 r)+                 -> (Int, [MnistDataLinearR r])+                 -> IO (Concrete (XParams widthHidden widthHidden2 r))+        runBatch !params (k, chunk) = do+          let res = fst $ sgd gamma f chunk params+              trainScore = ftest chunk (fromTarget res)+              testScore = ftest testData (fromTarget res)+              lenChunk = length chunk+          unless (widthHiddenInt < 10) $ do+            hPutStrLn stderr $+              printf "\n%s: (Batch %d with %d points)"+                     prefix k lenChunk+            hPutStrLn stderr $+              printf "%s: Training error:   %.2f%%"+                     prefix ((1 - trainScore) * 100)+            hPutStrLn stderr $+              printf "%s: Validation error: %.2f%%"+                     prefix ((1 - testScore ) * 100)+          return res+    let runEpoch :: Int+                 -> Concrete (XParams widthHidden widthHidden2 r)+                 -> IO (Concrete (XParams widthHidden widthHidden2 r))+        runEpoch n params | n > epochs = return params+        runEpoch n !params = do+          unless (widthHiddenInt < 10) $+            hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+          let trainDataShuffled = shuffle (mkStdGen $ n + 1) trainData+              chunks = take maxBatches+                       $ zip [1 ..] $ chunksOf batchSize+                       $ map mkMnistDataLinearR trainDataShuffled+          res <- foldM runBatch params chunks+          runEpoch (succ n) res+    res <- runEpoch 1 targetInit+    let testErrorFinal = 1 - ftest testData (fromTarget res)+    testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCase1VTI+  :: String+  -> Int -> Int -> Int -> Int -> Double -> Int -> Double+  -> TestTree #-}++tensorIntermediateMnistTests :: TestTree+tensorIntermediateMnistTests = testGroup "Ranked Intermediate MNIST tests"+  [ mnistTestCase1VTI "VTI1 1 epoch, 1 batch" 1 1 300 100 0.02 5000+                      (0.2116 :: Double)+  , mnistTestCase1VTI "VTI1 artificial 1 2 3 4 5" 1 2 3 4 5 5000+                      (0.9108 :: Float)+  , mnistTestCase1VTI "VTI1 1 epoch, 0 batch" 1 0 300 100 0.02 5000+                      (1 :: Float)+  ]++-- JAX differentiation, Ast term built and differentiated only once+-- and the result interpreted with different inputs in each gradient+-- descent iteration.+mnistTestCase1VTO+  :: forall r.+     ( Differentiable r, GoodScalar r, ADTensorScalar r ~ r+     , PrintfArg r, AssertEqualUpToEpsilon r)+  => String+  -> Int -> Int -> Int -> Int -> Double -> Int -> r+  -> TestTree+mnistTestCase1VTO prefix epochs maxBatches widthHiddenInt widthHidden2Int+                  gamma batchSize expected =+  withSNat widthHiddenInt $ \(widthHiddenSNat :: SNat widthHidden) ->+  withSNat widthHidden2Int $ \(widthHidden2SNat :: SNat widthHidden2) ->+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[SizeMnistGlyph] r)) (SNat @widthHidden)) $+  withKnownSTK+    (stkOfListR (knownSTK @(TKS '[widthHidden] Float)) (SNat @widthHidden2)) $+  let valsInit :: MnistFcnnRanked1.ADFcnnMnist1Parameters+                    Concrete widthHidden widthHidden2 r+      valsInit = fst $ randomValue 1 (mkStdGen 44)+      targetInit :: Concrete (XParams widthHidden widthHidden2 r)+      targetInit = toTarget @Concrete valsInit+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show widthHiddenInt, show widthHidden2Int+                        , show $ widthSTK+                          $ knownSTK @(XParams widthHidden widthHidden2 r)+                        , show (tsize knownSTK targetInit)+                        , show gamma ]+      ftest :: [MnistDataLinearR r]+            -> MnistFcnnRanked1.ADFcnnMnist1Parameters+                 Concrete widthHidden widthHidden2 r+            -> r+      ftest = MnistFcnnRanked1.afcnnMnistTest1 widthHiddenSNat widthHidden2SNat+  in testCase name $ do+    hPutStrLn stderr $+      printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+             prefix epochs maxBatches+    trainData <- loadMnistData trainGlyphsPath trainLabelsPath+    testData <- map mkMnistDataLinearR . take (batchSize * maxBatches)+                <$> loadMnistData testGlyphsPath testLabelsPath+    let dataInit = case testData of+          d : _ -> (rconcrete *** rconcrete) d+          [] -> error "empty test data"+        f :: ( MnistFcnnRanked1.ADFcnnMnist1Parameters+                 (AstTensor AstMethodLet FullSpan)+                 widthHidden widthHidden2 r+             , ( AstTensor AstMethodLet FullSpan (TKR 1 r)+               , AstTensor AstMethodLet FullSpan (TKR 1 r) ) )+          -> AstTensor AstMethodLet FullSpan (TKScalar r)+        f = \ (pars, (glyphR, labelR)) ->+          MnistFcnnRanked1.afcnnMnistLoss1+            widthHiddenSNat widthHidden2SNat+            (glyphR, labelR) pars+        artRaw = gradArtifact f (valsInit, dataInit)+        art = simplifyArtifactGradient artRaw+        go :: [MnistDataLinearR r]+           -> Concrete (XParams widthHidden widthHidden2 r)+           -> Concrete (XParams widthHidden widthHidden2 r)+        go [] parameters = parameters+        go ((glyph, label) : rest) !parameters =+          let parametersAndInput =+                tpair parameters (tpair (rconcrete glyph) (rconcrete label))+              gradient = tproject1 $ fst+                         $ revInterpretArtifact art parametersAndInput Nothing+          in go rest (updateWithGradient gamma knownSTK parameters gradient)+    let runBatch :: Concrete (XParams widthHidden widthHidden2 r)+                 -> (Int, [MnistDataLinearR r])+                 -> IO (Concrete (XParams widthHidden widthHidden2 r))+        runBatch !params (k, chunk) = do+          let res = go chunk params+              trainScore = ftest chunk (fromTarget res)+              testScore = ftest testData (fromTarget res)+              lenChunk = length chunk+          unless (widthHiddenInt < 10) $ do+            hPutStrLn stderr $+              printf "\n%s: (Batch %d with %d points)"+                     prefix k lenChunk+            hPutStrLn stderr $+              printf "%s: Training error:   %.2f%%"+                     prefix ((1 - trainScore) * 100)+            hPutStrLn stderr $+              printf "%s: Validation error: %.2f%%"+                     prefix ((1 - testScore ) * 100)+          return res+    let runEpoch :: Int+                 -> Concrete (XParams widthHidden widthHidden2 r)+                 -> IO (Concrete (XParams widthHidden widthHidden2 r))+        runEpoch n params | n > epochs = return params+        runEpoch n !params = do+          unless (widthHiddenInt < 10) $+            hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+          let trainDataShuffled = shuffle (mkStdGen $ n + 1) trainData+              chunks = take maxBatches+                       $ zip [1 ..] $ chunksOf batchSize+                       $ map mkMnistDataLinearR trainDataShuffled+          res <- foldM runBatch params chunks+          runEpoch (succ n) res+    res <- runEpoch 1 targetInit+    let testErrorFinal = 1 - ftest testData (fromTarget res)+    testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCase1VTO+  :: String+  -> Int -> Int -> Int -> Int -> Double -> Int -> Double+  -> TestTree #-}++tensorADOnceMnistTests :: TestTree+tensorADOnceMnistTests = testGroup "Ranked Once MNIST tests"+  [ mnistTestCase1VTO "VTO1 1 epoch, 1 batch" 1 1 300 100 0.02 5000+                      (0.2116 :: Double)+  , mnistTestCase1VTO "VTO1 artificial 1 2 3 4 5" 1 2 3 4 5 5000+                      (0.9108 :: Float)+  , mnistTestCase1VTO "VTO1 1 epoch, 0 batch" 1 0 300 100 0.02 5000+                      (1 :: Float)+  ]+++-- * Using matrices, which is rank 2++-- POPL differentiation, straight via the ADVal instance of RankedTensor,+-- which side-steps vectorization.+mnistTestCase2VTA+  :: forall r.+     ( Differentiable r, GoodScalar r+     , PrintfArg r, AssertEqualUpToEpsilon r )+  => String+  -> Int -> Int -> Int -> Int -> Double -> Int -> r+  -> TestTree+mnistTestCase2VTA prefix epochs maxBatches widthHidden widthHidden2+                  gamma batchSize expected =+  withSNat widthHidden $ \(SNat @widthHidden) ->+  withSNat widthHidden2 $ \(SNat @widthHidden2) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue+            @(Concrete (X (MnistFcnnRanked2.ADFcnnMnist2ParametersShaped+                             Concrete widthHidden widthHidden2 r Float)))+            1 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show widthHidden, show widthHidden2+                        , show $ widthSTK $ knownSTK @(XParams2 r Float)+                        , show (tsize knownSTK targetInit)+                        , show gamma ]+  in testCase name $ do+    hPutStrLn stderr $+      printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+             prefix epochs maxBatches+    trainData <- loadMnistData trainGlyphsPath trainLabelsPath+    testData <- map mkMnistDataLinearR . take (batchSize * maxBatches)+                <$> loadMnistData testGlyphsPath testLabelsPath+    let f :: MnistDataLinearR r -> ADVal Concrete (XParams2 r Float)+          -> ADVal Concrete (TKScalar r)+        f (glyph, label) adinputs =+          MnistFcnnRanked2.afcnnMnistLoss2+            (rconcrete glyph, rconcrete label) (fromTarget adinputs)+    let runBatch :: Concrete (XParams2 r Float) -> (Int, [MnistDataLinearR r])+                 -> IO (Concrete (XParams2 r Float))+        runBatch !params (k, chunk) = do+          let res = fst $ sgd gamma f chunk params+              trainScore =+                MnistFcnnRanked2.afcnnMnistTest2 chunk (fromTarget res)+              testScore =+                MnistFcnnRanked2.afcnnMnistTest2 testData (fromTarget res)+              lenChunk = length chunk+          unless (widthHidden < 10) $ do+            hPutStrLn stderr $+              printf "\n%s: (Batch %d with %d points)"+                     prefix k lenChunk+            hPutStrLn stderr $+              printf "%s: Training error:   %.2f%%"+                     prefix ((1 - trainScore) * 100)+            hPutStrLn stderr $+              printf "%s: Validation error: %.2f%%"+                     prefix ((1 - testScore ) * 100)+          return res+    let runEpoch :: Int -> Concrete (XParams2 r Float)+                 -> IO (Concrete (XParams2 r Float))+        runEpoch n params | n > epochs = return params+        runEpoch n !params = do+          unless (widthHidden < 10) $+            hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+          let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+              chunks = take maxBatches+                       $ zip [1 ..] $ chunksOf batchSize+                       $ map mkMnistDataLinearR trainDataShuffled+          res <- foldM runBatch params chunks+          runEpoch (succ n) res+    res <- runEpoch 1 targetInit+    let testErrorFinal =+          1 - MnistFcnnRanked2.afcnnMnistTest2 testData (fromTarget res)+    testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCase2VTA+  :: String+  -> Int -> Int -> Int -> Int -> Double -> Int -> Double+  -> TestTree #-}++tensorADValMnistTests2 :: TestTree+tensorADValMnistTests2 = testGroup "Ranked2 ADVal MNIST tests"+  [ mnistTestCase2VTA "VTA2 1 epoch, 1 batch" 1 1 300 100 0.02 5000+                       (0.21299999999999997 :: Double)+  , mnistTestCase2VTA "VTA2 artificial 1 2 3 4 5" 1 2 3 4 5 5000+                       (0.8972 :: Float)+  , mnistTestCase2VTA "VTA2 artificial 5 4 3 2 1" 5 4 3 2 1 5000+                       (0.6805:: Double)+  , mnistTestCase2VTA "VTA2 1 epoch, 0 batch" 1 0 300 100 0.02 5000+                       (1 :: Float)+  ]++-- POPL differentiation, with Ast term defined and vectorized only once,+-- but differentiated anew in each gradient descent iteration.+mnistTestCase2VTI+  :: forall r.+     ( Differentiable r, GoodScalar r+     , PrintfArg r, AssertEqualUpToEpsilon r )+  => String+  -> Int -> Int -> Int -> Int -> Double -> Int -> r+  -> TestTree+mnistTestCase2VTI prefix epochs maxBatches widthHidden widthHidden2+                  gamma batchSize expected =+  withSNat widthHidden $ \(SNat @widthHidden) ->+  withSNat widthHidden2 $ \(SNat @widthHidden2) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue+            @(Concrete (X (MnistFcnnRanked2.ADFcnnMnist2ParametersShaped+                             Concrete widthHidden widthHidden2 r Float)))+            1 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show widthHidden, show widthHidden2+                        , show $ widthSTK $ knownSTK @(XParams2 r Float)+                        , show (tsize knownSTK targetInit)+                        , show gamma ]+  in testCase name $ do+    hPutStrLn stderr $+      printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+             prefix epochs maxBatches+    trainData <- loadMnistData trainGlyphsPath trainLabelsPath+    testData <- map mkMnistDataLinearR . take (batchSize * maxBatches)+                <$> loadMnistData testGlyphsPath testLabelsPath+    let ftk = tftk @Concrete (knownSTK @(XParams2 r Float)) targetInit+    (_, _, var, varAst) <- funToAstRevIO ftk+    (varGlyph, astGlyph) <-+      funToAstIO (FTKR (sizeMnistGlyphInt :$: ZSR) FTKScalar) id+    (varLabel, astLabel) <-+      funToAstIO (FTKR (sizeMnistLabelInt :$: ZSR) FTKScalar) id+    let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)+        ast = simplifyInline+              $ MnistFcnnRanked2.afcnnMnistLoss2+                  (astGlyph, astLabel)+                  (fromTarget varAst)+        f :: MnistDataLinearR r -> ADVal Concrete (XParams2 r Float)+          -> ADVal Concrete (TKScalar r)+        f (glyph, label) varInputs =+          let env = extendEnv var varInputs emptyEnv+              envMnist = extendEnv varGlyph (rconcrete glyph)+                         $ extendEnv varLabel (rconcrete label) env+          in interpretAstFull envMnist ast+    let runBatch :: Concrete (XParams2 r Float) -> (Int, [MnistDataLinearR r])+                 -> IO (Concrete (XParams2 r Float))+        runBatch !params (k, chunk) = do+          let res = fst $ sgd gamma f chunk params+              trainScore =+                MnistFcnnRanked2.afcnnMnistTest2 chunk (fromTarget res)+              testScore =+                MnistFcnnRanked2.afcnnMnistTest2 testData (fromTarget res)+              lenChunk = length chunk+          unless (widthHidden < 10) $ do+            hPutStrLn stderr $+              printf "\n%s: (Batch %d with %d points)"+                     prefix k lenChunk+            hPutStrLn stderr $+              printf "%s: Training error:   %.2f%%"+                     prefix ((1 - trainScore) * 100)+            hPutStrLn stderr $+              printf "%s: Validation error: %.2f%%"+                     prefix ((1 - testScore ) * 100)+          return res+    let runEpoch :: Int -> Concrete (XParams2 r Float)+                 -> IO (Concrete (XParams2 r Float))+        runEpoch n params | n > epochs = return params+        runEpoch n !params = do+          unless (widthHidden < 10) $+            hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+          let trainDataShuffled = shuffle (mkStdGen $ n + 1) trainData+              chunks = take maxBatches+                       $ zip [1 ..] $ chunksOf batchSize+                       $ map mkMnistDataLinearR trainDataShuffled+          res <- foldM runBatch params chunks+          runEpoch (succ n) res+    res <- runEpoch 1 targetInit+    let testErrorFinal =+          1 - MnistFcnnRanked2.afcnnMnistTest2 testData (fromTarget res)+    testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCase2VTI+  :: String+  -> Int -> Int -> Int -> Int -> Double -> Int -> Double+  -> TestTree #-}++tensorIntermediateMnistTests2 :: TestTree+tensorIntermediateMnistTests2 = testGroup "Ranked2 Intermediate MNIST tests"+  [ mnistTestCase2VTI "VTI2 1 epoch, 1 batch" 1 1 300 100 0.02 5000+                       (0.20779999999999998 :: Double)+  , mnistTestCase2VTI "VTI2 artificial 1 2 3 4 5" 1 2 3 4 5 5000+                       (0.9108 :: Float)+  , mnistTestCase2VTI "VTI2 artificial 5 4 3 2 1" 5 4 3 2 1 5000+                       (0.8129 :: Double)+  , mnistTestCase2VTI "VTI2 1 epoch, 0 batch" 1 0 300 100 0.02 5000+                       (1 :: Float)+  ]++-- JAX differentiation, Ast term built and differentiated only once+-- and the result interpreted with different inputs in each gradient+-- descent iteration.+mnistTestCase2VTO+  :: forall r.+     ( Differentiable r, GoodScalar r+     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+  => String+  -> Int -> Int -> Int -> Int -> Double -> Int -> r+  -> TestTree+mnistTestCase2VTO prefix epochs maxBatches widthHidden widthHidden2+                  gamma batchSize expected =+  let (!targetInit, !artRaw) =+        MnistFcnnRanked2.mnistTrainBench2VTOGradient+          @r (Proxy @Float) IgnoreIncomingCotangent+          1 (mkStdGen 44) widthHidden widthHidden2+      !art = simplifyArtifactGradient artRaw+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show widthHidden, show widthHidden2+                        , show $ widthSTK $ knownSTK @(XParams2 r Float)+                        , show (tsize knownSTK targetInit)+                        , show gamma ]+  in testCase name $ do+    hPutStrLn stderr $+      printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+             prefix epochs maxBatches+    trainData <- loadMnistData trainGlyphsPath trainLabelsPath+    testData <- map mkMnistDataLinearR . take (batchSize * maxBatches)+                <$> loadMnistData testGlyphsPath testLabelsPath+    let go :: [MnistDataLinearR r] -> Concrete (XParams2 r Float)+           -> Concrete (XParams2 r Float)+        go [] parameters = parameters+        go ((glyph, label) : rest) !parameters =+          let parametersAndInput =+                tpair parameters (tpair (rconcrete glyph) (rconcrete label))+              gradient = tproject1 $ fst+                         $ revInterpretArtifact art parametersAndInput Nothing+          in go rest (updateWithGradient gamma knownSTK parameters gradient)+    let runBatch :: Concrete (XParams2 r Float) -> (Int, [MnistDataLinearR r])+                 -> IO (Concrete (XParams2 r Float))+        runBatch !params (k, chunk) = do+          let res = go chunk params+              trainScore =+                MnistFcnnRanked2.afcnnMnistTest2 chunk (fromTarget res)+              testScore =+                MnistFcnnRanked2.afcnnMnistTest2 testData (fromTarget res)+              lenChunk = length chunk+          unless (widthHidden < 10) $ do+            hPutStrLn stderr $+              printf "\n%s: (Batch %d with %d points)"+                     prefix k lenChunk+            hPutStrLn stderr $+              printf "%s: Training error:   %.2f%%"+                     prefix ((1 - trainScore) * 100)+            hPutStrLn stderr $+              printf "%s: Validation error: %.2f%%"+                     prefix ((1 - testScore ) * 100)+          return res+    let runEpoch :: Int -> Concrete (XParams2 r Float)+                 -> IO (Concrete (XParams2 r Float))+        runEpoch n params | n > epochs = return params+        runEpoch n !params = do+          unless (widthHidden < 10) $+            hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+          let trainDataShuffled = shuffle (mkStdGen $ n + 1) trainData+              chunks = take maxBatches+                       $ zip [1 ..] $ chunksOf batchSize+                       $ map mkMnistDataLinearR trainDataShuffled+          res <- foldM runBatch params chunks+          runEpoch (succ n) res+    res <- runEpoch 1 targetInit+    let testErrorFinal =+          1 - MnistFcnnRanked2.afcnnMnistTest2 testData (fromTarget res)+    testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCase2VTO+  :: String+  -> Int -> Int -> Int -> Int -> Double -> Int -> Double+  -> TestTree #-}++tensorADOnceMnistTests2 :: TestTree+tensorADOnceMnistTests2 = testGroup "Ranked2 Once MNIST tests"+  [ mnistTestCase2VTO "VTO2 1 epoch, 1 batch" 1 1 300 100 0.02 5000+                       (0.20779999999999998 :: Double)+  , mnistTestCase2VTO "VTO2 artificial 1 2 3 4 5" 1 2 3 4 5 5000+                       (0.9108 :: Float)+  , mnistTestCase2VTO "VTO2 artificial 5 4 3 2 1" 5 4 3 2 1 5000+                       (0.8129 :: Double)+  , mnistTestCase2VTO "VTO2 1 epoch, 0 batch" 1 0 300 100 0.02 5000+                       (1 :: Float)+  , testProperty "VTO2 grad vs fwd" $+    \seed0 ->+    forAllShrink (chooseInt (0, 600)) shrinkIntegral $ \width1Hidden ->+    forAllShrink (chooseInt (0, 200)) shrinkIntegral $ \width1Hidden2 ->+    forAllShrink (chooseInt (0, 5)) shrinkIntegral $ \simp ->+    forAll (choose (0.01, 1)) $ \range ->+    forAll (choose (0.01, 1)) $ \range2 ->+    forAll (choose (0.5, 1.5)) $ \dt ->+    forAll (choose (0, 1e-7)) $ \(perturbation :: Double) ->+    withSNat (1 + width1Hidden) $ \(SNat @widthHidden) ->+    withSNat (1 + width1Hidden2) $ \(SNat @widthHidden2) ->+    let (glyph0, seed2) = randomValue @(Concrete (TKS '[SizeMnistGlyph] Double))+                                      0.5 (mkStdGen seed0)+        (label0, seed3) = randomValue @(Concrete (TKS '[SizeMnistLabel] Double))+                                      5 seed2+        (glyph, label) = ( rmap1 (rscalar 0.5 +) $ forgetShape glyph0+                         , rmap1 (rscalar 5 + ) $ forgetShape label0 )+        ds :: Concrete (XParams2 Double Double)+        (ds, seed4) = first forgetShape $+          randomValue+            @(Concrete (X (MnistFcnnRanked2.ADFcnnMnist2ParametersShaped+                             Concrete widthHidden widthHidden2 Double Double)))+            range seed3+        (targetInit, artRaw) =+          MnistFcnnRanked2.mnistTrainBench2VTOGradient+            @Double (Proxy @Double) UseIncomingCotangent+            range2 seed4 (1 + width1Hidden) (1 + width1Hidden2)+        art = iterate simplifyArtifactGradient artRaw !! simp+        stk = knownSTK @(XParams2 Double Double)+        ftk = tftk @Concrete stk targetInit+        parametersAndInput = tpair targetInit (tpair glyph label)+        (_gradient0, value0) = first tproject1 $+          revInterpretArtifact art parametersAndInput Nothing+        (gradient1, value1) = first tproject1 $+          revInterpretArtifact art parametersAndInput (Just $ kconcrete dt)+        f :: ADVal Concrete (XParams2 Double Double)+          -> ADVal Concrete (TKScalar Double)+        f adinputs =+          MnistFcnnRanked2.afcnnMnistLoss2+            (rfromPrimal glyph, rfromPrimal label) (fromTarget adinputs)+        (derivative2, value2) = cfwdBoth f targetInit ds+--        goodDt :: forall r. GoodScalar r => r+--        goodDt = ifDifferentiable @r (realToFrac dt) 0+--        targetDt :: Concrete (XParams2 Double Double)+--        targetDt = replTarget goodDt ftk+        goodPerturbation :: forall r. GoodScalar r => r+        goodPerturbation = ifDifferentiable @r (realToFrac perturbation) 0+        targetPerturbed :: Concrete (XParams2 Double Double)+        targetPerturbed = treplTarget goodPerturbation ftk+        targetInitPerturbed :: Concrete (XParams2 Double Double)+        targetInitPerturbed = taddTarget stk targetInit targetPerturbed+        (derivative3, value3) = cfwdBoth f targetInit targetPerturbed+        value4 :: Concrete (TKScalar Double)+        value4 = MnistFcnnRanked2.afcnnMnistLoss2+                   (rfromPrimal glyph, rfromPrimal label)+                   (fromTarget targetInitPerturbed)+    in+      conjoin+        [ counterexample+            ("Objective function value from grad and jvp matches: "+             ++ show (value1, value2, value1 - value2))+            (abs (value1 - value2) < 1e-10)+        , counterexample+            ("Gradient and derivative agrees: "+             ++ show ( dt, derivative2, tdot0Target ftk gradient1 ds+                     , tdot0Target FTKScalar (kconcrete dt) derivative2+                       - tdot0Target ftk gradient1 ds ))+            (abs (tdot0Target FTKScalar (kconcrete dt) derivative2+                  - tdot0Target ftk gradient1 ds) < 1e-10)+--        , counterexample  -- this is implied by the other clauses+--            "Gradient is a linear function"+--            (gradient1 === tmultTarget stk targetDt gradient0)+        , counterexample+            "Objective function value unaffected by incoming cotangent"+            (value0 === value1)+        , counterexample+            "Objective function value unaffected by derivative perturbation"+            (value2 === value3)+        , counterexample+            ("Derivative approximates the perturbation of value: "+             ++ show ( value2, derivative3, value4+                     , (value3 + derivative3) - value4) )+            (abs ((value3 + derivative3) - value4) < 1e-6)+        ]+  ]
+ test/simplified/TestMnistPP.hs view
@@ -0,0 +1,726 @@+{-# LANGUAGE OverloadedLists #-}+-- | Tests of MNIST nns that pretty-print resulting gradient and primal terms.+module TestMnistPP+  ( testTrees+  ) where++import Prelude++import GHC.Exts (IsList (..))+import System.Random+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstInterpret+import HordeAd.Core.Ops (treplicate)++import MnistCnnRanked2 qualified+import MnistCnnShaped2 qualified+import MnistData+import MnistFcnnRanked1 qualified+import MnistFcnnRanked2 (XParams2)+import MnistFcnnRanked2 qualified+import MnistRnnRanked2 (ADRnnMnistParameters)+import MnistRnnRanked2 qualified++testTrees :: [TestTree]+testTrees = [ tensorMnistPPFCNNR+            , tensorMnistPPRNNR+            , tensorMnistCNNRPP+            ]++-- * FCNNR tests++type XParams widthHidden widthHidden2 r =+  X (MnistFcnnRanked1.ADFcnnMnist1Parameters+       Concrete widthHidden widthHidden2 r)++tensorMnistPPFCNNR :: TestTree+tensorMnistPPFCNNR = testGroup "PP and Ast tests for Short Ranked MNIST"+  [ testCase "VTO1 PP Lin" testVTOPP+  , testCase "VTO1 Ast Lin" testVTOAst+  , testCase "VTO1 PP NonLin" testVTOPPNonLin+  , testCase "VTO1 Ast NonLin" testVTOAstNonLin+  , testCase "VTO2 PP Lin" testVT2OPP+  , testCase "VTO2 Ast Lin" testVT2OAst+  , testCase "VTO2 PP NonLin" testVT2OPPNonLin+  , testCase "VTO2 PP NonLin2" testVT2OPPNonLin2+  , testCase "VTO2 Ast NonLin2" testVT2OAstNonLin2+  , testCase "VTO2 PP NonLin3" testVT2OPPNonLin3+  , testCase "VTO2 Ast NonLin3" testVT2OAstNonLin3+  ]++valsInitVTOPP :: (Num r, Enum r, Nested.PrimElt r)+              => MnistFcnnRanked1.ADFcnnMnist1Parameters Concrete 3 4 r+valsInitVTOPP =+  ( ( fromList (replicate 3 (Concrete+                             $ Nested.sfromListPrim+                                 (SNat @SizeMnistGlyph)+                                 [1 .. fromIntegral sizeMnistGlyphInt]))+    , Concrete $ Nested.sfromListPrim (SNat @3) [1, 2, 3] )+  , ( fromList (replicate 4 (Concrete $ Nested.sfromListPrim+                                          (SNat @3) [1, 2, 3]))+    , Concrete $ Nested.sfromListPrim (SNat @4) [1, 2, 3, 4] )+  , ( fromList (replicate sizeMnistLabelInt+                          (Concrete $ Nested.sfromListPrim+                                        (SNat @4) [1, 2, 3, 4]))+    , Concrete $ Nested.sfromListPrim (SNat @SizeMnistLabel)+                                      [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] ) )++testVTOPP :: Assertion+testVTOPP = do+  resetVarCounter+  let blackGlyph = treplicate (SNat @SizeMnistGlyph) knownSTK+                   $ fromPrimal $ rconcrete $ Nested.rscalar 7+      afcnn2T :: MnistFcnnRanked1.ADFcnnMnist1Parameters+                   (AstTensor AstMethodLet FullSpan) 3 4 Float+              -> AstTensor AstMethodLet FullSpan (TKR 1 Float)+      afcnn2T =+        MnistFcnnRanked1.afcnnMnist1 id id+                                     (SNat @3) (SNat @4) (sfromR blackGlyph)+      ftk = tftk @Concrete (knownSTK @(XParams 3 4 Float))+                           (toTarget @Concrete valsInitVTOPP)+      artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2T ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\v1 -> rfromS (let v4 = sfromVector (fromList [sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject1 (tproject1 (tproject1 v1)))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v4, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v4]) + tproject2 (tproject2 (tproject1 v1)) in sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 v1))) v5, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 v1)))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) v5, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))))) v5]) + tproject2 (tproject2 v1))"+  printArtifactPrimalPretty artifactRev+    @?= "\\v1 -> let v4 = sfromVector (fromList [ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject1 (tproject1 (tproject1 v1)))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector (fromList [ssum @3 (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v4), ssum @3 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v4), ssum @3 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v4), ssum @3 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v4)]) + tproject2 (tproject2 (tproject1 v1)) in rfromS (sfromVector (fromList [ssum @4 (tproject1 (tproject1 (tproject2 v1)) * v5), ssum @4 (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v5), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v5)]) + tproject2 (tproject2 v1))"+  printArtifactPretty artifactRev+    @?= "\\dret v1 -> let v4 = sfromVector (fromList [ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject1 (tproject1 (tproject1 v1)))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector (fromList [ssum @3 (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v4), ssum @3 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v4), ssum @3 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v4), ssum @3 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v4)]) + tproject2 (tproject2 (tproject1 v1)) ; v7 = sreplicate @4 (sfromR dret !$ [9]) ; v8 = sreplicate @4 (sfromR dret !$ [8]) ; v9 = sreplicate @4 (sfromR dret !$ [7]) ; v10 = sreplicate @4 (sfromR dret !$ [6]) ; v11 = sreplicate @4 (sfromR dret !$ [5]) ; v12 = sreplicate @4 (sfromR dret !$ [4]) ; v13 = sreplicate @4 (sfromR dret !$ [3]) ; v14 = sreplicate @4 (sfromR dret !$ [2]) ; v15 = sreplicate @4 (sfromR dret !$ [1]) ; v16 = sreplicate @4 (sfromR dret !$ [0]) ; v17 = tproject1 (tproject1 (tproject2 v1)) * v16 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v15 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v14 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v13 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v12 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v11 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v10 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v9 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v8 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v7)))))))) ; v18 = sreplicate @3 (v17 !$ [3]) ; v19 = sreplicate @3 (v17 !$ [2]) ; v20 = sreplicate @3 (v17 !$ [1]) ; v21 = sreplicate @3 (v17 !$ [0]) ; v22 = tproject1 (tproject1 (tproject2 (tproject1 v1))) * v21 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v20 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v19 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v18)) in tpair (tpair (tpair (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [0])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [1])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [2])) Z1))) v22) (tpair (tpair (v4 * v21) (tpair (v4 * v20) (tpair (v4 * v19) (tpair (v4 * v18) Z1)))) v17)) (tpair (tpair (v5 * v16) (tpair (v5 * v15) (tpair (v5 * v14) (tpair (v5 * v13) (tpair (v5 * v12) (tpair (v5 * v11) (tpair (v5 * v10) (tpair (v5 * v9) (tpair (v5 * v8) (tpair (v5 * v7) Z1)))))))))) (sfromR dret))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret v1 -> let v4 = sfromVector (fromList [sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject1 (tproject1 (tproject1 v1)))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v5 = sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v4, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v4, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v4]) + tproject2 (tproject2 (tproject1 v1)) ; x7 = sfromR dret !$ [9] ; x8 = sfromR dret !$ [8] ; x9 = sfromR dret !$ [7] ; x10 = sfromR dret !$ [6] ; x11 = sfromR dret !$ [5] ; x12 = sfromR dret !$ [4] ; x13 = sfromR dret !$ [3] ; x14 = sfromR dret !$ [2] ; x15 = sfromR dret !$ [1] ; x16 = sfromR dret !$ [0] ; v17 = tproject1 (tproject1 (tproject2 v1)) * sreplicate @4 x16 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * sreplicate @4 x15 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * sreplicate @4 x14 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * sreplicate @4 x13 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * sreplicate @4 x12 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * sreplicate @4 x11 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * sreplicate @4 x10 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * sreplicate @4 x9 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * sreplicate @4 x8 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * sreplicate @4 x7)))))))) ; x18 = v17 !$ [3] ; x19 = v17 !$ [2] ; x20 = v17 !$ [1] ; x21 = v17 !$ [0] ; v22 = tproject1 (tproject1 (tproject2 (tproject1 v1))) * sreplicate @3 x21 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * sreplicate @3 x20 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * sreplicate @3 x19 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * sreplicate @3 x18)) in tpair (tpair (tpair (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [0])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [1])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v22 !$ [2])) Z1))) v22) (tpair (tpair (v4 * sreplicate @3 x21) (tpair (v4 * sreplicate @3 x20) (tpair (v4 * sreplicate @3 x19) (tpair (v4 * sreplicate @3 x18) Z1)))) v17)) (tpair (tpair (v5 * sreplicate @4 x16) (tpair (v5 * sreplicate @4 x15) (tpair (v5 * sreplicate @4 x14) (tpair (v5 * sreplicate @4 x13) (tpair (v5 * sreplicate @4 x12) (tpair (v5 * sreplicate @4 x11) (tpair (v5 * sreplicate @4 x10) (tpair (v5 * sreplicate @4 x9) (tpair (v5 * sreplicate @4 x8) (tpair (v5 * sreplicate @4 x7) Z1)))))))))) (sfromR dret))"++testVTOAst :: Assertion+testVTOAst = do+  let ftk = tftk @Concrete (knownSTK @(XParams 3 4 Float))+                 (toTarget @Concrete valsInitVTOPP)+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var :: AstTensor AstMethodLet FullSpan (XParams 3 4 Float)+      var = AstVar varName+      vals = toTarget @Concrete valsInitVTOPP+      env = extendEnv varName vals emptyEnv+      blackGlyph = treplicate (SNat @SizeMnistGlyph) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => MnistFcnnRanked1.ADFcnnMnist1Parameters f 3 4 Float+             -> f (TKR 1 Float)+      afcnn2 = MnistFcnnRanked1.afcnnMnist1+                 id id (SNat @3) (SNat @4)+                 (sfromR $ rconcrete $ unConcrete blackGlyph)+      afcnn1 = afcnn2 $ fromTarget var+  interpretAstFull @Concrete env afcnn1+    @?= afcnn2 valsInitVTOPP+  interpretAstFull @Concrete env+                   (simplifyInline @(TKR 1 Float) afcnn1)+    @?= afcnn2 valsInitVTOPP+  interpretAstFull @Concrete env+                   (simplifyInlineContract @(TKR 1 Float) afcnn1)+    @?= afcnn2 valsInitVTOPP++testVTOPPNonLin :: Assertion+testVTOPPNonLin = do+  resetVarCounter+  let blackGlyph = treplicate (SNat @SizeMnistGlyph) knownSTK+                   $ fromPrimal $ rconcrete $ Nested.rscalar 7+      afcnn2TnonLin :: MnistFcnnRanked1.ADFcnnMnist1Parameters+                         (AstTensor AstMethodLet FullSpan) 3 4 Double+                    -> AstTensor AstMethodLet FullSpan (TKR 1 Double)+      afcnn2TnonLin =+        MnistFcnnRanked1.afcnnMnist1 logisticS softMax1S+                                     (SNat @3) (SNat @4) (sfromR blackGlyph)+      ftk = tftk @Concrete (knownSTK @(XParams 3 4 Double))+                           (toTarget @Concrete valsInitVTOPP)+      artifactRevnonLin =+        revArtifactAdapt UseIncomingCotangent afcnn2TnonLin ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRevnonLin)+    @?= "\\v1 -> rfromS (let v15 = scast (recip (sconcrete (sreplicate [3] 1.0) + exp (negate (sfromVector (fromList [sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject1 (tproject1 (tproject1 v1)))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))])) + negate (tproject2 (tproject1 (tproject1 v1)))))) ; v19 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (scast (sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v15, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v15, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v15, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v15]))) + negate (tproject2 (tproject2 (tproject1 v1))))) ; v22 = exp (sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 v1))) v19, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 v1)))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))))) v19]) + tproject2 (tproject2 v1)) in sreplicate @10 (recip (ssum0 v22)) * v22)"+  printArtifactPrimalPretty artifactRevnonLin+    @?= "\\v1 -> let v9 = sfromVector (fromList [ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject1 (tproject1 (tproject1 v1)))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v10 = exp (negate v9) ; v11 = sconcrete (sreplicate [3] 1.0) + v10 ; v12 = recip v11 ; v15 = scast v12 ; v16 = scast (sfromVector (fromList [ssum @3 (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v15), ssum @3 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v15), ssum @3 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v15), ssum @3 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v15)])) + tproject2 (tproject2 (tproject1 v1)) ; v17 = exp (negate v16) ; v18 = sconcrete (sreplicate [4] 1.0) + v17 ; v19 = recip v18 ; v22 = exp (sfromVector (fromList [ssum @4 (tproject1 (tproject1 (tproject2 v1)) * v19), ssum @4 (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v19)]) + tproject2 (tproject2 v1)) ; x23 = ssum @10 v22 ; v24 = sreplicate @10 (recip x23) in rfromS (v24 * v22)"+  printArtifactPretty artifactRevnonLin+    @?= "\\dret v1 -> let v9 = sfromVector (fromList [ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject1 (tproject1 (tproject1 v1)))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), ssum @784 (sconcrete (sreplicate [784] 7.0) * tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))]) + tproject2 (tproject1 (tproject1 v1)) ; v10 = exp (negate v9) ; v11 = sconcrete (sreplicate [3] 1.0) + v10 ; v12 = recip v11 ; v13 = sconcrete (sreplicate [3] 1.0) + negate v12 ; v14 = v12 * v13 ; v15 = scast v12 ; v16 = scast (sfromVector (fromList [ssum @3 (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v15), ssum @3 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v15), ssum @3 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v15), ssum @3 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v15)])) + tproject2 (tproject2 (tproject1 v1)) ; v17 = exp (negate v16) ; v18 = sconcrete (sreplicate [4] 1.0) + v17 ; v19 = recip v18 ; v20 = sconcrete (sreplicate [4] 1.0) + negate v19 ; v21 = v19 * v20 ; v22 = exp (sfromVector (fromList [ssum @4 (tproject1 (tproject1 (tproject2 v1)) * v19), ssum @4 (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v19), ssum @4 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v19)]) + tproject2 (tproject2 v1)) ; x23 = ssum @10 v22 ; v24 = sreplicate @10 (recip x23) ; v26 = v22 * (sreplicate @10 (negate (recip (x23 * x23)) * ssum @10 (v22 * sfromR dret)) + v24 * sfromR dret) ; v27 = sreplicate @4 (v26 !$ [9]) ; v28 = sreplicate @4 (v26 !$ [8]) ; v29 = sreplicate @4 (v26 !$ [7]) ; v30 = sreplicate @4 (v26 !$ [6]) ; v31 = sreplicate @4 (v26 !$ [5]) ; v32 = sreplicate @4 (v26 !$ [4]) ; v33 = sreplicate @4 (v26 !$ [3]) ; v34 = sreplicate @4 (v26 !$ [2]) ; v35 = sreplicate @4 (v26 !$ [1]) ; v36 = sreplicate @4 (v26 !$ [0]) ; v37 = v21 * (tproject1 (tproject1 (tproject2 v1)) * v36 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * v35 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * v34 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * v33 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * v32 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * v31 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * v30 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * v29 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * v28 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * v27))))))))) ; v38 = scast v37 ; v39 = sreplicate @3 (v38 !$ [3]) ; v40 = sreplicate @3 (v38 !$ [2]) ; v41 = sreplicate @3 (v38 !$ [1]) ; v42 = sreplicate @3 (v38 !$ [0]) ; v43 = v14 * scast (tproject1 (tproject1 (tproject2 (tproject1 v1))) * v42 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * v41 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * v40 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * v39))) in tpair (tpair (tpair (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [0])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [1])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [2])) Z1))) v43) (tpair (tpair (v15 * v42) (tpair (v15 * v41) (tpair (v15 * v40) (tpair (v15 * v39) Z1)))) v37)) (tpair (tpair (v19 * v36) (tpair (v19 * v35) (tpair (v19 * v34) (tpair (v19 * v33) (tpair (v19 * v32) (tpair (v19 * v31) (tpair (v19 * v30) (tpair (v19 * v29) (tpair (v19 * v28) (tpair (v19 * v27) Z1)))))))))) v26)"+  printArtifactPretty (simplifyArtifact artifactRevnonLin)+    @?= "\\dret v1 -> let v12 = recip (sconcrete (sreplicate [3] 1.0) + exp (negate (sfromVector (fromList [sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject1 (tproject1 (tproject1 v1)))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject1 (tproject1 (tproject1 v1))))), sdot0 (sconcrete (sreplicate [784] 7.0)) (tproject1 (tproject2 (tproject2 (tproject1 (tproject1 (tproject1 v1))))))])) + negate (tproject2 (tproject1 (tproject1 v1))))) ; v15 = scast v12 ; v19 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (scast (sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 (tproject1 v1)))) v15, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) v15, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) v15, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))))) v15]))) + negate (tproject2 (tproject2 (tproject1 v1))))) ; v22 = exp (sfromVector (fromList [sdot0 (tproject1 (tproject1 (tproject2 v1))) v19, sdot0 (tproject1 (tproject2 (tproject1 (tproject2 v1)))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) v19, sdot0 (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))))) v19]) + tproject2 (tproject2 v1)) ; x23 = ssum0 v22 ; v26 = v22 * (sreplicate @10 (negate (recip (x23 * x23)) * sdot0 v22 (sfromR dret)) + sreplicate @10 (recip x23) * sfromR dret) ; x27 = v26 !$ [9] ; x28 = v26 !$ [8] ; x29 = v26 !$ [7] ; x30 = v26 !$ [6] ; x31 = v26 !$ [5] ; x32 = v26 !$ [4] ; x33 = v26 !$ [3] ; x34 = v26 !$ [2] ; x35 = v26 !$ [1] ; x36 = v26 !$ [0] ; v37 = (v19 * (sconcrete (sreplicate [4] 1.0) + negate v19)) * (tproject1 (tproject1 (tproject2 v1)) * sreplicate @4 x36 + (tproject1 (tproject2 (tproject1 (tproject2 v1))) * sreplicate @4 x35 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 v1)))) * sreplicate @4 x34 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))) * sreplicate @4 x33 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))) * sreplicate @4 x32 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))) * sreplicate @4 x31 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))) * sreplicate @4 x30 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))) * sreplicate @4 x29 + (tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1)))))))))) * sreplicate @4 x28 + tproject1 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 v1))))))))))) * sreplicate @4 x27))))))))) ; v38 = scast v37 ; x39 = v38 !$ [3] ; x40 = v38 !$ [2] ; x41 = v38 !$ [1] ; x42 = v38 !$ [0] ; v43 = (v12 * (sconcrete (sreplicate [3] 1.0) + negate v12)) * scast (tproject1 (tproject1 (tproject2 (tproject1 v1))) * sreplicate @3 x42 + (tproject1 (tproject2 (tproject1 (tproject2 (tproject1 v1)))) * sreplicate @3 x41 + (tproject1 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1))))) * sreplicate @3 x40 + tproject1 (tproject2 (tproject2 (tproject2 (tproject1 (tproject2 (tproject1 v1)))))) * sreplicate @3 x39))) in tpair (tpair (tpair (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [0])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [1])) (tpair (sconcrete (sreplicate [784] 7.0) * sreplicate @784 (v43 !$ [2])) Z1))) v43) (tpair (tpair (v15 * sreplicate @3 x42) (tpair (v15 * sreplicate @3 x41) (tpair (v15 * sreplicate @3 x40) (tpair (v15 * sreplicate @3 x39) Z1)))) v37)) (tpair (tpair (v19 * sreplicate @4 x36) (tpair (v19 * sreplicate @4 x35) (tpair (v19 * sreplicate @4 x34) (tpair (v19 * sreplicate @4 x33) (tpair (v19 * sreplicate @4 x32) (tpair (v19 * sreplicate @4 x31) (tpair (v19 * sreplicate @4 x30) (tpair (v19 * sreplicate @4 x29) (tpair (v19 * sreplicate @4 x28) (tpair (v19 * sreplicate @4 x27) Z1)))))))))) v26)"++testVTOAstNonLin :: Assertion+testVTOAstNonLin = do+  let ftk = tftk @Concrete (knownSTK @(XParams 3 4 Double))+                 (toTarget @Concrete valsInitVTOPP)+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var :: AstTensor AstMethodLet FullSpan (XParams 3 4 Double)+      var = AstVar varName+      vals = toTarget @Concrete valsInitVTOPP+      env = extendEnv varName vals emptyEnv+      blackGlyph = treplicate (SNat @SizeMnistGlyph) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => MnistFcnnRanked1.ADFcnnMnist1Parameters f 3 4 Double+             -> f (TKR 1 Double)+      afcnn2 = MnistFcnnRanked1.afcnnMnist1+                 logisticS softMax1S (SNat @3) (SNat @4)+                 (sfromR $ rconcrete $ unConcrete blackGlyph)+      afcnn1 = afcnn2 $ fromTarget var+  interpretAstFull @Concrete env afcnn1+    @?= afcnn2 valsInitVTOPP+  interpretAstFull @Concrete env+                   (simplifyInline @(TKR 1 Double) afcnn1)+    @?= afcnn2 valsInitVTOPP+  interpretAstFull @Concrete env+                   (simplifyInlineContract @(TKR 1 Double) afcnn1)+    @?= afcnn2 valsInitVTOPP++valsInitVT2OPP :: MnistFcnnRanked2.ADFcnnMnist2Parameters Concrete Double Float+valsInitVT2OPP =+  ( ( Concrete $ Nested.rfromListPrimLinear [4, 3]+               (concat $ replicate 4 [1, 2, 3])+    , Concrete $ Nested.rfromListPrimLinear [4] [1, 2, 3, 4] )+  , ( Concrete $ Nested.rfromListPrimLinear [5, 4]+               (concat $ replicate 5 [1, 2, 3, 4])+    , Concrete $ Nested.rfromListPrimLinear [5] [1, 2, 3, 4, 5] )+  , ( Concrete $ Nested.rfromListPrimLinear [2, 5]+               (concat $ replicate 2 [1, 2, 3, 4, 5])+    , Concrete $ Nested.rfromListPrimLinear [2] [1, 2] ) )++testVT2OPP :: Assertion+testVT2OPP = do+  resetVarCounter+  let blackGlyph = treplicate (SNat @3) knownSTK+                   $ fromPrimal $ rconcrete $ Nested.rscalar 7+      afcnn2T :: MnistFcnnRanked2.ADFcnnMnist2Parameters+                   (AstTensor AstMethodLet FullSpan) Double Float+              -> AstTensor AstMethodLet FullSpan (TKR 1 Double)+      afcnn2T = MnistFcnnRanked2.afcnnMnist2 id id blackGlyph+      ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))+                           (toTarget @Concrete valsInitVT2OPP)+      artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2T ftk+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> let m5 = str (sreplicate @5 (scast (ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m6 = str (sreplicate @2 (scast (ssum @4 (m5 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))))) in rfromS (ssum @5 (m6 * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1)))"+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let m5 = str (sreplicate @5 (scast (ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m6 = str (sreplicate @2 (scast (ssum @4 (m5 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))))) ; v8 = ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 (sfromR dret))) ; m9 = sreplicate @4 (scast v8) ; v10 = scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m9))) in tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [3,4] 7.0) * sreplicate @3 v10))) (rfromS v10)) (tpair (rfromS (str (m5 * m9))) (rfromS v8))) (tpair (rfromS (str (m6 * sreplicate @5 (sfromR dret)))) dret)"+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let m5 = str (sreplicate @5 (scast (ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; m6 = str (sreplicate @2 (scast (ssum @4 (m5 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))))) ; v8 = ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 (sfromR dret))) ; m9 = sreplicate @4 (scast v8) ; v10 = scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m9))) in tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [3,4] 7.0) * sreplicate @3 v10))) (rfromS v10)) (tpair (rfromS (str (m5 * m9))) (rfromS v8))) (tpair (rfromS (str (m6 * sreplicate @5 (sfromR dret)))) dret)"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4] FTKScalar)) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,5] FTKScalar)) ConvSX)) ConvId)) (STKProduct (STKProduct (STKProduct (STKS [4,3] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [5,4] STKScalar) (STKS [5] STKScalar))) (STKProduct (STKS [2,5] STKScalar) (STKR (SNat @1) STKScalar))) (let v5 = scast (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1)))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; v8 = sdot1In (str (sfromR (tproject1 (tproject2 m1)))) (sreplicate @5 (sfromR dret)) ; v9 = scast v8 ; v10 = scast (sdot1In (str (sfromR (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @4 v9)) in tpair (tpair (tpair (sconcrete (sreplicate [4,3] 7.0) * str (sreplicate @3 v10)) v10) (tpair (sreplicate @5 v5 * str (sreplicate @4 v9)) v8)) (tpair (sreplicate @2 (scast (sdot1In (sreplicate @5 v5) (sfromR (tproject1 (tproject2 (tproject1 m1))))) + sfromR (tproject2 (tproject2 (tproject1 m1)))) * str (sreplicate @5 (sfromR dret))) dret))"++testVT2OAst :: Assertion+testVT2OAst = do+  let ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))+                 (toTarget @Concrete valsInitVT2OPP)+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var :: AstTensor AstMethodLet FullSpan (XParams2 Double Float)+      var = AstVar varName+      vals = toTarget @Concrete valsInitVT2OPP+      env = extendEnv varName vals emptyEnv+      blackGlyph = treplicate (SNat @3) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => MnistFcnnRanked2.ADFcnnMnist2Parameters f Double Float+             -> f (TKR 1 Double)+      afcnn2 = MnistFcnnRanked2.afcnnMnist2+                 id id+                 (rconcrete $ unConcrete blackGlyph)+      afcnn1 = afcnn2 $ fromTarget var+  interpretAstFull @Concrete env afcnn1+    @?= afcnn2 valsInitVT2OPP+  interpretAstFull @Concrete env+                   (simplifyInline @(TKR 1 Double) afcnn1)+    @?= afcnn2 valsInitVT2OPP+  interpretAstFull @Concrete env+                   (simplifyInlineContract @(TKR 1 Double) afcnn1)+    @?= afcnn2 valsInitVT2OPP++testVT2OPPNonLin :: Assertion+testVT2OPPNonLin = do+  resetVarCounter+  let blackGlyph = treplicate (SNat @3) knownSTK+                   $ fromPrimal $ rconcrete $ Nested.rscalar 7+      afcnn2TnonLin :: MnistFcnnRanked2.ADFcnnMnist2Parameters+                         (AstTensor AstMethodLet FullSpan) Float Float+                    -> AstTensor AstMethodLet FullSpan (TKR 1 Float)+      afcnn2TnonLin = MnistFcnnRanked2.afcnnMnist2 logistic softMax1 blackGlyph+      constant =+        let ((a1, a2), (a3, a4), (a5, a6)) = valsInitVT2OPP+        in ( ( rcast $ fromPrimal $ rconcrete $ unConcrete a1+             , rcast $ fromPrimal $ rconcrete $ unConcrete a2 )+           , ( fromPrimal $ rcast $ rconcrete $ unConcrete a3+             , fromPrimal $ rcast $ rconcrete $ unConcrete a4 )+           , ( rcast $ fromPrimal $ rconcrete $ unConcrete a5+             , fromPrimal $ rcast $ rconcrete $ unConcrete a6 ) )+      ast3 = fun1ToAst (FTKR (0 :$: ZSR) (FTKScalar @Float))+                       (const $ afcnn2TnonLin constant)+  "\\dummy" ++ " -> " ++ printAstSimple ast3+    @?= "\\dummy -> rfromS (tlet (exp (ssum @5 (str (sreplicate @2 (tlet (ssum @4 (str (sreplicate @5 (tlet (tfromPrimal (STKS [4] STKScalar) (ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (scast (sconcrete (sfromListLinear [4,3] [1.0,2.0,3.0,1.0,2.0,3.0,1.0,2.0,3.0,1.0,2.0,3.0])))) + scast (sconcrete (sfromListLinear [4] [1.0,2.0,3.0,4.0])))) (\\v5 -> ttletPrimal (recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sfromR (tprimalPart (rfromS v5)))))) (\\v6 -> tfromPrimal (STKS [4] STKScalar) v6 + sfromR (tfromDual (tdualPart (STKR (SNat @1) STKScalar) (rfromS (tfromPrimal (STKS [4] STKScalar) (v6 * (sconcrete (sreplicate [4] 1.0) + negate v6)) * sfromR (tfromDual (tdualPart (STKR (SNat @1) STKScalar) (rfromS v5))))))))))) * tfromPrimal (STKS [4,5] STKScalar) (sconcrete (sfromListLinear [4,5] [1.0,1.0,1.0,1.0,1.0,2.0,2.0,2.0,2.0,2.0,3.0,3.0,3.0,3.0,3.0,4.0,4.0,4.0,4.0,4.0]))) + tfromPrimal (STKS [5] STKScalar) (scast (sconcrete (sfromListLinear [5] [1.0,2.0,3.0,4.0,5.0])))) (\\v7 -> ttletPrimal (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (sfromR (tprimalPart (rfromS v7)))))) (\\v8 -> tfromPrimal (STKS [5] STKScalar) v8 + sfromR (tfromDual (tdualPart (STKR (SNat @1) STKScalar) (rfromS (tfromPrimal (STKS [5] STKScalar) (v8 * (sconcrete (sreplicate [5] 1.0) + negate v8)) * sfromR (tfromDual (tdualPart (STKR (SNat @1) STKScalar) (rfromS v7))))))))))) * tfromPrimal (STKS [5,2] STKScalar) (str (scast (sconcrete (sfromListLinear [2,5] [1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0]))))) + tfromPrimal (STKS [2] STKScalar) (scast (sconcrete (sfromListLinear [2] [1.0,2.0]))))) (\\v9 -> sreplicate @2 (recip (ssum @2 v9)) * v9))"+  "\\dummy" ++ " -> " ++ printAstSimple (simplifyInlineContract ast3)+    @?= "\\dummy -> rfromS (tlet (exp (sdot1In (sreplicate @2 (tlet (sdot1In (sreplicate @5 (ttletPrimal (recip (sconcrete (sreplicate [4] 1.0) + exp (sconcrete (sfromListLinear [4] [-43.0,-44.0,-45.0,-46.0])))) (\\v6 -> tfromPrimal (STKS [4] STKScalar) v6 + tfromDual (tdualPart (STKS [4] STKScalar) (tfromPrimal (STKS [4] STKScalar) (v6 * (sconcrete (sreplicate [4] 1.0) + negate v6)) * tfromDual (tdualPart (STKS [4] STKScalar) (tfromPrimal (STKS [4] STKScalar) (sconcrete (sreplicate [4] 0.0))))))))) (tfromPrimal (STKS [5,4] STKScalar) (sconcrete (sfromListLinear [5,4] [1.0,2.0,3.0,4.0,1.0,2.0,3.0,4.0,1.0,2.0,3.0,4.0,1.0,2.0,3.0,4.0,1.0,2.0,3.0,4.0]))) + tfromPrimal (STKS [5] STKScalar) (sconcrete (sfromListLinear [5] [1.0,2.0,3.0,4.0,5.0]))) (\\v7 -> ttletPrimal (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (tprimalPart v7)))) (\\v8 -> tfromPrimal (STKS [5] STKScalar) v8 + tfromDual (tdualPart (STKS [5] STKScalar) (tfromPrimal (STKS [5] STKScalar) (v8 * (sconcrete (sreplicate [5] 1.0) + negate v8)) * tfromDual (tdualPart (STKS [5] STKScalar) v7))))))) (tfromPrimal (STKS [2,5] STKScalar) (sconcrete (sfromListLinear [2,5] [1.0,2.0,3.0,4.0,5.0,1.0,2.0,3.0,4.0,5.0]))) + tfromPrimal (STKS [2] STKScalar) (sconcrete (sfromListLinear [2] [1.0,2.0])))) (\\v9 -> sreplicate @2 (recip (ssum0 v9)) * v9))"++testVT2OPPNonLin2 :: Assertion+testVT2OPPNonLin2 = do+  resetVarCounter+  let blackGlyph = treplicate (SNat @3) knownSTK+                   $ fromPrimal $ rconcrete $ Nested.rscalar 7+      afcnn2TnonLin :: MnistFcnnRanked2.ADFcnnMnist2Parameters+                         (AstTensor AstMethodLet FullSpan) Double Float+                    -> AstTensor AstMethodLet FullSpan (TKR 1 Double)+      afcnn2TnonLin = MnistFcnnRanked2.afcnnMnist2 logistic softMax1 blackGlyph+      ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))+                           (toTarget @Concrete valsInitVT2OPP)+      artifactRevnonLin =+        revArtifactAdapt UseIncomingCotangent afcnn2TnonLin ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRevnonLin)+    @?= "\\m1 -> rfromS (let v23 = exp (sdot1In (sreplicate @2 (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 (scast (recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1))))))))) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))))) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) in sreplicate @2 (recip (ssum0 v23)) * v23)"+  printArtifactPrimalPretty artifactRevnonLin+    @?= "\\m1 -> let v10 = ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v11 = exp (negate v10) ; v12 = sconcrete (sreplicate [4] 1.0) + v11 ; v13 = recip v12 ; m16 = str (sreplicate @5 (scast v13)) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = sconcrete (sreplicate [5] 1.0) + v18 ; v20 = recip v19 ; v23 = exp (ssum @5 (str (sreplicate @2 v20) * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum @2 v23 ; v25 = sreplicate @2 (recip x24) in rfromS (v25 * v23)"+  printArtifactPretty artifactRevnonLin+    @?= "\\dret m1 -> let v10 = ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v11 = exp (negate v10) ; v12 = sconcrete (sreplicate [4] 1.0) + v11 ; v13 = recip v12 ; v14 = sconcrete (sreplicate [4] 1.0) + negate v13 ; v15 = v13 * v14 ; m16 = str (sreplicate @5 (scast v13)) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = sconcrete (sreplicate [5] 1.0) + v18 ; v20 = recip v19 ; v21 = sconcrete (sreplicate [5] 1.0) + negate v20 ; v22 = v20 * v21 ; v23 = exp (ssum @5 (str (sreplicate @2 v20) * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum @2 v23 ; v25 = sreplicate @2 (recip x24) ; v27 = v23 * (sreplicate @2 (negate (recip (x24 * x24)) * ssum @2 (v23 * sfromR dret)) + v25 * sfromR dret) ; v28 = v22 * ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 v27)) ; m29 = sreplicate @4 (scast v28) ; v30 = v15 * scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m29))) in tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [3,4] 7.0) * sreplicate @3 v30))) (rfromS v30)) (tpair (rfromS (str (m16 * m29))) (rfromS v28))) (tpair (rfromS (str (str (sreplicate @2 v20) * sreplicate @5 v27))) (rfromS v27))"+  printArtifactPretty (simplifyArtifact artifactRevnonLin)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4] FTKScalar)) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2] FTKScalar)) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKS [4,3] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [5,4] STKScalar) (STKS [5] STKScalar))) (STKProduct (STKS [2,5] STKScalar) (STKS [2] STKScalar))) (let v13 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v16 = scast v13 ; v20 = recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 v16) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; v23 = exp (sdot1In (sreplicate @2 v20) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum0 v23 ; v27 = v23 * (sreplicate @2 (negate (recip (x24 * x24)) * sdot0 v23 (sfromR dret)) + sreplicate @2 (recip x24) * sfromR dret) ; v28 = (v20 * (sconcrete (sreplicate [5] 1.0) + negate v20)) * sdot1In (str (sfromR (tproject1 (tproject2 m1)))) (sreplicate @5 v27) ; v29 = scast v28 ; v30 = (v13 * (sconcrete (sreplicate [4] 1.0) + negate v13)) * scast (sdot1In (str (sfromR (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @4 v29)) in tpair (tpair (tpair (sconcrete (sreplicate [4,3] 7.0) * str (sreplicate @3 v30)) v30) (tpair (sreplicate @5 v16 * str (sreplicate @4 v29)) v28)) (tpair (sreplicate @2 v20 * str (sreplicate @5 v27)) v27))"++testVT2OAstNonLin2 :: Assertion+testVT2OAstNonLin2 = do+  let ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))+                 (toTarget @Concrete valsInitVT2OPP)+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var :: AstTensor AstMethodLet FullSpan (XParams2 Double Float)+      var = AstVar varName+      vals = toTarget @Concrete valsInitVT2OPP+      env = extendEnv varName vals emptyEnv+      blackGlyph = treplicate (SNat @3) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => MnistFcnnRanked2.ADFcnnMnist2Parameters f Double Float+             -> f (TKR 1 Double)+      afcnn2 = MnistFcnnRanked2.afcnnMnist2+                 logistic softMax1+                 (rconcrete $ unConcrete blackGlyph)+      afcnn1 = afcnn2 $ fromTarget var+  interpretAstFull @Concrete env afcnn1+    @?= afcnn2 valsInitVT2OPP+  interpretAstFull @Concrete env+                   (simplifyInline @(TKR 1 Double) afcnn1)+    @?= afcnn2 valsInitVT2OPP+  interpretAstFull @Concrete env+                   (simplifyInlineContract @(TKR 1 Double) afcnn1)+    @?= afcnn2 valsInitVT2OPP++testVT2OPPNonLin3 :: Assertion+testVT2OPPNonLin3 = do+  resetVarCounter+  let blackGlyph = treplicate (SNat @3) knownSTK+                   $ fromPrimal $ rconcrete $ Nested.rscalar 7+      blackLabel = treplicate (SNat @2) knownSTK+                   $ fromPrimal $ rconcrete $ Nested.rscalar 8+      afcnn2TnonLin :: MnistFcnnRanked2.ADFcnnMnist2Parameters+                         (AstTensor AstMethodLet FullSpan) Double Float+                    -> AstTensor AstMethodLet FullSpan (TKScalar Double)+      afcnn2TnonLin = MnistFcnnRanked2.afcnnMnistLoss2 (blackGlyph, blackLabel)+      ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))+                           (toTarget @Concrete valsInitVT2OPP)+      artifactRevnonLin =+        revArtifactAdapt UseIncomingCotangent afcnn2TnonLin ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRevnonLin)+    @?= "\\m1 -> let v23 = exp (sdot1In (sreplicate @2 (recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 (scast (recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1))))))))) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))))) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) in kfromS (negate (sdot0 (sconcrete (sreplicate [2] 8.0)) (log (sreplicate @2 (recip (ssum0 v23)) * v23))))"+  printArtifactPrimalPretty artifactRevnonLin+    @?= "\\m1 -> let v10 = ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v11 = exp (negate v10) ; v12 = sconcrete (sreplicate [4] 1.0) + v11 ; v13 = recip v12 ; m16 = str (sreplicate @5 (scast v13)) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = sconcrete (sreplicate [5] 1.0) + v18 ; v20 = recip v19 ; v23 = exp (ssum @5 (str (sreplicate @2 v20) * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum @2 v23 ; v25 = sreplicate @2 (recip x24) ; v26 = v25 * v23 ; v27 = log v26 in kfromS (negate (ssum @2 (sconcrete (sreplicate [2] 8.0) * v27)))"+  printArtifactPretty artifactRevnonLin+    @?= "\\dret m1 -> let v10 = ssum @3 (sconcrete (sreplicate [3,4] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1))) ; v11 = exp (negate v10) ; v12 = sconcrete (sreplicate [4] 1.0) + v11 ; v13 = recip v12 ; v14 = sconcrete (sreplicate [4] 1.0) + negate v13 ; v15 = v13 * v14 ; m16 = str (sreplicate @5 (scast v13)) ; v17 = scast (ssum @4 (m16 * str (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + sfromR (tproject2 (tproject2 (tproject1 m1))) ; v18 = exp (negate v17) ; v19 = sconcrete (sreplicate [5] 1.0) + v18 ; v20 = recip v19 ; v21 = sconcrete (sreplicate [5] 1.0) + negate v20 ; v22 = v20 * v21 ; v23 = exp (ssum @5 (str (sreplicate @2 v20) * str (sfromR (tproject1 (tproject2 m1)))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum @2 v23 ; v25 = sreplicate @2 (recip x24) ; v26 = v25 * v23 ; v29 = sconcrete (sreplicate [2] 8.0) * (recip v26 * sreplicate @2 (sscalar (-1.0) * sfromK dret)) ; v30 = v23 * (sreplicate @2 (negate (recip (x24 * x24)) * ssum @2 (v23 * v29)) + v25 * v29) ; v31 = v22 * ssum @2 (str (str (sfromR (tproject1 (tproject2 m1))) * sreplicate @5 v30)) ; m32 = sreplicate @4 (scast v31) ; v33 = v15 * scast (ssum @5 (str (str (sfromR (tproject1 (tproject2 (tproject1 m1)))) * m32))) in tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [3,4] 7.0) * sreplicate @3 v33))) (rfromS v33)) (tpair (rfromS (str (m16 * m32))) (rfromS v31))) (tpair (rfromS (str (str (sreplicate @2 v20) * sreplicate @5 v30))) (rfromS v30))"+  printArtifactPretty (simplifyArtifact artifactRevnonLin)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4,3] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [4] FTKScalar)) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5,4] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [5] FTKScalar)) ConvSX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2,5] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [2] FTKScalar)) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKS [4,3] STKScalar) (STKS [4] STKScalar)) (STKProduct (STKS [5,4] STKScalar) (STKS [5] STKScalar))) (STKProduct (STKS [2,5] STKScalar) (STKS [2] STKScalar))) (let v13 = recip (sconcrete (sreplicate [4] 1.0) + exp (negate (sdot1In (sconcrete (sreplicate [4,3] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 m1))))) + negate (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v16 = scast v13 ; v20 = recip (sconcrete (sreplicate [5] 1.0) + exp (negate (scast (sdot1In (sreplicate @5 v16) (sfromR (tproject1 (tproject2 (tproject1 m1)))))) + negate (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; v23 = exp (sdot1In (sreplicate @2 v20) (sfromR (tproject1 (tproject2 m1))) + sfromR (tproject2 (tproject2 m1))) ; x24 = ssum0 v23 ; x25 = recip x24 ; v29 = sconcrete (sreplicate [2] 8.0) * (recip (sreplicate @2 x25 * v23) * sreplicate @2 (sscalar (-1.0) * sfromK dret)) ; v30 = v23 * (sreplicate @2 (negate (recip (x24 * x24)) * sdot0 v23 v29) + sreplicate @2 x25 * v29) ; v31 = (v20 * (sconcrete (sreplicate [5] 1.0) + negate v20)) * sdot1In (str (sfromR (tproject1 (tproject2 m1)))) (sreplicate @5 v30) ; v32 = scast v31 ; v33 = (v13 * (sconcrete (sreplicate [4] 1.0) + negate v13)) * scast (sdot1In (str (sfromR (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @4 v32)) in tpair (tpair (tpair (sconcrete (sreplicate [4,3] 7.0) * str (sreplicate @3 v33)) v33) (tpair (sreplicate @5 v16 * str (sreplicate @4 v32)) v31)) (tpair (sreplicate @2 v20 * str (sreplicate @5 v30)) v30))"++testVT2OAstNonLin3 :: Assertion+testVT2OAstNonLin3 = do+  let ftk = tftk @Concrete (knownSTK @(XParams2 Double Float))+                 (toTarget @Concrete valsInitVT2OPP)+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var :: AstTensor AstMethodLet FullSpan (XParams2 Double Float)+      var = AstVar varName+      vals = toTarget @Concrete valsInitVT2OPP+      env = extendEnv varName vals emptyEnv+      blackGlyph = treplicate (SNat @3) knownSTK $ rscalar 7+      blackLabel = treplicate (SNat @2) knownSTK $ rscalar 8+      afcnn2 :: ADReady f+             => MnistFcnnRanked2.ADFcnnMnist2Parameters f Double Float+             -> f (TKScalar Double)+      afcnn2 = MnistFcnnRanked2.afcnnMnistLoss2+                 ( rconcrete $ unConcrete blackGlyph+                 , rconcrete $ unConcrete blackLabel )+      afcnn1 = afcnn2 $ fromTarget var+  interpretAstFull @Concrete env afcnn1+    @?= afcnn2 valsInitVT2OPP+  interpretAstFull @Concrete env+                   (simplifyInline @(TKScalar Double) afcnn1)+    @?= afcnn2 valsInitVT2OPP+  interpretAstFull @Concrete env+                   (simplifyInlineContract @(TKScalar Double) afcnn1)+    @?= afcnn2 valsInitVT2OPP+++-- * RNNR tests++tensorMnistPPRNNR :: TestTree+tensorMnistPPRNNR = testGroup "PP and Ast tests for RNNR MNIST"+  [ testCase "RNNO PP" testRNNOPP+  , testCase "RNNO Ast" testRNNOAst+  , testCase "RNNO PP 2" testRNNOPP2+  , testCase "RNNO Ast 2" testRNNOAst2+  ]++valsInitRNNOPP+  :: Int -> Int -> ADRnnMnistParameters Concrete Double+valsInitRNNOPP out_width sizeMnistHeightI =+  ( ( Concrete+      $ Nested.rfromListPrimLinear [out_width, sizeMnistHeightI]+          (map fromIntegral [0 .. out_width * sizeMnistHeightI - 1])+    , Concrete+      $ Nested.rfromListPrimLinear [out_width, out_width]+          (map fromIntegral [0 .. out_width * out_width - 1])+    , Concrete+      $ Nested.rfromListPrimLinear [out_width]+          (map fromIntegral [0 .. out_width - 1]) )+  , ( Concrete+      $ Nested.rfromListPrimLinear [out_width, out_width]+          (map fromIntegral [0 .. out_width * out_width - 1])+    , Concrete+      $ Nested.rfromListPrimLinear [out_width, out_width]+          (map fromIntegral [0 .. out_width * out_width - 1])+    , Concrete+      $ Nested.rfromListPrimLinear [out_width]+          (map fromIntegral [0 .. out_width - 1]) )+  , ( Concrete+      $ Nested.rfromListPrimLinear [sizeMnistLabelInt, out_width]+          (map fromIntegral [0 .. sizeMnistLabelInt * out_width - 1])+    , Concrete+      $ Nested.rfromListPrimLinear [sizeMnistLabelInt]+          (map fromIntegral [0 .. sizeMnistLabelInt - 1]) ) )++testRNNOPP :: Assertion+testRNNOPP = do+  resetVarCounter+  let batch_size = 1+      sizeMnistHeightI = 1+      blackGlyph :: AstTensor AstMethodLet PrimalSpan (TKR 3 Double)+      blackGlyph = AstReplicate (SNat @1) knownSTK+                   $ AstReplicate (SNat @1) knownSTK+                   $ AstReplicate (SNat @1) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: ADRnnMnistParameters (AstTensor AstMethodLet FullSpan)+                                      Double+              -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      afcnn2T = MnistRnnRanked2.rnnMnistZeroR batch_size blackGlyph+      ftk = tftk @Concrete+                 (knownSTK @(X (ADRnnMnistParameters Concrete Double)))+                 (toTarget @Concrete $ valsInitRNNOPP 1 sizeMnistHeightI)+      artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2T ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (str (sreplicate @1 (str (sfromR (tproject1 (tproject2 m1))) !$ [0] * sreplicate @10 (tanh (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) !$ [0, 0] * tanh (sscalar 7.0 * sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) !$ [0, 0] + sfromR (tproject2 (tproject1 (tproject1 m1))) !$ [0]) + sfromR (tproject2 (tproject2 (tproject1 m1))) !$ [0])))) + str (sreplicate @1 (sfromR (tproject2 (tproject2 m1)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> let x16 = sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) !$ [0, 0] ; x18 = tanh (sscalar 7.0 * x16 + sfromR (tproject2 (tproject1 (tproject1 m1))) !$ [0]) ; x19 = sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) !$ [0, 0] ; x21 = tanh (x19 * x18 + sfromR (tproject2 (tproject2 (tproject1 m1))) !$ [0]) ; v22 = str (sfromR (tproject1 (tproject2 m1))) !$ [0] in rfromS (str (sreplicate @1 (v22 * sreplicate @10 x21)) + str (sreplicate @1 (sfromR (tproject2 (tproject2 m1)))))"+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let x16 = sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) !$ [0, 0] ; x18 = tanh (sscalar 7.0 * x16 + sfromR (tproject2 (tproject1 (tproject1 m1))) !$ [0]) ; x19 = sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) !$ [0, 0] ; x21 = tanh (x19 * x18 + sfromR (tproject2 (tproject2 (tproject1 m1))) !$ [0]) ; v22 = str (sfromR (tproject1 (tproject2 m1))) !$ [0] ; x24 = (sscalar 1.0 + negate x21 * x21) * ssum @10 (v22 * ssum @1 (str (sfromR dret))) ; x25 = (sscalar 1.0 + negate x18 * x18) * (x19 * x24) in tpair (tpair (tpair (tpair (rfromS (soneHot (sscalar 7.0 * x25) [0, 0])) (rfromS (soneHot (sscalar 0.0) [0, 0]))) (rfromS (soneHot x25 [0]))) (tpair (tpair (rfromS (soneHot (x18 * x24) [0, 0])) (rfromS (soneHot (sscalar 0.0) [0, 0]))) (rfromS (soneHot x24 [0])))) (tpair (rfromS (str (soneHot (sreplicate @10 x21 * ssum @1 (str (sfromR dret))) [0]))) (rfromS (ssum @1 (str (sfromR dret)))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1,1] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1,1] FTKScalar)) ConvSX))) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1] FTKScalar)) ConvSX))) (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1,1] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1,1] FTKScalar)) ConvSX))) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [1] FTKScalar)) ConvSX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [10,1] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [10] FTKScalar)) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKProduct (STKS [1,1] STKScalar) (STKS [1,1] STKScalar)) (STKS [1] STKScalar)) (STKProduct (STKProduct (STKS [1,1] STKScalar) (STKS [1,1] STKScalar)) (STKS [1] STKScalar))) (STKProduct (STKS [10,1] STKScalar) (STKS [10] STKScalar))) (let x18 = tanh (sscalar 7.0 * sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))) !$ [0, 0] + sfromR (tproject2 (tproject1 (tproject1 m1))) !$ [0]) ; x19 = sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))) !$ [0, 0] ; x21 = tanh (x19 * x18 + sfromR (tproject2 (tproject2 (tproject1 m1))) !$ [0]) ; x24 = (sscalar 1.0 + negate x21 * x21) * sdot0 (str (sfromR (tproject1 (tproject2 m1))) !$ [0]) (str (sfromR dret) !$ [0]) ; x25 = (sscalar 1.0 + negate x18 * x18) * (x19 * x24) in tpair (tpair (tpair (tpair (sreplicate @1 (sreplicate @1 (sscalar 7.0 * x25))) (sconcrete (sfromListLinear [1,1] [0.0]))) (sreplicate @1 x25)) (tpair (tpair (sreplicate @1 (sreplicate @1 (x18 * x24))) (sconcrete (sfromListLinear [1,1] [0.0]))) (sreplicate @1 x24))) (tpair (str (sreplicate @1 (sreplicate @10 x21 * str (sfromR dret) !$ [0]))) (str (sfromR dret) !$ [0])))"++testRNNOAst :: Assertion+testRNNOAst = do+  let batch_size = 1+      sizeMnistHeightI = 1+      ftk = tftk @Concrete+                 (knownSTK @(X (ADRnnMnistParameters Concrete Double)))+                 (toTarget @Concrete $ valsInitRNNOPP 1 sizeMnistHeightI)+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var :: AstTensor AstMethodLet FullSpan+                       (X (ADRnnMnistParameters Concrete Double))+      var = AstVar varName+      vals = toTarget @Concrete $ valsInitRNNOPP 1 sizeMnistHeightI+      env = extendEnv varName vals emptyEnv+      blackGlyph = treplicate (SNat @1) knownSTK+                   $ treplicate (SNat @1) knownSTK+                   $ treplicate (SNat @1) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => ADRnnMnistParameters f Double+             -> f (TKR 2 Double)+      afcnn2 = MnistRnnRanked2.rnnMnistZeroR+                 batch_size (rconcrete $ unConcrete blackGlyph)+      afcnn1 = afcnn2 $ fromTarget var+  interpretAstFull @Concrete env afcnn1+    @?= afcnn2 (valsInitRNNOPP 1 sizeMnistHeightI)+  interpretAstFull @Concrete env+                   (simplifyInline @(TKR 2 Double) afcnn1)+    @?= afcnn2 (valsInitRNNOPP 1 sizeMnistHeightI)+  interpretAstFull @Concrete env+                   (simplifyInlineContract @(TKR 2 Double) afcnn1)+    @?= afcnn2 (valsInitRNNOPP 1 sizeMnistHeightI)++testRNNOPP2 :: Assertion+testRNNOPP2 = do+  resetVarCounter+  let batch_size = 2+      sizeMnistHeightI = 2+      blackGlyph :: AstTensor AstMethodLet PrimalSpan (TKR 3 Double)+      blackGlyph = AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                   $ AstReplicate (SNat @2) knownSTK+                       (rconcrete $ Nested.rscalar 7+                        :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double))+      afcnn2T :: ADRnnMnistParameters (AstTensor AstMethodLet FullSpan)+                                      Double+              -> AstTensor AstMethodLet FullSpan (TKR 2 Double)+      afcnn2T = MnistRnnRanked2.rnnMnistZeroR batch_size blackGlyph+      ftk = tftk @Concrete+                 (knownSTK @(X (ADRnnMnistParameters Concrete Double)))+                 (toTarget @Concrete $ valsInitRNNOPP 2 sizeMnistHeightI)+      artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2T ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\m1 -> rfromS (let m40 = sappend (tanh (str (sreplicate @2 (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))))) (tanh (str (sreplicate @2 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @2 (tanh (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1))))))) in smatmul2 (sfromR (tproject1 (tproject2 m1))) (tanh ((smatmul2 (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (tanh ((str (sreplicate @2 (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + smatmul2 (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))) (sslice (SNat @0) (SNat @2) m40)) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))))) + smatmul2 (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))) (sslice (SNat @2) (SNat @2) m40)) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 m1)))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\m1 -> let m37 = tanh ((str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))))) + str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 0.0))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v38 = tanh ((ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) + ssum @2 (sconcrete (sreplicate [2,2] 0.0))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; m39 = tanh ((str (sreplicate @2 (ssum @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) * str (sreplicate @2 v38)))) + str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 0.0))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m40 = sappend m37 m39 ; m41 = tanh ((sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + ssum @2 (str (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) * stranspose @[2,1,0] (sreplicate @2 (str (sslice (SNat @0) (SNat @2) m40))))) + sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))) ; m42 = tanh ((ssum @2 (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2,0,1] (sreplicate @2 m41)) + ssum @2 (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2,0,1] (sreplicate @2 (str (sslice (SNat @2) (SNat @2) m40))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) in rfromS (ssum @2 (stranspose @[2,1,0] (sreplicate @2 (sfromR (tproject1 (tproject2 m1)))) * str (sreplicate @10 m42)) + str (sreplicate @2 (sfromR (tproject2 (tproject2 m1)))))"+  printArtifactPretty artifactRev+    @?= "\\dret m1 -> let m37 = tanh ((str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))))) + str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 0.0))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v38 = tanh ((ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) + ssum @2 (sconcrete (sreplicate [2,2] 0.0))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; m39 = tanh ((str (sreplicate @2 (ssum @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) * str (sreplicate @2 v38)))) + str (sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 0.0))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m40 = sappend m37 m39 ; m41 = tanh ((sreplicate @2 (ssum @2 (sconcrete (sreplicate [2,2] 7.0) * str (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + ssum @2 (str (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) * stranspose @[2,1,0] (sreplicate @2 (str (sslice (SNat @0) (SNat @2) m40))))) + sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))) ; m42 = tanh ((ssum @2 (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2,0,1] (sreplicate @2 m41)) + ssum @2 (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))))) * stranspose @[2,0,1] (sreplicate @2 (str (sslice (SNat @2) (SNat @2) m40))))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m44 = (sconcrete (sreplicate [2,2] 1.0) + negate m42 * m42) * ssum @10 (str (stranspose @[2,1,0] (sreplicate @2 (sfromR (tproject1 (tproject2 m1)))) * sreplicate @2 (sfromR dret))) ; m45 = (sconcrete (sreplicate [2,2] 1.0) + negate m41 * m41) * ssum @2 (stranspose @[1,2,0] (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))))) * sreplicate @2 m44)) ; m46 = sappend (sconcrete (sfromListLinear [0,2] [])) (sappend (str (ssum @2 (stranspose @[2,1,0] (str (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) * sreplicate @2 m45)))) (sconcrete (sreplicate [2,2] 0.0))) + sappend (sconcrete (sreplicate [2,2] 0.0)) (sappend (str (ssum @2 (stranspose @[1,2,0] (stranspose @[1,2,0] (sreplicate @2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))))) * sreplicate @2 m44)))) (sconcrete (sfromListLinear [0,2] []))) ; m47 = (sconcrete (sreplicate [2,2] 1.0) + negate m39 * m39) * sslice (SNat @2) (SNat @2) m46 ; m48 = sreplicate @2 (ssum @2 (str m47)) ; v49 = (sconcrete (sreplicate [2] 1.0) + negate v38 * v38) * ssum @2 (str (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) * m48)) ; m50 = (sconcrete (sreplicate [2,2] 1.0) + negate m37 * m37) * sslice (SNat @0) (SNat @2) m46 in tpair (tpair (tpair (tpair (rfromS (str (sconcrete (sreplicate [2,2] 7.0) * sreplicate @2 (ssum @2 (str m50))) + (str (sconcrete (sreplicate [2,2] 7.0) * sreplicate @2 v49) + str (sconcrete (sreplicate [2,2] 7.0) * sreplicate @2 (ssum @2 m45))))) (rfromS (str (sconcrete (sreplicate [2,2] 0.0)) + (str (sconcrete (sreplicate [2,2] 0.0)) + str (ssum @2 (str (stranspose @[2,1,0] (sreplicate @2 (str (sslice (SNat @0) (SNat @2) m40))) * sreplicate @2 m45))))))) (rfromS (ssum @2 (str m50) + (v49 + ssum @2 m45)))) (tpair (tpair (rfromS (str (str (sreplicate @2 v38) * m48) + str (ssum @2 (stranspose @[2,0,1] (stranspose @[2,0,1] (sreplicate @2 m41) * sreplicate @2 m44))))) (rfromS (str (sconcrete (sreplicate [2,2] 0.0)) + str (ssum @2 (stranspose @[2,0,1] (stranspose @[2,0,1] (sreplicate @2 (str (sslice (SNat @2) (SNat @2) m40))) * sreplicate @2 m44)))))) (rfromS (ssum @2 (str m47) + ssum @2 (str m44))))) (tpair (rfromS (ssum @2 (stranspose @[2,1,0] (str (sreplicate @10 m42) * sreplicate @2 (sfromR dret))))) (rfromS (ssum @2 (str (sfromR dret)))))"+  printArtifactPretty (simplifyArtifact artifactRev)+    @?= "\\dret m1 -> tconvert (ConvT2 (ConvT2 (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvT2 (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [10,2] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [10] FTKScalar)) ConvSX)))) (STKProduct (STKProduct (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2] STKScalar)) (STKProduct (STKProduct (STKS [2,2] STKScalar) (STKS [2,2] STKScalar)) (STKS [2] STKScalar))) (STKProduct (STKS [10,2] STKScalar) (STKS [10] STKScalar))) (let m37 = tanh (str (sreplicate @2 (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))))) + str (sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1)))))) ; v38 = tanh (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1))))) + sfromR (tproject2 (tproject1 (tproject1 m1)))) ; m39 = tanh (str (sreplicate @2 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (sreplicate @2 v38))) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m40 = sappend m37 m39 ; m41 = tanh ((sreplicate @2 (sdot1In (sconcrete (sreplicate [2,2] 7.0)) (sfromR (tproject1 (tproject1 (tproject1 (tproject1 m1)))))) + smatmul2 (str (sslice (SNat @0) (SNat @2) m40)) (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1))))))) + sreplicate @2 (sfromR (tproject2 (tproject1 (tproject1 m1))))) ; m42 = tanh ((smatmul2 (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) (str m41) + smatmul2 (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1))))) (sslice (SNat @2) (SNat @2) m40)) + str (sreplicate @2 (sfromR (tproject2 (tproject2 (tproject1 m1)))))) ; m44 = (sconcrete (sreplicate [2,2] 1.0) + negate m42 * m42) * smatmul2 (str (sfromR (tproject1 (tproject2 m1)))) (sfromR dret) ; m45 = (sconcrete (sreplicate [2,2] 1.0) + negate m41 * m41) * smatmul2 (str m44) (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1))))) ; m46 = sappend (smatmul2 (str (sfromR (tproject2 (tproject1 (tproject1 (tproject1 m1)))))) (str m45)) (sconcrete (sreplicate [2,2] 0.0)) + sappend (sconcrete (sreplicate [2,2] 0.0)) (smatmul2 (str (sfromR (tproject2 (tproject1 (tproject2 (tproject1 m1)))))) m44) ; m47 = (sconcrete (sreplicate [2,2] 1.0) + negate m39 * m39) * sslice (SNat @2) (SNat @2) m46 ; v48 = ssum @2 (str m47) ; v49 = (sconcrete (sreplicate [2] 1.0) + negate v38 * v38) * sdot1In (str (sfromR (tproject1 (tproject1 (tproject2 (tproject1 m1)))))) (sreplicate @2 v48) ; m50 = (sconcrete (sreplicate [2,2] 1.0) + negate m37 * m37) * sslice (SNat @0) (SNat @2) m46 in tpair (tpair (tpair (tpair (sconcrete (sreplicate [2,2] 7.0) * str (sreplicate @2 (ssum @2 (str m50))) + (sconcrete (sreplicate [2,2] 7.0) * str (sreplicate @2 v49) + sconcrete (sreplicate [2,2] 7.0) * str (sreplicate @2 (ssum @2 m45)))) (smatmul2 (str m45) (str (sslice (SNat @0) (SNat @2) m40)))) (ssum @2 (str m50) + (v49 + ssum @2 m45))) (tpair (tpair (sreplicate @2 v38 * str (sreplicate @2 v48) + smatmul2 m44 m41) (smatmul2 m44 (str (sslice (SNat @2) (SNat @2) m40)))) (ssum @2 (str m47) + ssum @2 (str m44)))) (tpair (smatmul2 (sfromR dret) (str m42)) (ssum @2 (str (sfromR dret)))))"++testRNNOAst2 :: Assertion+testRNNOAst2 = do+  let batch_size = 2+      sizeMnistHeightI = 2+      ftk = tftk @Concrete+                 (knownSTK @(X (ADRnnMnistParameters Concrete Double)))+                 (toTarget @Concrete $ valsInitRNNOPP 2 sizeMnistHeightI)+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var :: AstTensor AstMethodLet FullSpan+                       (X (ADRnnMnistParameters Concrete Double))+      var = AstVar varName+      vals = toTarget @Concrete $ valsInitRNNOPP 2 sizeMnistHeightI+      env = extendEnv varName vals emptyEnv+      blackGlyph = treplicate (SNat @2) knownSTK+                   $ treplicate (SNat @2) knownSTK+                   $ treplicate (SNat @2) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => ADRnnMnistParameters f Double+             -> f (TKR 2 Double)+      afcnn2 = MnistRnnRanked2.rnnMnistZeroR+                 batch_size (rconcrete $ unConcrete blackGlyph)+      afcnn1 = afcnn2 $ fromTarget var+  interpretAstFull @Concrete env afcnn1+    @?= afcnn2 (valsInitRNNOPP 2 sizeMnistHeightI)+  interpretAstFull @Concrete env+                   (simplifyInline @(TKR 2 Double) afcnn1)+    @?= afcnn2 (valsInitRNNOPP 2 sizeMnistHeightI)+  interpretAstFull @Concrete env+                   (simplifyInlineContract @(TKR 2 Double) afcnn1)+    @?= afcnn2 (valsInitRNNOPP 2 sizeMnistHeightI)+++-- * CNNR tests++tensorMnistCNNRPP :: TestTree+tensorMnistCNNRPP = testGroup "Ast tests for CNNR MNIST"+  [ testCase "CNNO PP 1" testCNNOPP1+  , testCase "CNNO Ast 1" testCNNOAst1+  , testCase "CNNO PP 2" testCNNOPP2+  , testCase "CNNO Ast 2" testCNNOAst2+  , testCase "CNNO PP 2S" testCNNOPP2S+  ]++testCNNOPP1 :: Assertion+testCNNOPP1 = do+  resetVarCounter+  let batch_size = 5+      sizeMnistWidthI = 7+      sizeMnistHeightI = 9+      ftk = tftk @Concrete+                 (knownSTK @(X (MnistCnnRanked2.ADCnnMnistParameters+                                  Concrete Double)))+                 vals+      valsInit :: MnistCnnRanked2.ADCnnMnistParameters Concrete Double+      valsInit =+        forgetShape $ fst+        $ randomValue @(MnistCnnRanked2.ADCnnMnistParametersShaped+                         Concrete 7 9  -- see sizeMnistWidthI, etc.+                         1 1 1 1 Double)+                      0.4 (mkStdGen 44)+      vals = toTarget @Concrete valsInit+      blackGlyph = treplicate (SNat @5) knownSTK+                   $ treplicate (SNat @1) knownSTK+                   $ treplicate (SNat @7) knownSTK+                   $ treplicate (SNat @9) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => MnistCnnRanked2.ADCnnMnistParameters f Double+             -> f (TKR 2 Double)+      afcnn2 = MnistCnnRanked2.convMnistTwoR+                 sizeMnistHeightI sizeMnistWidthI batch_size+                 (rconcrete $ unConcrete blackGlyph)+      artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2 ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (let t189 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[7,9,4] (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [7,9] 7.0)) (\\[i358, i359] -> [i358 + i359]))) (\\[i185, i186] -> [i185 + i186])))))) * sreplicate @7 (sreplicate @9 (sreplicate @1 (sreplicate @1 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0])))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @9 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; t204 = sreshape @[3,4,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i192, i193, i194, i195] -> [ifH (sscalar -0.0 <=. negate (t189 !$ [0, kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i192, i194]), kfromS (sconcrete (sfromListLinear [4,2] [0,1,2,3,4,5,6,7]) !$ [i193, i195])])) 0 1]) * sgather (t189 !$ [0]) (\\[i198, i199, i200, i201] -> [kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i198, i200]), kfromS (sconcrete (sfromListLinear [4,2] [0,1,2,3,4,5,6,7]) !$ [i199, i201])])) ; t213 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[3,4,4] (stranspose @[2,0,3,1] (sgather (stranspose @[2,3,4,0,1] (sgather (stranspose @[3,5,0,4,1,2] (sgather (stranspose @[3,2,1,6,5,4,0] (sreplicate @2 (stranspose @[5,4,3,0,1,2] (sreplicate @4 (sreplicate @3 (sreplicate @2 (stranspose @[2,1,0] t204))))))) (\\[i346, i350] -> [kfromS (smaxIndex (t204 !$ [i346, i350])), i350, i346]))) (\\[i355, i356] -> [i355, i356, i355 + i356]))) (\\[i209, i210] -> [i209, i209 + i210, i210])) * sreplicate @3 (sreplicate @4 (sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0])))))) + stranspose @[2,0,1] (sreplicate @3 (sreplicate @4 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m223 = sreshape @[2,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i215, i216, i217] -> [ifH (sscalar -0.0 <=. negate (t213 !$ [0, i216, kfromS (sconcrete (sfromListLinear [2,2] [0,1,2,3]) !$ [i215, i217])])) 0 1]) * stranspose @[0,2,1] (sgather (str (sslice (SNat @0) (SNat @2) (t213 !$ [0]))) (\\[i219, i220] -> [kfromS (sconcrete (sfromListLinear [2,2] [0,1,2,3]) !$ [i219, i220])]))) ; m227 = sreplicate @1 (sreplicate @5 (sdot0 (sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0]) (sgather m223 (\\[i224] -> [i224, kfromS (smaxIndex (m223 !$ [i224]))])))) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) in str (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) * sreplicate @10 (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i228] -> [ifH (sscalar -0.0 <=. negate (m227 !$ [0, i228])) 0 1]) * m227 !$ [0]) + str (sreplicate @5 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w187 = stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [7,9] 7.0)) (\\[i183, i184] -> [i183 + i184]))) (\\[i185, i186] -> [i185 + i186])))))) ; w188 = sreplicate @7 (sreplicate @9 (sreplicate @1 (sreplicate @1 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0])))) ; t189 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[7,9,4] (w187 * w188)))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @9 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; m190 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m191 = str (sreplicate @2 (sconcrete (sreplicate [4] 2) * siota (SNat @4))) + sreplicate @4 (siota (SNat @2)) ; u202 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i192, i193, i194, i195] -> [let x196 = m190 !$ [i192, i194] ; x197 = m191 !$ [i193, i195] in ifH (sscalar -0.0 <=. negate (t189 !$ [0, kfromS x196, kfromS x197])) 0 1]) ; u203 = sgather (t189 !$ [0]) (\\[i198, i199, i200, i201] -> [kfromS (m190 !$ [i198, i200]), kfromS (m191 !$ [i199, i201])]) ; t204 = sreshape @[3,4,4] (u202 * u203) ; u211 = stranspose @[2,0,3,1] (sgather (stranspose @[2,3,4,0,1] (sgather (stranspose @[3,5,0,4,1,2] (sgather (stranspose @[3,2,1,6,5,4,0] (sreplicate @2 (stranspose @[5,4,3,0,1,2] (sreplicate @4 (sreplicate @3 (sreplicate @2 (stranspose @[2,1,0] t204))))))) (\\[i205, i206] -> [kfromS (smaxIndex (t204 !$ [i205, i206])), i206, i205]))) (\\[i207, i208] -> [i207, i208, i207 + i208]))) (\\[i209, i210] -> [i209, i209 + i210, i210])) ; u212 = sreplicate @3 (sreplicate @4 (sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0])) ; t213 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[3,4,4] (u211 * u212)))) + stranspose @[2,0,1] (sreplicate @3 (sreplicate @4 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m214 = str (sreplicate @2 (sconcrete (sreplicate [2] 2) * siota (SNat @2))) + sreplicate @2 (siota (SNat @2)) ; t221 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i215, i216, i217] -> [let x218 = m214 !$ [i215, i217] in ifH (sscalar -0.0 <=. negate (t213 !$ [0, i216, kfromS x218])) 0 1]) ; t222 = stranspose @[0,2,1] (sgather (str (sslice (SNat @0) (SNat @2) (t213 !$ [0]))) (\\[i219, i220] -> [kfromS (m214 !$ [i219, i220])])) ; m223 = sreshape @[2,4] (t221 * t222) ; v225 = sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0] ; v226 = sgather m223 (\\[i224] -> [i224, kfromS (smaxIndex (m223 !$ [i224]))]) ; m227 = sreplicate @1 (sreplicate @5 (ssum @2 (v225 * v226))) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; v229 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i228] -> [ifH (sscalar -0.0 <=. negate (m227 !$ [0, i228])) 0 1]) ; v230 = m227 !$ [0] ; m231 = str (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) ; m232 = sreplicate @10 (v229 * v230) in rfromS (m231 * m232 + str (sreplicate @5 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w187 = stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [7,9] 7.0)) (\\[i183, i184] -> [i183 + i184]))) (\\[i185, i186] -> [i185 + i186])))))) ; w188 = sreplicate @7 (sreplicate @9 (sreplicate @1 (sreplicate @1 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0])))) ; t189 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[7,9,4] (w187 * w188)))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @9 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; m190 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m191 = str (sreplicate @2 (sconcrete (sreplicate [4] 2) * siota (SNat @4))) + sreplicate @4 (siota (SNat @2)) ; u202 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i192, i193, i194, i195] -> [let x196 = m190 !$ [i192, i194] ; x197 = m191 !$ [i193, i195] in ifH (sscalar -0.0 <=. negate (t189 !$ [0, kfromS x196, kfromS x197])) 0 1]) ; u203 = sgather (t189 !$ [0]) (\\[i198, i199, i200, i201] -> [kfromS (m190 !$ [i198, i200]), kfromS (m191 !$ [i199, i201])]) ; t204 = sreshape @[3,4,4] (u202 * u203) ; u211 = stranspose @[2,0,3,1] (sgather (stranspose @[2,3,4,0,1] (sgather (stranspose @[3,5,0,4,1,2] (sgather (stranspose @[3,2,1,6,5,4,0] (sreplicate @2 (stranspose @[5,4,3,0,1,2] (sreplicate @4 (sreplicate @3 (sreplicate @2 (stranspose @[2,1,0] t204))))))) (\\[i205, i206] -> [kfromS (smaxIndex (t204 !$ [i205, i206])), i206, i205]))) (\\[i207, i208] -> [i207, i208, i207 + i208]))) (\\[i209, i210] -> [i209, i209 + i210, i210])) ; u212 = sreplicate @3 (sreplicate @4 (sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0])) ; t213 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[3,4,4] (u211 * u212)))) + stranspose @[2,0,1] (sreplicate @3 (sreplicate @4 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m214 = str (sreplicate @2 (sconcrete (sreplicate [2] 2) * siota (SNat @2))) + sreplicate @2 (siota (SNat @2)) ; t221 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i215, i216, i217] -> [let x218 = m214 !$ [i215, i217] in ifH (sscalar -0.0 <=. negate (t213 !$ [0, i216, kfromS x218])) 0 1]) ; t222 = stranspose @[0,2,1] (sgather (str (sslice (SNat @0) (SNat @2) (t213 !$ [0]))) (\\[i219, i220] -> [kfromS (m214 !$ [i219, i220])])) ; m223 = sreshape @[2,4] (t221 * t222) ; v225 = sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0] ; v226 = sgather m223 (\\[i224] -> [i224, kfromS (smaxIndex (m223 !$ [i224]))]) ; m227 = sreplicate @1 (sreplicate @5 (ssum @2 (v225 * v226))) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; v229 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i228] -> [ifH (sscalar -0.0 <=. negate (m227 !$ [0, i228])) 0 1]) ; v230 = m227 !$ [0] ; m231 = str (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) ; m232 = sreplicate @10 (v229 * v230) ; m234 = soneHot (v229 * ssum @10 (m231 * sfromR dret)) [0] ; v235 = sreplicate @2 (ssum @5 (ssum @1 m234)) ; t239 = soneHot (sappend (sconcrete (sfromListLinear [0,4] [])) (sappend (str (sscatter (stranspose @[0,2,1] (t221 * sreshape @[2,2,2] (sscatter (v225 * v235) (\\[i236] -> [i236, kfromS (smaxIndex (m223 !$ [i236]))])))) (\\[i237, i238] -> [kfromS (m214 !$ [i237, i238])]))) (sconcrete (sreplicate [1,4] 0.0)))) [0] ; u240 = sreshape @[3,4,2,2] (stranspose @[1,2,0] (sreplicate @4 (ssum @1 t239))) ; t251 = soneHot (sscatter (u202 * sreshape @[3,4,2,2] (stranspose @[2,1,0] (ssum @2 (ssum @3 (ssum @4 (stranspose @[3,4,5,2,1,0] (ssum @2 (stranspose @[6,2,1,0,5,4,3] (sscatter (stranspose @[2,4,5,0,3,1] (sscatter (stranspose @[3,4,0,1,2] (sscatter (stranspose @[1,3,0,2] (u212 * u240)) (\\[i241, i242] -> [i241, i241 + i242, i242]))) (\\[i243, i244] -> [i243, i244, i243 + i244]))) (\\[i245, i246] -> [kfromS (smaxIndex (t204 !$ [i245, i246])), i246, i245])))))))))) (\\[i247, i248, i249, i250] -> [kfromS (m190 !$ [i247, i249]), kfromS (m191 !$ [i248, i250])])) [0] in tpair (tpair (tpair (rfromS (soneHot (ssum @1 (ssum @1 (ssum @9 (ssum @7 (w187 * sreshape @[7,9,1,1,2,2] (stranspose @[1,2,0] (sreplicate @4 (ssum @1 t251)))))))) [0, 0])) (rfromS (ssum @9 (ssum @7 (stranspose @[1,2,0] t251))))) (tpair (rfromS (soneHot (ssum @4 (ssum @3 (u211 * u240))) [0, 0])) (rfromS (ssum @4 (ssum @3 (stranspose @[1,2,0] t239)))))) (tpair (tpair (rfromS (soneHot (v226 * v235) [0])) (rfromS (ssum @5 (str m234)))) (tpair (rfromS (str (soneHot (ssum @5 (str (m232 * sfromR dret))) [0]))) (rfromS (ssum @5 (str (sfromR dret))))))"+-- TODO: different test result with GHC 9.10:   printArtifactPretty (simplifyArtifact artifactRev)+--    @?= "\\dret u1 -> tfromS (STKProduct (STKProduct (STKProduct (STKR (SNat @4) STKScalar) (STKR (SNat @1) STKScalar)) (STKProduct (STKR (SNat @4) STKScalar) (STKR (SNat @1) STKScalar))) (STKProduct (STKProduct (STKR (SNat @2) STKScalar) (STKR (SNat @1) STKScalar)) (STKProduct (STKR (SNat @2) STKScalar) (STKR (SNat @1) STKScalar)))) (let u181 = sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [7,9] 7.0)) (\\[i296, i297] -> [i296 + i297]))) (\\[i179, i180] -> [i179 + i180]) ; t183 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[7,9,4] (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] u181)))) * sreplicate @7 (sreplicate @9 (sreplicate @1 (sreplicate @1 (sfromR (tproject1 (tproject1 (tproject1 u1))) !$ [0, 0])))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @9 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; u192 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i184, i185, i186, i187] -> [ifH (sscalar -0.0 <=. negate (t183 !$ [0, 2 * i184 + i186, 2 * i185 + i187])) 0 1]) ; t194 = sreshape @[3,4,4] (u192 * sgather (t183 !$ [0]) (\\[i188, i189, i190, i191] -> [2 * i188 + i190, 2 * i189 + i191])) ; u201 = sgather (stranspose @[2,3,4,0,1] (sgather (stranspose @[3,5,0,4,1,2] (sgather (stranspose @[3,2,1,6,5,4,0] (sreplicate @2 (stranspose @[5,4,3,0,1,2] (sreplicate @4 (sreplicate @3 (sreplicate @2 (stranspose @[2,1,0] t194))))))) (\\[i284, i288] -> [kfromS (smaxIndex (t194 !$ [i284, i288])), i288, i284]))) (\\[i293, i294] -> [i293, i294, i293 + i294]))) (\\[i199, i200] -> [i199, i199 + i200, i200]) ; m202 = sfromR (tproject1 (tproject2 (tproject1 u1))) !$ [0, 0] ; t203 = sreplicate @1 (ssum @4 (stranspose @[2,0,1] (sreshape @[3,4,4] (stranspose @[2,0,3,1] u201 * sreplicate @3 (sreplicate @4 m202))))) + stranspose @[2,0,1] (sreplicate @3 (sreplicate @4 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; t209 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i204, i205, i206] -> [ifH (sscalar -0.0 <=. negate (t203 !$ [0, i205, 2 * i204 + i206])) 0 1]) ; m211 = sreshape @[2,4] (t209 * stranspose @[0,2,1] (sgather (str (sslice (SNat @0) (SNat @2) (t203 !$ [0]))) (\\[i207, i208] -> [2 * i207 + i208]))) ; v213 = sfromR (tproject1 (tproject1 (tproject2 u1))) !$ [0] ; v214 = sgather m211 (\\[i212] -> [i212, kfromS (smaxIndex (m211 !$ [i212]))]) ; m215 = sreplicate @1 (sreplicate @5 (sdot0 v213 v214)) + str (sreplicate @5 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; v217 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i216] -> [ifH (sscalar -0.0 <=. negate (m215 !$ [0, i216])) 0 1]) ; v222 = v217 * sdot1In (sreplicate @5 (str (sfromR (tproject1 (tproject2 (tproject2 u1)))) !$ [0])) (str (sfromR dret)) ; x223 = ssum0 v222 ; m227 = sappend (str (sscatter (stranspose @[0,2,1] t209 * stranspose @[0,2,1] (sreshape @[2,2,2] (sscatter (v213 * sreplicate @2 x223) (\\[i224] -> [i224, kfromS (smaxIndex (m211 !$ [i224]))])))) (\\[i225, i226] -> [2 * i225 + i226]))) (sconcrete (sreplicate [1,4] 0.0)) ; u228 = sreshape @[3,4,2,2] (stranspose @[1,2,0] (sreplicate @4 m227)) ; m239 = sscatter (u192 * sreshape @[3,4,2,2] (ssum @2 (ssum @3 (ssum @4 (ssum @2 (stranspose @[6,5,4,3,2,1,0] (sscatter (stranspose @[2,4,5,0,3,1] (sscatter (stranspose @[3,4,0,1,2] (sscatter (stranspose @[1,3,0,2] (sreplicate @3 (sreplicate @4 m202)) * stranspose @[1,3,0,2] u228) (\\[i229, i230] -> [i229, i229 + i230, i230]))) (\\[i231, i232] -> [i231, i232, i231 + i232]))) (\\[i233, i234] -> [kfromS (smaxIndex (t194 !$ [i233, i234])), i234, i233])))))))) (\\[i235, i236, i237, i238] -> [2 * i235 + i237, 2 * i236 + i238]) in tpair (tpair (tpair (sreplicate @1 (sreplicate @1 (ssum @9 (sdot1In (stranspose @[0,3,1,2] u181) (stranspose @[2,3,1,4,5,0] (sreshape @[7,9,1,1,2,2] (stranspose @[1,2,0] (sreplicate @4 m239))) !$ [0, 0]))))) (ssum @9 (ssum @7 (stranspose @[1,2,0] (sreplicate @1 m239))))) (tpair (sreplicate @1 (sreplicate @1 (ssum @4 (sdot1In (stranspose @[0,3,1,2] u201) (stranspose @[1,2,3,0] u228))))) (ssum @4 (ssum @3 (stranspose @[1,2,0] (sreplicate @1 m227)))))) (tpair (tpair (sreplicate @1 (v214 * sreplicate @2 x223)) (sreplicate @1 (ssum0 v222))) (tpair (str (sreplicate @1 (sdot1In (sreplicate @10 (v217 * m215 !$ [0])) (sfromR dret)))) (ssum @5 (str (sfromR dret))))))"++testCNNOAst1 :: Assertion+testCNNOAst1 = do+  let batch_size = 5+      sizeMnistWidthI = 7+      sizeMnistHeightI = 9+      ftk = tftk @Concrete+                 (knownSTK @(X (MnistCnnRanked2.ADCnnMnistParameters+                                  Concrete Double)))+                 vals+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var :: AstTensor AstMethodLet FullSpan+                       (X (MnistCnnRanked2.ADCnnMnistParameters+                             Concrete Double))+      var = AstVar varName+      valsInit :: MnistCnnRanked2.ADCnnMnistParameters Concrete Double+      valsInit =+        forgetShape $ fst+        $ randomValue @(MnistCnnRanked2.ADCnnMnistParametersShaped+                         Concrete 7 9  -- see sizeMnistWidthI, etc.+                         1 1 1 1 Double)+                      0.4 (mkStdGen 44)+      vals = toTarget @Concrete valsInit+      env = extendEnv varName vals emptyEnv+      blackGlyph = treplicate (SNat @5) knownSTK+                   $ treplicate (SNat @1) knownSTK+                   $ treplicate (SNat @7) knownSTK+                   $ treplicate (SNat @9) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => MnistCnnRanked2.ADCnnMnistParameters f Double+             -> f (TKR 2 Double)+      afcnn2 = MnistCnnRanked2.convMnistTwoR+                 sizeMnistHeightI sizeMnistWidthI batch_size+                 (rconcrete $ unConcrete blackGlyph)+      afcnn1 = afcnn2 $ fromTarget var+  interpretAstFull @Concrete env afcnn1+    @?= afcnn2 valsInit+  interpretAstFull @Concrete env+                   (simplifyInline @(TKR 2 Double) afcnn1)+    @?= afcnn2 valsInit+  interpretAstFull @Concrete env+                   (simplifyInlineContract @(TKR 2 Double) afcnn1)+    @?= afcnn2 valsInit++testCNNOPP2 :: Assertion+testCNNOPP2 = do+  resetVarCounter+  let batch_size = 7+      sizeMnistWidthI = 14+      sizeMnistHeightI = 23+      ftk = tftk @Concrete+                 (knownSTK @(X (MnistCnnRanked2.ADCnnMnistParameters+                                  Concrete Double)))+                 vals+      valsInit :: MnistCnnRanked2.ADCnnMnistParameters Concrete Double+      valsInit =+        forgetShape $ fst+        $ randomValue @(MnistCnnRanked2.ADCnnMnistParametersShaped+                         Concrete 14 23  -- see sizeMnistWidthI, etc.+                         2 3 4 5 Double)+                      0.4 (mkStdGen 44)+      vals = toTarget @Concrete valsInit+      blackGlyph = treplicate (SNat @7) knownSTK+                   $ treplicate (SNat @1) knownSTK+                   $ treplicate (SNat @14) knownSTK+                   $ treplicate (SNat @23) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => MnistCnnRanked2.ADCnnMnistParameters f Double+             -> f (TKR 2 Double)+      afcnn2 = MnistCnnRanked2.convMnistTwoR+                 sizeMnistHeightI sizeMnistWidthI batch_size+                 (rconcrete $ unConcrete blackGlyph)+      artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2 ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> rfromS (let t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i473, i474] -> [i473 + i474]))) (\\[i251, i252] -> [i251 + i252]))))))) * str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0]))))))))))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; u271 = sreshape @[4,7,11,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS (sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) !$ [i259, i261]), kfromS (sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) !$ [i260, i262])])) 0 1]) * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) !$ [i265, i267]), kfromS (sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) !$ [i266, i268])]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i458, i460, i461] -> [kfromS (smaxIndex (u271 !$ [i461, i458, i460])), i460, i458, i461]))) (\\[i466, i468] -> [i466, i468, i466 + i468]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) * str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; u296 = sreshape @[4,3,5,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i284, i286]), kfromS (sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) !$ [i285, i287])])) 0 1]) * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i290, i292]), kfromS (sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) !$ [i291, i293])]))) ; m301 = str (sreplicate @7 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 u1)))) (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) in smatmul2 (sfromR (tproject1 (tproject2 (tproject2 u1)))) (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) * m301) + str (sreplicate @7 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w253 = sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i249, i250] -> [i249 + i250]))) (\\[i251, i252] -> [i251 + i252]))))))) ; w254 = str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0])))))))) ; t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (w253 * w254))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; m256 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; m257 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; w269 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [let x263 = m256 !$ [i259, i261] ; x264 = m257 !$ [i260, i262] in ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS x263, kfromS x264])) 0 1]) ; w270 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (m256 !$ [i265, i267]), kfromS (m257 !$ [i266, i268])])) ; u271 = sreshape @[4,7,11,4] (w269 * w270) ; w279 = sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i272, i273, i274] -> [kfromS (smaxIndex (u271 !$ [i274, i272, i273])), i273, i272, i274]))) (\\[i275, i276] -> [i275, i276, i275 + i276]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (w279 * str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m281 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m282 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; w294 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [let x288 = m281 !$ [i284, i286] ; x289 = m282 !$ [i285, i287] in ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS x288, kfromS x289])) 0 1]) ; w295 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (m281 !$ [i290, i292]), kfromS (m282 !$ [i291, i293])])) ; u296 = sreshape @[4,3,5,4] (w294 * w295) ; m300 = str (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))) ; m301 = str (sreplicate @7 (ssum @60 (str (sfromR (tproject1 (tproject1 (tproject2 u1)))) * m300))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; m304 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) ; t305 = str (sreplicate @10 (m304 * m301)) in rfromS (ssum @5 (stranspose @[2,1,0] (sreplicate @7 (sfromR (tproject1 (tproject2 (tproject2 u1))))) * t305) + str (sreplicate @7 (sfromR (tproject2 (tproject2 (tproject2 u1))))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w253 = sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i249, i250] -> [i249 + i250]))) (\\[i251, i252] -> [i251 + i252]))))))) ; w254 = str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0])))))))) ; t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (w253 * w254))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; m256 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; m257 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; w269 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [let x263 = m256 !$ [i259, i261] ; x264 = m257 !$ [i260, i262] in ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS x263, kfromS x264])) 0 1]) ; w270 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (m256 !$ [i265, i267]), kfromS (m257 !$ [i266, i268])])) ; u271 = sreshape @[4,7,11,4] (w269 * w270) ; w279 = sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i272, i273, i274] -> [kfromS (smaxIndex (u271 !$ [i274, i272, i273])), i273, i272, i274]))) (\\[i275, i276] -> [i275, i276, i275 + i276]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (w279 * str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; m281 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m282 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; w294 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [let x288 = m281 !$ [i284, i286] ; x289 = m282 !$ [i285, i287] in ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS x288, kfromS x289])) 0 1]) ; w295 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (m281 !$ [i290, i292]), kfromS (m282 !$ [i291, i293])])) ; u296 = sreshape @[4,3,5,4] (w294 * w295) ; m300 = str (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))) ; m301 = str (sreplicate @7 (ssum @60 (str (sfromR (tproject1 (tproject1 (tproject2 u1)))) * m300))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; m304 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) ; t305 = str (sreplicate @10 (m304 * m301)) ; m307 = m304 * ssum @10 (str (stranspose @[2,1,0] (sreplicate @7 (sfromR (tproject1 (tproject2 (tproject2 u1))))) * sreplicate @5 (sfromR dret))) ; m308 = sreplicate @60 (ssum @7 (str m307)) ; t316 = stranspose @[2,0,1] (sscatter (stranspose @[1,2,3,4,0] (w294 * sreshape @[4,3,5,2,2] (sscatter (sreshape @[4,3,5] (ssum @5 (str (str (sfromR (tproject1 (tproject1 (tproject2 u1)))) * m308)))) (\\[i309, i310, i311] -> [i309, i310, i311, kfromS (smaxIndex (u296 !$ [i309, i310, i311]))])))) (\\[i312, i313, i314, i315] -> [kfromS (m281 !$ [i312, i314]), kfromS (m282 !$ [i313, i315])])) ; w317 = sreshape @[4,7,11,4,3,4] (stranspose @[1,2,3,0] (sreplicate @48 t316)) ; t329 = stranspose @[2,0,1] (sscatter (stranspose @[1,2,3,4,0] (w269 * sreshape @[4,7,11,2,2] (stranspose @[3,2,1,0] (ssum @4 (ssum @7 (ssum @11 (stranspose @[4,5,6,3,2,1,0] (ssum @3 (stranspose @[7,3,2,1,0,6,5,4] (sscatter (stranspose @[2,5,4,6,0,3,1] (sscatter (stranspose @[3,5,0,4,1,2] (sscatter (stranspose @[1,4,0,2,3] (ssum @4 (str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1))))))) * w317))) (\\[i318, i319] -> [i318, i318 + i319, i319]))) (\\[i320, i321] -> [i320, i321, i320 + i321]))) (\\[i322, i323, i324] -> [kfromS (smaxIndex (u271 !$ [i324, i322, i323])), i323, i322, i324]))))))))))) (\\[i325, i326, i327, i328] -> [kfromS (m256 !$ [i325, i327]), kfromS (m257 !$ [i326, i328])])) in tpair (tpair (tpair (rfromS (str (soneHot (ssum @1 (str (ssum @1 (str (ssum @23 (str (ssum @14 (str (w253 * sreshape @[4,14,23,1,1,3,4] (stranspose @[1,2,3,0] (sreplicate @12 t329))))))))))) [0]))) (rfromS (ssum @23 (ssum @14 (stranspose @[1,2,0] t329))))) (tpair (rfromS (ssum @11 (str (ssum @7 (str (w279 * w317)))))) (rfromS (ssum @11 (ssum @7 (stranspose @[1,2,0] t316)))))) (tpair (tpair (rfromS (str (m300 * m308))) (rfromS (ssum @7 (str m307)))) (tpair (rfromS (ssum @7 (stranspose @[2,1,0] (t305 * sreplicate @5 (sfromR dret))))) (rfromS (ssum @7 (str (sfromR dret))))))"+-- TODO: different test result with GHC 9.10:   printArtifactPretty (simplifyArtifact artifactRev)+--    @?= "\\dret u1 -> tfromS (STKProduct (STKProduct (STKProduct (STKR (SNat @4) STKScalar) (STKR (SNat @1) STKScalar)) (STKProduct (STKR (SNat @4) STKScalar) (STKR (SNat @1) STKScalar))) (STKProduct (STKProduct (STKR (SNat @2) STKScalar) (STKR (SNat @1) STKScalar)) (STKProduct (STKR (SNat @2) STKScalar) (STKR (SNat @1) STKScalar)))) (let u247 = sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i397, i398] -> [i397 + i398]))) (\\[i245, i246] -> [i245 + i246]) ; t249 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] u247))))) * str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (sfromR (tproject1 (tproject1 (tproject1 u1)))) !$ [0]))))))))))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (sfromR (tproject2 (tproject1 (tproject1 u1)))))) ; w259 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i250, i251, i252, i253, i254] -> [ifH (sscalar -0.0 <=. negate (t249 !$ [i250, 2 * i251 + i253, 2 * i252 + i254])) 0 1]) ; u261 = sreshape @[4,7,11,4] (w259 * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t249) (\\[i255, i256, i257, i258] -> [2 * i255 + i257, 2 * i256 + i258]))) ; w269 = sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u261))))))) (\\[i380, i382, i383] -> [kfromS (smaxIndex (u261 !$ [i383, i380, i382])), i382, i380, i383]))) (\\[i388, i390] -> [i388, i390, i388 + i390]))) (\\[i267, i268] -> [i267, i267 + i268, i268]) ; t270 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (sreplicate @4 (stranspose @[2,0,3,4,1] w269) * str (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (sfromR (tproject2 (tproject2 (tproject1 u1)))))) ; w280 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i271, i272, i273, i274, i275] -> [ifH (sscalar -0.0 <=. negate (t270 !$ [i271, 2 * i272 + i274, 2 * i273 + i275])) 0 1]) ; u282 = sreshape @[4,3,5,4] (w280 * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t270) (\\[i276, i277, i278, i279] -> [2 * i276 + i278, 2 * i277 + i279]))) ; v286 = sreshape @[60] (sgather u282 (\\[i283, i284, i285] -> [i283, i284, i285, kfromS (smaxIndex (u282 !$ [i283, i284, i285]))])) ; m287 = str (sreplicate @7 (sdot1In (sfromR (tproject1 (tproject1 (tproject2 u1)))) (sreplicate @5 v286))) + str (sreplicate @7 (sfromR (tproject2 (tproject1 (tproject2 u1))))) ; m290 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i288, i289] -> [ifH (sscalar -0.0 <=. negate (m287 !$ [i288, i289])) 0 1]) ; m293 = m290 * smatmul2 (str (sfromR (tproject1 (tproject2 (tproject2 u1))))) (sfromR dret) ; v294 = ssum @7 (str m293) ; t302 = sscatter (stranspose @[1,2,3,4,0] w280 * stranspose @[1,2,3,4,0] (sreshape @[4,3,5,2,2] (sscatter (sreshape @[4,3,5] (sdot1In (str (sfromR (tproject1 (tproject1 (tproject2 u1))))) (sreplicate @60 v294))) (\\[i295, i296, i297] -> [i295, i296, i297, kfromS (smaxIndex (u282 !$ [i295, i296, i297]))])))) (\\[i298, i299, i300, i301] -> [2 * i298 + i300, 2 * i299 + i301]) ; w303 = sreshape @[4,7,11,4,3,4] (stranspose @[1,2,3,0] (sreplicate @48 (stranspose @[2,0,1] t302))) ; t315 = sscatter (stranspose @[1,2,3,4,0] w259 * stranspose @[1,2,3,4,0] (sreshape @[4,7,11,2,2] (ssum @4 (ssum @7 (ssum @11 (ssum @3 (stranspose @[7,6,5,4,3,2,1,0] (sscatter (stranspose @[2,5,4,6,0,3,1] (sscatter (stranspose @[3,5,0,4,1,2] (sscatter (sdot1In (stranspose @[2,5,0,3,4,1] (sreplicate @7 (str (sreplicate @11 (sfromR (tproject1 (tproject2 (tproject1 u1)))))))) (stranspose @[2,5,1,3,4,0] w303)) (\\[i304, i305] -> [i304, i304 + i305, i305]))) (\\[i306, i307] -> [i306, i307, i306 + i307]))) (\\[i308, i309, i310] -> [kfromS (smaxIndex (u261 !$ [i310, i308, i309])), i309, i308, i310]))))))))) (\\[i311, i312, i313, i314] -> [2 * i311 + i313, 2 * i312 + i314]) in tpair (tpair (tpair (str (sreplicate @1 (ssum @23 (sdot1In (stranspose @[4,0,3,2,1] (sreplicate @4 (stranspose @[2,1,3,0] u247))) (stranspose @[3,4,2,0,5,6,1] (sreshape @[4,14,23,1,1,3,4] (stranspose @[1,2,3,0] (sreplicate @12 (stranspose @[2,0,1] t315)))) !$ [0, 0]))))) (ssum @23 (ssum @14 t315))) (tpair (ssum @11 (sdot1In (stranspose @[2,0,3,4,5,1] (sreplicate @4 (stranspose @[2,0,3,4,1] w269))) (stranspose @[2,0,3,4,5,1] w303))) (ssum @11 (ssum @7 t302)))) (tpair (tpair (sreplicate @5 v286 * str (sreplicate @60 v294)) (ssum @7 (str m293))) (tpair (smatmul2 (sfromR dret) (str m290 * str m287)) (ssum @7 (str (sfromR dret))))))"++testCNNOAst2 :: Assertion+testCNNOAst2 = do+  let batch_size = 7+      sizeMnistWidthI = 14+      sizeMnistHeightI = 23+      ftk = tftk @Concrete+                 (knownSTK @(X (MnistCnnRanked2.ADCnnMnistParameters+                                  Concrete Double)))+                 vals+      varName = mkAstVarName ftk Nothing . intToAstVarId $ 100000000+      var :: AstTensor AstMethodLet FullSpan+                       (X (MnistCnnRanked2.ADCnnMnistParameters+                             Concrete Double))+      var = AstVar varName+      valsInit :: MnistCnnRanked2.ADCnnMnistParameters Concrete Double+      valsInit =+        forgetShape $ fst+        $ randomValue @(MnistCnnRanked2.ADCnnMnistParametersShaped+                         Concrete 14 23  -- see sizeMnistWidthI, etc.+                         2 3 4 5 Double)+                      0.4 (mkStdGen 44)+      vals = toTarget @Concrete valsInit+      env = extendEnv varName vals emptyEnv+      blackGlyph = treplicate (SNat @7) knownSTK+                   $ treplicate (SNat @1) knownSTK+                   $ treplicate (SNat @14) knownSTK+                   $ treplicate (SNat @23) knownSTK $ rscalar 7+      afcnn2 :: ADReady f+             => MnistCnnRanked2.ADCnnMnistParameters f Double+             -> f (TKR 2 Double)+      afcnn2 = MnistCnnRanked2.convMnistTwoR+                 sizeMnistHeightI sizeMnistWidthI batch_size+                 (rconcrete $ unConcrete blackGlyph)+      afcnn1 = afcnn2 $ fromTarget var+  interpretAstFull @Concrete env afcnn1+    @?= afcnn2 valsInit+  interpretAstFull @Concrete env+                   (simplifyInline @(TKR 2 Double) afcnn1)+    @?= afcnn2 valsInit+  interpretAstFull @Concrete env+                   (simplifyInlineContract @(TKR 2 Double) afcnn1)+    @?= afcnn2 valsInit++testCNNOPP2S :: Assertion+testCNNOPP2S = do+  resetVarCounter+  let batch_size = SNat @7+      sizeMnistWidthI = SNat @14+      sizeMnistHeightI = SNat @23+      ftk = tftk @Concrete+                 (knownSTK @(X (MnistCnnShaped2.ADCnnMnistParametersShaped+                                  Concrete 14 23 2 3 4 5 Double)))+                 vals+      valsInit :: MnistCnnShaped2.ADCnnMnistParametersShaped Concrete 14 23 2 3 4 5 Double+      valsInit =+        fst+        $ randomValue @(MnistCnnShaped2.ADCnnMnistParametersShaped+                         Concrete 14 23  -- see sizeMnistWidthI, etc.+                         2 3 4 5 Double)+                      0.4 (mkStdGen 44)+      vals = toTarget @Concrete valsInit+      blackGlyph = treplicate batch_size knownSTK+                   $ treplicate (SNat @1) knownSTK+                   $ treplicate sizeMnistWidthI knownSTK+                   $ treplicate sizeMnistHeightI knownSTK $ sscalar 7+      afcnn2 :: ADReady f+             => MnistCnnShaped2.ADCnnMnistParametersShaped f 14 23 2 3 4 5 Double+             -> f (TKS '[SizeMnistLabel, 7] Double)+      afcnn2 = MnistCnnShaped2.convMnistTwoS+                 (SNat @2) (SNat @3) sizeMnistWidthI sizeMnistHeightI+                 (SNat @4) (SNat @5) batch_size+                 (sconcrete $ unConcrete blackGlyph)+      artifactRev = revArtifactAdapt UseIncomingCotangent afcnn2 ftk+  printArtifactPrimalPretty (simplifyArtifact artifactRev)+    @?= "\\u1 -> let t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i473, i474] -> [i473 + i474]))) (\\[i251, i252] -> [i251 + i252]))))))) * str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0]))))))))))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (tproject2 (tproject1 (tproject1 u1))))) ; u271 = sreshape @[4,7,11,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS (sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) !$ [i259, i261]), kfromS (sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) !$ [i260, i262])])) 0 1]) * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (sconcrete (sfromListLinear [7,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13]) !$ [i265, i267]), kfromS (sconcrete (sfromListLinear [11,2] [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]) !$ [i266, i268])]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i458, i460, i461] -> [kfromS (smaxIndex (u271 !$ [i461, i458, i460])), i460, i458, i461]))) (\\[i466, i468] -> [i466, i468, i466 + i468]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) * str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (tproject2 (tproject2 (tproject1 u1))))) ; u296 = sreshape @[4,3,5,4] (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i284, i286]), kfromS (sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) !$ [i285, i287])])) 0 1]) * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (sconcrete (sfromListLinear [3,2] [0,1,2,3,4,5]) !$ [i290, i292]), kfromS (sconcrete (sfromListLinear [5,2] [0,1,2,3,4,5,6,7,8,9]) !$ [i291, i293])]))) ; m301 = str (sreplicate @7 (sdot1In (tproject1 (tproject1 (tproject2 u1))) (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) in smatmul2 (tproject1 (tproject2 (tproject2 u1))) (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) * m301) + str (sreplicate @7 (tproject2 (tproject2 (tproject2 u1))))"+  printArtifactPrimalPretty artifactRev+    @?= "\\u1 -> let w253 = sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i249, i250] -> [i249 + i250]))) (\\[i251, i252] -> [i251 + i252]))))))) ; w254 = str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0])))))))) ; t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (w253 * w254))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (tproject2 (tproject1 (tproject1 u1))))) ; m256 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; m257 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; w269 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [let x263 = m256 !$ [i259, i261] ; x264 = m257 !$ [i260, i262] in ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS x263, kfromS x264])) 0 1]) ; w270 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (m256 !$ [i265, i267]), kfromS (m257 !$ [i266, i268])])) ; u271 = sreshape @[4,7,11,4] (w269 * w270) ; w279 = sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i272, i273, i274] -> [kfromS (smaxIndex (u271 !$ [i274, i272, i273])), i273, i272, i274]))) (\\[i275, i276] -> [i275, i276, i275 + i276]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (w279 * str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (tproject2 (tproject2 (tproject1 u1))))) ; m281 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m282 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; w294 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [let x288 = m281 !$ [i284, i286] ; x289 = m282 !$ [i285, i287] in ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS x288, kfromS x289])) 0 1]) ; w295 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (m281 !$ [i290, i292]), kfromS (m282 !$ [i291, i293])])) ; u296 = sreshape @[4,3,5,4] (w294 * w295) ; m300 = str (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))) ; m301 = str (sreplicate @7 (ssum @60 (str (tproject1 (tproject1 (tproject2 u1))) * m300))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) ; m304 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) ; t305 = str (sreplicate @10 (m304 * m301)) in ssum @5 (stranspose @[2,1,0] (sreplicate @7 (tproject1 (tproject2 (tproject2 u1)))) * t305) + str (sreplicate @7 (tproject2 (tproject2 (tproject2 u1))))"+  printArtifactPretty artifactRev+    @?= "\\dret u1 -> let w253 = sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] (sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i249, i250] -> [i249 + i250]))) (\\[i251, i252] -> [i251 + i252]))))))) ; w254 = str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0])))))))) ; t255 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (w253 * w254))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (tproject2 (tproject1 (tproject1 u1))))) ; m256 = str (sreplicate @2 (sconcrete (sreplicate [7] 2) * siota (SNat @7))) + sreplicate @7 (siota (SNat @2)) ; m257 = str (sreplicate @2 (sconcrete (sreplicate [11] 2) * siota (SNat @11))) + sreplicate @11 (siota (SNat @2)) ; w269 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i258, i259, i260, i261, i262] -> [let x263 = m256 !$ [i259, i261] ; x264 = m257 !$ [i260, i262] in ifH (sscalar -0.0 <=. negate (t255 !$ [i258, kfromS x263, kfromS x264])) 0 1]) ; w270 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t255) (\\[i265, i266, i267, i268] -> [kfromS (m256 !$ [i265, i267]), kfromS (m257 !$ [i266, i268])])) ; u271 = sreshape @[4,7,11,4] (w269 * w270) ; w279 = sreplicate @4 (stranspose @[2,0,3,4,1] (sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u271))))))) (\\[i272, i273, i274] -> [kfromS (smaxIndex (u271 !$ [i274, i272, i273])), i273, i272, i274]))) (\\[i275, i276] -> [i275, i276, i275 + i276]))) (\\[i277, i278] -> [i277, i277 + i278, i278]))) ; t280 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (w279 * str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (tproject2 (tproject2 (tproject1 u1))))) ; m281 = str (sreplicate @2 (sconcrete (sreplicate [3] 2) * siota (SNat @3))) + sreplicate @3 (siota (SNat @2)) ; m282 = str (sreplicate @2 (sconcrete (sreplicate [5] 2) * siota (SNat @5))) + sreplicate @5 (siota (SNat @2)) ; w294 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i283, i284, i285, i286, i287] -> [let x288 = m281 !$ [i284, i286] ; x289 = m282 !$ [i285, i287] in ifH (sscalar -0.0 <=. negate (t280 !$ [i283, kfromS x288, kfromS x289])) 0 1]) ; w295 = stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t280) (\\[i290, i291, i292, i293] -> [kfromS (m281 !$ [i290, i292]), kfromS (m282 !$ [i291, i293])])) ; u296 = sreshape @[4,3,5,4] (w294 * w295) ; m300 = str (sreplicate @5 (sreshape @[60] (sgather u296 (\\[i297, i298, i299] -> [i297, i298, i299, kfromS (smaxIndex (u296 !$ [i297, i298, i299]))])))) ; m301 = str (sreplicate @7 (ssum @60 (str (tproject1 (tproject1 (tproject2 u1))) * m300))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) ; m304 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i302, i303] -> [ifH (sscalar -0.0 <=. negate (m301 !$ [i302, i303])) 0 1]) ; t305 = str (sreplicate @10 (m304 * m301)) ; m307 = m304 * ssum @10 (str (stranspose @[2,1,0] (sreplicate @7 (tproject1 (tproject2 (tproject2 u1)))) * sreplicate @5 dret)) ; m308 = sreplicate @60 (ssum @7 (str m307)) ; t316 = stranspose @[2,0,1] (sscatter (stranspose @[1,2,3,4,0] (w294 * sreshape @[4,3,5,2,2] (sscatter (sreshape @[4,3,5] (ssum @5 (str (str (tproject1 (tproject1 (tproject2 u1))) * m308)))) (\\[i309, i310, i311] -> [i309, i310, i311, kfromS (smaxIndex (u296 !$ [i309, i310, i311]))])))) (\\[i312, i313, i314, i315] -> [kfromS (m281 !$ [i312, i314]), kfromS (m282 !$ [i313, i315])])) ; w317 = sreshape @[4,7,11,4,3,4] (stranspose @[1,2,3,0] (sreplicate @48 t316)) ; t329 = stranspose @[2,0,1] (sscatter (stranspose @[1,2,3,4,0] (w269 * sreshape @[4,7,11,2,2] (stranspose @[3,2,1,0] (ssum @4 (ssum @7 (ssum @11 (stranspose @[4,5,6,3,2,1,0] (ssum @3 (stranspose @[7,3,2,1,0,6,5,4] (sscatter (stranspose @[2,5,4,6,0,3,1] (sscatter (stranspose @[3,5,0,4,1,2] (sscatter (stranspose @[1,4,0,2,3] (ssum @4 (str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1)))))) * w317))) (\\[i318, i319] -> [i318, i318 + i319, i319]))) (\\[i320, i321] -> [i320, i321, i320 + i321]))) (\\[i322, i323, i324] -> [kfromS (smaxIndex (u271 !$ [i324, i322, i323])), i323, i322, i324]))))))))))) (\\[i325, i326, i327, i328] -> [kfromS (m256 !$ [i325, i327]), kfromS (m257 !$ [i326, i328])])) in tpair (tpair (tpair (str (soneHot (ssum @1 (str (ssum @1 (str (ssum @23 (str (ssum @14 (str (w253 * sreshape @[4,14,23,1,1,3,4] (stranspose @[1,2,3,0] (sreplicate @12 t329))))))))))) [0])) (ssum @23 (ssum @14 (stranspose @[1,2,0] t329)))) (tpair (ssum @11 (str (ssum @7 (str (w279 * w317))))) (ssum @11 (ssum @7 (stranspose @[1,2,0] t316))))) (tpair (tpair (str (m300 * m308)) (ssum @7 (str m307))) (tpair (ssum @7 (stranspose @[2,1,0] (t305 * sreplicate @5 dret))) (ssum @7 (str dret))))"+-- TODO: different test result with GHC 9.10: printArtifactPretty (simplifyArtifact artifactRev)+--    @?= "\\dret u1 -> let u247 = sgather (stranspose @[2,0,1] (sgather (sconcrete (sreplicate [14,23] 7.0)) (\\[i397, i398] -> [i397 + i398]))) (\\[i245, i246] -> [i245 + i246]) ; t249 = ssum @12 (stranspose @[3,0,1,2] (sreshape @[4,14,23,12] (sreplicate @4 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[1,2,0] (sreplicate @1 (stranspose @[2,0,3,1] u247))))) * str (sreplicate @14 (str (sreplicate @23 (str (sreplicate @1 (str (sreplicate @1 (str (tproject1 (tproject1 (tproject1 u1))) !$ [0]))))))))))) + stranspose @[2,0,1] (sreplicate @14 (sreplicate @23 (tproject2 (tproject1 (tproject1 u1))))) ; w259 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i250, i251, i252, i253, i254] -> [ifH (sscalar -0.0 <=. negate (t249 !$ [i250, 2 * i251 + i253, 2 * i252 + i254])) 0 1]) ; u261 = sreshape @[4,7,11,4] (w259 * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t249) (\\[i255, i256, i257, i258] -> [2 * i255 + i257, 2 * i256 + i258]))) ; w269 = sgather (stranspose @[2,4,5,0,3,1] (sgather (stranspose @[4,6,0,5,2,1,3] (sgather (stranspose @[4,3,2,1,7,6,5,0] (sreplicate @3 (stranspose @[6,5,4,3,0,1,2] (sreplicate @11 (sreplicate @7 (sreplicate @4 (stranspose @[3,2,1,0] u261))))))) (\\[i380, i382, i383] -> [kfromS (smaxIndex (u261 !$ [i383, i380, i382])), i382, i380, i383]))) (\\[i388, i390] -> [i388, i390, i388 + i390]))) (\\[i267, i268] -> [i267, i267 + i268, i268]) ; t270 = ssum @48 (stranspose @[3,0,1,2] (sreshape @[4,7,11,48] (sreplicate @4 (stranspose @[2,0,3,4,1] w269) * str (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))))) + stranspose @[2,0,1] (sreplicate @7 (sreplicate @11 (tproject2 (tproject2 (tproject1 u1))))) ; w280 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i271, i272, i273, i274, i275] -> [ifH (sscalar -0.0 <=. negate (t270 !$ [i271, 2 * i272 + i274, 2 * i273 + i275])) 0 1]) ; u282 = sreshape @[4,3,5,4] (w280 * stranspose @[4,0,1,2,3] (sgather (stranspose @[1,2,0] t270) (\\[i276, i277, i278, i279] -> [2 * i276 + i278, 2 * i277 + i279]))) ; v286 = sreshape @[60] (sgather u282 (\\[i283, i284, i285] -> [i283, i284, i285, kfromS (smaxIndex (u282 !$ [i283, i284, i285]))])) ; m287 = str (sreplicate @7 (sdot1In (tproject1 (tproject1 (tproject2 u1))) (sreplicate @5 v286))) + str (sreplicate @7 (tproject2 (tproject1 (tproject2 u1)))) ; m290 = sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i288, i289] -> [ifH (sscalar -0.0 <=. negate (m287 !$ [i288, i289])) 0 1]) ; m293 = m290 * smatmul2 (str (tproject1 (tproject2 (tproject2 u1)))) dret ; v294 = ssum @7 (str m293) ; t302 = sscatter (stranspose @[1,2,3,4,0] w280 * stranspose @[1,2,3,4,0] (sreshape @[4,3,5,2,2] (sscatter (sreshape @[4,3,5] (sdot1In (str (tproject1 (tproject1 (tproject2 u1)))) (sreplicate @60 v294))) (\\[i295, i296, i297] -> [i295, i296, i297, kfromS (smaxIndex (u282 !$ [i295, i296, i297]))])))) (\\[i298, i299, i300, i301] -> [2 * i298 + i300, 2 * i299 + i301]) ; w303 = sreshape @[4,7,11,4,3,4] (stranspose @[1,2,3,0] (sreplicate @48 (stranspose @[2,0,1] t302))) ; t315 = sscatter (stranspose @[1,2,3,4,0] w259 * stranspose @[1,2,3,4,0] (sreshape @[4,7,11,2,2] (ssum @4 (ssum @7 (ssum @11 (ssum @3 (stranspose @[7,6,5,4,3,2,1,0] (sscatter (stranspose @[2,5,4,6,0,3,1] (sscatter (stranspose @[3,5,0,4,1,2] (sscatter (sdot1In (stranspose @[2,5,0,3,4,1] (sreplicate @7 (str (sreplicate @11 (tproject1 (tproject2 (tproject1 u1))))))) (stranspose @[2,5,1,3,4,0] w303)) (\\[i304, i305] -> [i304, i304 + i305, i305]))) (\\[i306, i307] -> [i306, i307, i306 + i307]))) (\\[i308, i309, i310] -> [kfromS (smaxIndex (u261 !$ [i310, i308, i309])), i309, i308, i310]))))))))) (\\[i311, i312, i313, i314] -> [2 * i311 + i313, 2 * i312 + i314]) in tpair (tpair (tpair (str (sreplicate @1 (ssum @23 (sdot1In (stranspose @[4,0,3,2,1] (sreplicate @4 (stranspose @[2,1,3,0] u247))) (stranspose @[3,4,2,0,5,6,1] (sreshape @[4,14,23,1,1,3,4] (stranspose @[1,2,3,0] (sreplicate @12 (stranspose @[2,0,1] t315)))) !$ [0, 0]))))) (ssum @23 (ssum @14 t315))) (tpair (ssum @11 (sdot1In (stranspose @[2,0,3,4,5,1] (sreplicate @4 (stranspose @[2,0,3,4,1] w269))) (stranspose @[2,0,3,4,5,1] w303))) (ssum @11 (ssum @7 t302)))) (tpair (tpair (sreplicate @5 v286 * str (sreplicate @60 v294)) (ssum @7 (str m293))) (tpair (smatmul2 dret (str m290 * str m287)) (ssum @7 (str dret))))"
+ test/simplified/TestMnistRNNR.hs view
@@ -0,0 +1,427 @@+{-# LANGUAGE OverloadedLists #-}+-- | Tests of "MnistRnnRanked2" recurrent neural networks using a few different+-- optimization pipelines.+--+-- Not LSTM.+-- Doesn't train without Adam, regardless of whether mini-batches used. It does+-- train with Adam, but only after very carefully tweaking initialization.+-- This is extremely sensitive to initial parameters, more than to anything+-- else. Probably, gradient is vanishing if parameters are initialized+-- with a probability distribution that doesn't have the right variance. See+-- https://stats.stackexchange.com/questions/301285/what-is-vanishing-gradient.+-- Regularization/normalization might help as well.+module TestMnistRNNR+  ( testTrees+  ) where++import Prelude++import Control.Monad (foldM, unless)+import System.IO (hPutStrLn, stderr)+import System.Random+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)+import Text.Printf++import Data.Array.Nested.Ranked.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstInterpret++import EqEpsilon++import MnistData+import MnistRnnRanked2 (ADRnnMnistParameters, ADRnnMnistParametersShaped)+import MnistRnnRanked2 qualified++-- TODO: optimize enough that it can run for one full epoch in reasonable time+-- and then verify it trains down to ~20% validation error in a short enough+-- time to include such a training run in tests.++testTrees :: [TestTree]+testTrees = [ tensorADValMnistTestsRNNRA+            , tensorADValMnistTestsRNNRI+            , tensorADValMnistTestsRNNRO+            ]++-- POPL differentiation, straight via the ADVal instance of RankedTensor,+-- which side-steps vectorization.+mnistTestCaseRNNRA+  :: forall r.+     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+  => String+  -> Int -> Int -> Int -> Int -> Int -> r+  -> TestTree+mnistTestCaseRNNRA prefix epochs maxBatches width miniBatchSize totalBatchSize+                   expected =+  withSNat width $ \(SNat @width) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue @(Concrete (X (ADRnnMnistParametersShaped+                                       Concrete width r)))+                      0.23 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show width, show miniBatchSize+                        , show $ widthSTK+                          $ knownSTK @(X (ADRnnMnistParameters Concrete r))+                        , show (tsize knownSTK targetInit) ]+      ftest :: Int -> MnistDataBatchR r+            -> Concrete (X (ADRnnMnistParameters Concrete r))+            -> r+      ftest batch_size mnistData pars =+        MnistRnnRanked2.rnnMnistTestR+          batch_size mnistData (fromTarget @Concrete pars)+  in testCase name $ do+       hPutStrLn stderr $+         printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+                prefix epochs maxBatches+       trainData <- map mkMnistDataR+                    <$> loadMnistData trainGlyphsPath trainLabelsPath+       testData <- map mkMnistDataR . take (totalBatchSize * maxBatches)+                   <$> loadMnistData testGlyphsPath testLabelsPath+       let testDataR = mkMnistDataBatchR testData+           f :: MnistDataBatchR r+             -> ADVal Concrete (X (ADRnnMnistParameters Concrete r))+             -> ADVal Concrete (TKScalar r)+           f (glyphR, labelR) adinputs =+             MnistRnnRanked2.rnnMnistLossFusedR+               miniBatchSize (rconcrete glyphR, rconcrete labelR)+               (fromTarget @(ADVal Concrete) adinputs)+           runBatch :: ( Concrete (X (ADRnnMnistParameters Concrete r))+                       , StateAdam (X (ADRnnMnistParameters Concrete r)) )+                    -> (Int, [MnistDataR r])+                    -> IO ( Concrete (X (ADRnnMnistParameters Concrete r))+                          , StateAdam (X (ADRnnMnistParameters Concrete r)) )+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkR = map mkMnistDataBatchR+                          $ filter (\ch -> length ch == miniBatchSize)+                          $ chunksOf miniBatchSize chunk+                 res@(parameters2, _) =+                   sgdAdam @(MnistDataBatchR r)+                               @(X (ADRnnMnistParameters Concrete r))+                               f chunkR parameters stateAdam+                 trainScore =+                   ftest (length chunk) (mkMnistDataBatchR chunk) parameters2+                 testScore =+                   ftest ((totalBatchSize * maxBatches) `min` 10000)+                         testDataR parameters2+                 lenChunk = length chunk+             unless (width < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> ( Concrete (X (ADRnnMnistParameters Concrete r))+                       , StateAdam (X (ADRnnMnistParameters Concrete r)) )+                    -> IO (Concrete (X (ADRnnMnistParameters Concrete r)))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (width < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+           ftk = tftk @Concrete+                      (knownSTK @(X (ADRnnMnistParameters Concrete r)))+                      targetInit+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal =+             1 - ftest ((totalBatchSize * maxBatches) `min` 10000)+                       testDataR res+       testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCaseRNNRA+  :: String+  -> Int -> Int -> Int -> Int -> Int -> Double+  -> TestTree #-}++tensorADValMnistTestsRNNRA :: TestTree+tensorADValMnistTestsRNNRA = testGroup "RNNR ADVal MNIST tests"+  [ mnistTestCaseRNNRA "RNNRA 1 epoch, 1 batch" 1 1 128 150 5000+                       (0.6026 :: Double)+  , mnistTestCaseRNNRA "RNNRA artificial 1 2 3 4 5" 2 3 4 5 50+                       (0.8933333 :: Float)+  , mnistTestCaseRNNRA "RNNRA artificial 5 4 3 2 1" 5 4 3 2 49+                       (0.8622448979591837 :: Double)+  , mnistTestCaseRNNRA "RNNRA 1 epoch, 0 batch" 1 0 128 150 50+                       (1.0 :: Float)+  ]++-- POPL differentiation, with Ast term defined and vectorized only once,+-- but differentiated anew in each gradient descent iteration.+mnistTestCaseRNNRI+  :: forall r.+     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+  => String+  -> Int -> Int -> Int -> Int -> Int -> r+  -> TestTree+mnistTestCaseRNNRI prefix epochs maxBatches width miniBatchSize totalBatchSize+                   expected =+  withSNat width $ \(SNat @width) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue @(Concrete (X (ADRnnMnistParametersShaped+                                       Concrete width r)))+                      0.23 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show width, show miniBatchSize+                        , show $ widthSTK+                          $ knownSTK @(X (ADRnnMnistParameters Concrete r))+                        , show (tsize knownSTK targetInit) ]+      ftest :: Int -> MnistDataBatchR r+            -> Concrete (X (ADRnnMnistParameters Concrete r))+            -> r+      ftest batch_size mnistData pars =+        MnistRnnRanked2.rnnMnistTestR+          batch_size mnistData (fromTarget @Concrete pars)+  in testCase name $ do+       hPutStrLn stderr $+         printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+                prefix epochs maxBatches+       trainData <- map mkMnistDataR+                    <$> loadMnistData trainGlyphsPath trainLabelsPath+       testData <- map mkMnistDataR . take (totalBatchSize * maxBatches)+                   <$> loadMnistData testGlyphsPath testLabelsPath+       let testDataR = mkMnistDataBatchR testData+           ftk = tftk @Concrete+                      (knownSTK @(X (ADRnnMnistParameters Concrete r)))+                      targetInit+       (_, _, var, varAst) <- funToAstRevIO ftk+       (varGlyph, astGlyph) <-+         funToAstIO (FTKR (miniBatchSize+                           :$: sizeMnistHeightInt+                           :$: sizeMnistWidthInt+                           :$: ZSR) FTKScalar) id+       (varLabel, astLabel) <-+         funToAstIO (FTKR (miniBatchSize+                           :$: sizeMnistLabelInt+                           :$: ZSR) FTKScalar) id+       let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)+           ast = simplifyInline+                 $ MnistRnnRanked2.rnnMnistLossFusedR+                     miniBatchSize (astGlyph, astLabel)+                     (fromTarget varAst)+           f :: MnistDataBatchR r+              -> ADVal Concrete (X (ADRnnMnistParameters Concrete r))+              -> ADVal Concrete (TKScalar r)+           f (glyph, label) varInputs =+             let env = extendEnv var varInputs emptyEnv+                 envMnist = extendEnv varGlyph (rconcrete glyph)+                            $ extendEnv varLabel (rconcrete label) env+             in interpretAstFull envMnist ast+           runBatch :: ( Concrete (X (ADRnnMnistParameters Concrete r))+                       , StateAdam (X (ADRnnMnistParameters Concrete r)) )+                    -> (Int, [MnistDataR r])+                    -> IO ( Concrete (X (ADRnnMnistParameters Concrete r))+                          , StateAdam (X (ADRnnMnistParameters Concrete r)) )+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkR = map mkMnistDataBatchR+                          $ filter (\ch -> length ch == miniBatchSize)+                          $ chunksOf miniBatchSize chunk+                 res@(parameters2, _) =+                   sgdAdam @(MnistDataBatchR r)+                               @(X (ADRnnMnistParameters Concrete r))+                               f chunkR parameters stateAdam+                 trainScore =+                   ftest (length chunk) (mkMnistDataBatchR chunk) parameters2+                 testScore =+                   ftest ((totalBatchSize * maxBatches) `min` 10000)+                         testDataR parameters2+                 lenChunk = length chunk+             unless (width < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> ( Concrete (X (ADRnnMnistParameters Concrete r))+                       , StateAdam (X (ADRnnMnistParameters Concrete r)) )+                    -> IO (Concrete (X (ADRnnMnistParameters Concrete r)))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (width < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal =+             1 - ftest ((totalBatchSize * maxBatches) `min` 10000)+                       testDataR res+       testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCaseRNNRI+  :: String+  -> Int -> Int -> Int -> Int -> Int -> Double+  -> TestTree #-}++tensorADValMnistTestsRNNRI :: TestTree+tensorADValMnistTestsRNNRI = testGroup "RNNR Intermediate MNIST tests"+  [ mnistTestCaseRNNRI "RNNRI 1 epoch, 1 batch" 1 1 128 150 5000+                       (0.6026 :: Double)+  , mnistTestCaseRNNRI "RNNRI artificial 1 2 3 4 5" 2 3 4 5 50+                       (0.8933333 :: Float)+  , mnistTestCaseRNNRI "RNNRI artificial 5 4 3 2 1" 5 4 3 2 49+                       (0.8622448979591837 :: Double)+  , mnistTestCaseRNNRI "RNNRI 1 epoch, 0 batch" 1 0 128 150 50+                       (1.0 :: Float)+  ]++-- JAX differentiation, Ast term built and differentiated only once+-- and the result interpreted with different inputs in each gradient+-- descent iteration.+mnistTestCaseRNNRO+  :: forall r.+     ( Differentiable r, GoodScalar r+     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+  => String+  -> Int -> Int -> Int -> Int -> Int -> r+  -> TestTree+mnistTestCaseRNNRO prefix epochs maxBatches width miniBatchSize totalBatchSize+                   expected =+  withSNat width $ \(SNat @width) ->+  let targetInit =+        forgetShape $ fst+        $ randomValue @(Concrete (X (ADRnnMnistParametersShaped+                                       Concrete width r)))+                      0.23 (mkStdGen 44)+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show width, show miniBatchSize+                        , show $ widthSTK+                          $ knownSTK @(X (ADRnnMnistParameters Concrete r))+                        , show (tsize knownSTK targetInit) ]+      ftest :: Int -> MnistDataBatchR r+            -> Concrete (X (ADRnnMnistParameters Concrete r))+            -> r+      ftest batch_size mnistData pars =+        MnistRnnRanked2.rnnMnistTestR+          batch_size mnistData (fromTarget @Concrete pars)+  in testCase name $ do+       hPutStrLn stderr $+         printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+                prefix epochs maxBatches+       trainData <- map mkMnistDataR+                    <$> loadMnistData trainGlyphsPath trainLabelsPath+       testData <- map mkMnistDataR . take (totalBatchSize * maxBatches)+                   <$> loadMnistData testGlyphsPath testLabelsPath+       let testDataR = mkMnistDataBatchR testData+           dataInit = case chunksOf miniBatchSize testData of+             d : _ -> let (dglyph, dlabel) = mkMnistDataBatchR d+                      in (rconcrete dglyph, rconcrete dlabel)+             [] -> error "empty test data"+           f :: ( ADRnnMnistParameters (AstTensor AstMethodLet FullSpan) r+                , ( AstTensor AstMethodLet FullSpan (TKR 3 r)+                  , AstTensor AstMethodLet FullSpan (TKR 2 r) ) )+             -> AstTensor AstMethodLet FullSpan (TKScalar r)+           f = \ (pars, (glyphR, labelR)) ->+             MnistRnnRanked2.rnnMnistLossFusedR+               miniBatchSize (rprimalPart glyphR, rprimalPart labelR) pars+           artRaw = gradArtifact f (fromTarget targetInit, dataInit)+           art = simplifyArtifactGradient artRaw+           go :: [MnistDataBatchR r]+              -> ( Concrete (X (ADRnnMnistParameters Concrete r))+                 , StateAdam (X (ADRnnMnistParameters Concrete r)) )+              -> ( Concrete (X (ADRnnMnistParameters Concrete r))+                 , StateAdam (X (ADRnnMnistParameters Concrete r)) )+           go [] (parameters, stateAdam) = (parameters, stateAdam)+           go ((glyph, label) : rest) (!parameters, !stateAdam) =+             let parametersAndInput =+                   tpair parameters (tpair (rconcrete glyph) (rconcrete label))+                 gradient = tproject1 $ fst+                            $ revInterpretArtifact+                                art parametersAndInput Nothing+             in go rest (updateWithGradientAdam+                           @(X (ADRnnMnistParameters Concrete r))+                           defaultArgsAdam stateAdam knownSTK parameters+                           gradient)+           runBatch :: ( Concrete (X (ADRnnMnistParameters Concrete r))+                       , StateAdam (X (ADRnnMnistParameters Concrete r)) )+                    -> (Int, [MnistDataR r])+                    -> IO ( Concrete (X (ADRnnMnistParameters Concrete r))+                          , StateAdam (X (ADRnnMnistParameters Concrete r)) )+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkR = map mkMnistDataBatchR+                          $ filter (\ch -> length ch == miniBatchSize)+                          $ chunksOf miniBatchSize chunk+                 res@(parameters2, _) = go chunkR (parameters, stateAdam)+                 trainScore =+                   ftest (length chunk) (mkMnistDataBatchR chunk) parameters2+                 testScore =+                   ftest ((totalBatchSize * maxBatches) `min` 10000)+                         testDataR parameters2+                 lenChunk = length chunk+             unless (width < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> ( Concrete (X (ADRnnMnistParameters Concrete r))+                       , StateAdam (X (ADRnnMnistParameters Concrete r)) )+                    -> IO (Concrete (X (ADRnnMnistParameters Concrete r)))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (width < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+           ftk = tftk @Concrete (knownSTK @(X (ADRnnMnistParameters+                                                 Concrete r)))+                      targetInit+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal =+             1 - ftest ((totalBatchSize * maxBatches) `min` 10000)+                       testDataR res+       assertEqualUpToEpsilon 1e-1 expected testErrorFinal++{-# SPECIALIZE mnistTestCaseRNNRO+  :: String+  -> Int -> Int -> Int -> Int -> Int -> Double+  -> TestTree #-}++tensorADValMnistTestsRNNRO :: TestTree+tensorADValMnistTestsRNNRO = testGroup "RNNR Once MNIST tests"+  [ mnistTestCaseRNNRO "RNNRO 1 epoch, 1 batch" 1 1 128 150 5000+                       (0.6026 :: Double)+  , mnistTestCaseRNNRO "RNNRO artificial 1 2 3 4 5" 2 3 4 5 50+                       (0.8933333 :: Float)+  , mnistTestCaseRNNRO "RNNRO artificial 5 4 3 2 1" 5 4 3 2 49+                       (0.8928571428571429 :: Double)+  , mnistTestCaseRNNRO "RNNRO 1 epoch, 0 batch" 1 0 128 150 50+                       (1.0 :: Float)+  ]
+ test/simplified/TestMnistRNNS.hs view
@@ -0,0 +1,415 @@+-- | Tests of "MnistRnnShaped2" recurrent neural networks using a few different+-- optimization pipelines.+--+-- Not LSTM.+-- Doesn't train without Adam, regardless of whether mini-batches used. It does+-- train with Adam, but only after very carefully tweaking initialization.+-- This is extremely sensitive to initial parameters, more than to anything+-- else. Probably, gradient is vanishing if parameters are initialized+-- with a probability distribution that doesn't have the right variance. See+-- https://stats.stackexchange.com/questions/301285/what-is-vanishing-gradient.+-- Regularization/normalization might help as well.+module TestMnistRNNS+  ( testTrees+  ) where++import Prelude++import Control.Monad (foldM, unless)+import Data.Proxy (Proxy (Proxy))+import Data.Type.Equality ((:~:) (Refl))+import GHC.TypeLits (KnownNat, sameNat)+import System.IO (hPutStrLn, stderr)+import System.Random+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)+import Text.Printf++import Data.Array.Nested.Shaped.Shape++import HordeAd+import HordeAd.Core.Adaptor+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstInterpret++import EqEpsilon++import MnistData+import MnistRnnShaped2 (ADRnnMnistParametersShaped)+import MnistRnnShaped2 qualified++-- TODO: optimize enough that it can run for one full epoch in reasonable time+-- and then verify it trains down to ~20% validation error in a short enough+-- time to include such a training run in tests.++type XParams out_width r =+ X (ADRnnMnistParametersShaped Concrete SizeMnistHeight out_width r)++testTrees :: [TestTree]+testTrees = [ tensorADValMnistTestsRNNSA+            , tensorADValMnistTestsRNNSI+            , tensorADValMnistTestsRNNSO+            ]++-- POPL differentiation, straight via the ADVal instance of RankedTensor,+-- which side-steps vectorization.+mnistTestCaseRNNSA+  :: forall width batch_size r.+     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+  => String+  -> Int -> Int -> SNat width -> SNat batch_size -> Int -> r+  -> TestTree+mnistTestCaseRNNSA prefix epochs maxBatches width@SNat batch_size@SNat+                   totalBatchSize expected =+  let targetInit =+        fst $ randomValue @(Concrete (XParams width r)) 0.23 (mkStdGen 44)+      miniBatchSize = sNatValue batch_size+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show (sNatValue width), show miniBatchSize+                        , show $ widthSTK+                          $ knownSTK @(XParams width r)+                        , show (tsize knownSTK targetInit) ]+      ftest :: forall batch_size2. KnownNat batch_size2+            => MnistDataBatchS batch_size2 r -> Concrete (XParams width r)+            -> r+      ftest _ _ | Just Refl <- sameNat (Proxy @0) (Proxy @batch_size2) = 0+      ftest mnistData testParams =+        MnistRnnShaped2.rnnMnistTestS+          width (SNat @batch_size2) mnistData (fromTarget @Concrete testParams)+  in testCase name $ do+    hPutStrLn stderr $+      printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+             prefix epochs maxBatches+    trainData <- map mkMnistDataS+                 <$> loadMnistData trainGlyphsPath trainLabelsPath+    testData <- map mkMnistDataS . take (totalBatchSize * maxBatches)+                <$> loadMnistData testGlyphsPath testLabelsPath+    withSNat ((totalBatchSize * maxBatches) `min` 10000)+     $ \(SNat @lenTestData) -> do+       let testDataS = mkMnistDataBatchS @lenTestData testData+           f :: MnistDataBatchS batch_size r+             -> ADVal Concrete (XParams width r)+             -> ADVal Concrete (TKScalar r)+           f (glyphS, labelS) adinputs =+             MnistRnnShaped2.rnnMnistLossFusedS+               width batch_size (sconcrete glyphS, sconcrete labelS)+               (fromTarget @(ADVal Concrete) adinputs)+           runBatch :: ( Concrete (XParams width r)+                       , StateAdam (XParams width r) )+                    -> (Int, [MnistDataS r])+                    -> IO ( Concrete (XParams width r)+                          , StateAdam (XParams width r) )+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkS = map mkMnistDataBatchS+                          $ filter (\ch -> length ch == miniBatchSize)+                          $ chunksOf miniBatchSize chunk+                 res@(parameters2, _) =+                   sgdAdam @(MnistDataBatchS batch_size r)+                               @(XParams width r)+                               f chunkS parameters stateAdam+                 trainScore = withSNat (length chunk) $ \(SNat @len) ->+                   ftest @len (mkMnistDataBatchS chunk) parameters2+                 testScore = ftest @lenTestData testDataS parameters2+                 lenChunk = length chunk+             unless (sNatValue width < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> ( Concrete (XParams width r)+                       , StateAdam (XParams width r) )+                    -> IO (Concrete (XParams width r))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (sNatValue width < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+           ftk = tftk @Concrete (knownSTK @(XParams width r))+                      targetInit+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal = 1 - ftest @lenTestData testDataS res+       testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCaseRNNSA+  :: String+  -> Int -> Int -> SNat width -> SNat batch_size -> Int -> Double+  -> TestTree #-}++tensorADValMnistTestsRNNSA :: TestTree+tensorADValMnistTestsRNNSA = testGroup "RNNS ADVal MNIST tests"+  [ mnistTestCaseRNNSA "RNNSA 1 epoch, 1 batch" 1 1 (SNat @128) (SNat @150) 5000+                       (0.6026 :: Double)+  , mnistTestCaseRNNSA "RNNSA artificial 1 2 3 4 5" 2 3 (SNat @4) (SNat @5) 50+                       (0.8933333 :: Float)+  , mnistTestCaseRNNSA "RNNSA artificial 5 4 3 2 1" 5 4 (SNat @3) (SNat @2) 49+                       (0.8622448979591837 :: Double)+  , mnistTestCaseRNNSA "RNNSA 1 epoch, 0 batch" 1 0 (SNat @128) (SNat @150) 50+                       (1.0 :: Float)+  ]++-- POPL differentiation, with Ast term defined and vectorized only once,+-- but differentiated anew in each gradient descent iteration.+mnistTestCaseRNNSI+  :: forall width batch_size r.+     (Differentiable r, GoodScalar r, PrintfArg r, AssertEqualUpToEpsilon r)+  => String+  -> Int -> Int -> SNat width -> SNat batch_size -> Int -> r+  -> TestTree+mnistTestCaseRNNSI prefix epochs maxBatches width@SNat batch_size@SNat+                   totalBatchSize expected =+  let targetInit =+        fst $ randomValue @(Concrete (XParams width r)) 0.23 (mkStdGen 44)+      miniBatchSize = sNatValue batch_size+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show (sNatValue width), show miniBatchSize+                        , show $ widthSTK+                          $ knownSTK @(XParams width r)+                        , show (tsize knownSTK targetInit) ]+      ftest :: forall batch_size2. KnownNat batch_size2+            => MnistDataBatchS batch_size2 r -> Concrete (XParams width r)+            -> r+      ftest _ _ | Just Refl <- sameNat (Proxy @0) (Proxy @batch_size2) = 0+      ftest mnistData testParams =+        MnistRnnShaped2.rnnMnistTestS+          width (SNat @batch_size2) mnistData (fromTarget @Concrete testParams)+  in testCase name $ do+    hPutStrLn stderr $+      printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+             prefix epochs maxBatches+    trainData <- map mkMnistDataS+                 <$> loadMnistData trainGlyphsPath trainLabelsPath+    testData <- map mkMnistDataS . take (totalBatchSize * maxBatches)+                <$> loadMnistData testGlyphsPath testLabelsPath+    withSNat ((totalBatchSize * maxBatches) `min` 10000)+     $ \(SNat @lenTestData) -> do+       let testDataS = mkMnistDataBatchS @lenTestData testData+           ftk = tftk @Concrete (knownSTK @(XParams width r)) targetInit+       (_, _, var, varAst) <- funToAstRevIO ftk+       (varGlyph, astGlyph) <- funToAstIO (FTKS knownShS FTKScalar) id+       (varLabel, astLabel) <- funToAstIO (FTKS knownShS FTKScalar) id+       let ast :: AstTensor AstMethodLet FullSpan (TKScalar r)+           ast = simplifyInline+                 $ MnistRnnShaped2.rnnMnistLossFusedS+                     width batch_size (astGlyph, astLabel)+                     (fromTarget varAst)+           f :: MnistDataBatchS batch_size r+             -> ADVal Concrete (XParams width r)+             -> ADVal Concrete (TKScalar r)+           f (glyph, label) varInputs =+             let env = extendEnv var varInputs emptyEnv+                 envMnist = extendEnv varGlyph (sconcrete glyph)+                            $ extendEnv varLabel (sconcrete label) env+             in interpretAstFull envMnist ast+           runBatch :: ( Concrete (XParams width r)+                       , StateAdam (XParams width r) )+                    -> (Int, [MnistDataS r])+                    -> IO ( Concrete (XParams width r)+                          , StateAdam (XParams width r) )+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkS = map mkMnistDataBatchS+                          $ filter (\ch -> length ch == miniBatchSize)+                          $ chunksOf miniBatchSize chunk+                 res@(parameters2, _) =+                   sgdAdam @(MnistDataBatchS batch_size r)+                               @(XParams width r)+                               f chunkS parameters stateAdam+                 trainScore = withSNat (length chunk) $ \(SNat @len) ->+                   ftest @len (mkMnistDataBatchS chunk) parameters2+                 testScore = ftest @lenTestData testDataS parameters2+                 lenChunk = length chunk+             unless (sNatValue width < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> ( Concrete (XParams width r)+                       , StateAdam (XParams width r) )+                    -> IO (Concrete (XParams width r))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (sNatValue width < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal = 1 - ftest @lenTestData testDataS res+       testErrorFinal @?~ expected++{-# SPECIALIZE mnistTestCaseRNNSI+  :: String+  -> Int -> Int -> SNat width -> SNat batch_size -> Int -> Double+  -> TestTree #-}++tensorADValMnistTestsRNNSI :: TestTree+tensorADValMnistTestsRNNSI = testGroup "RNNS Intermediate MNIST tests"+  [ mnistTestCaseRNNSI "RNNSI 1 epoch, 1 batch" 1 1 (SNat @128) (SNat @150) 5000+                       (0.6026 :: Double)+  , mnistTestCaseRNNSI "RNNSI artificial 1 2 3 4 5" 2 3 (SNat @4) (SNat @5) 50+                       (0.8933333 :: Float)+  , mnistTestCaseRNNSI "RNNSI artificial 5 4 3 2 1" 5 4 (SNat @3) (SNat @2) 49+                       (0.8622448979591837 :: Double)+  , mnistTestCaseRNNSI "RNNSI 1 epoch, 0 batch" 1 0 (SNat @128) (SNat @150) 50+                       (1.0 :: Float)+  ]++-- JAX differentiation, Ast term built and differentiated only once+-- and the result interpreted with different inputs in each gradient+-- descent iteration.+mnistTestCaseRNNSO+  :: forall width batch_size r.+     ( Differentiable r, GoodScalar r+     , PrintfArg r, AssertEqualUpToEpsilon r, ADTensorScalar r ~ r )+  => String+  -> Int -> Int -> SNat width -> SNat batch_size -> Int -> r+  -> TestTree+mnistTestCaseRNNSO prefix epochs maxBatches width@SNat batch_size@SNat+                   totalBatchSize expected =+  let targetInit =+        fst $ randomValue @(Concrete (XParams width r)) 0.23 (mkStdGen 44)+      miniBatchSize = sNatValue batch_size+      name = prefix ++ ": "+             ++ unwords [ show epochs, show maxBatches+                        , show (sNatValue width), show miniBatchSize+                        , show $ widthSTK+                          $ knownSTK @(XParams width r)+                        , show (tsize knownSTK targetInit) ]+      ftest :: forall batch_size2. KnownNat batch_size2+            => MnistDataBatchS batch_size2 r -> Concrete (XParams width r)+            -> r+      ftest _ _ | Just Refl <- sameNat (Proxy @0) (Proxy @batch_size2) = 0+      ftest mnistData testParams =+        MnistRnnShaped2.rnnMnistTestS+          width (SNat @batch_size2) mnistData (fromTarget @Concrete testParams)+  in testCase name $ do+    hPutStrLn stderr $+      printf "\n%s: Epochs to run/max batches per epoch: %d/%d"+             prefix epochs maxBatches+    trainData <- map mkMnistDataS+                 <$> loadMnistData trainGlyphsPath trainLabelsPath+    testData <- map mkMnistDataS . take (totalBatchSize * maxBatches)+                <$> loadMnistData testGlyphsPath testLabelsPath+    withSNat ((totalBatchSize * maxBatches) `min` 10000)+     $ \(SNat @lenTestData) -> do+       let testDataS = mkMnistDataBatchS @lenTestData testData+           ftk = tftk @Concrete (knownSTK @(XParams width r)) targetInit+           ftkData = FTKProduct (FTKS (batch_size+                                       :$$ sizeMnistHeight+                                       :$$ sizeMnistWidth+                                       :$$ ZSS) FTKScalar)+                                (FTKS (batch_size+                                       :$$ sizeMnistLabel+                                       :$$ ZSS) FTKScalar)+           f :: ( ADRnnMnistParametersShaped (AstTensor AstMethodLet FullSpan)+                    SizeMnistHeight width r+                , ( AstTensor AstMethodLet FullSpan+                      (TKS '[batch_size, SizeMnistHeight, SizeMnistWidth] r)+                  , AstTensor AstMethodLet FullSpan+                      (TKS '[batch_size, SizeMnistLabel] r) ) )+             -> AstTensor AstMethodLet FullSpan (TKScalar r)+           f = \ (pars, (glyphS, labelS)) ->+             MnistRnnShaped2.rnnMnistLossFusedS+               width batch_size (sprimalPart glyphS, sprimalPart labelS) pars+           artRaw = revArtifactAdapt IgnoreIncomingCotangent+                                     f (FTKProduct ftk ftkData)+           art = simplifyArtifactGradient artRaw+           go :: [MnistDataBatchS batch_size r]+              -> ( Concrete (XParams width r)+                 , StateAdam (XParams width r) )+              -> ( Concrete (XParams width r)+                 , StateAdam (XParams width r) )+           go [] (parameters, stateAdam) = (parameters, stateAdam)+           go ((glyph, label) : rest) (!parameters, !stateAdam) =+             let parametersAndInput =+                   tpair parameters (tpair (sconcrete glyph) (sconcrete label))+                 gradient = tproject1 $ fst+                            $ revInterpretArtifact+                                art parametersAndInput Nothing+             in go rest (updateWithGradientAdam+                           @(XParams width r)+                           defaultArgsAdam stateAdam knownSTK parameters+                           gradient)+           runBatch :: ( Concrete (XParams width r)+                       , StateAdam (XParams width r) )+                    -> (Int, [MnistDataS r])+                    -> IO ( Concrete (XParams width r)+                          , StateAdam (XParams width r) )+           runBatch (!parameters, !stateAdam) (k, chunk) = do+             let chunkS = map mkMnistDataBatchS+                          $ filter (\ch -> length ch == miniBatchSize)+                          $ chunksOf miniBatchSize chunk+                 res@(parameters2, _) = go chunkS (parameters, stateAdam)+                 trainScore = withSNat (length chunk) $ \(SNat @len) ->+                   ftest @len (mkMnistDataBatchS chunk) parameters2+                 testScore = ftest @lenTestData testDataS parameters2+                 lenChunk = length chunk+             unless (sNatValue width < 10) $ do+               hPutStrLn stderr $+                 printf "\n%s: (Batch %d with %d points)"+                        prefix k lenChunk+               hPutStrLn stderr $+                 printf "%s: Training error:   %.2f%%"+                        prefix ((1 - trainScore) * 100)+               hPutStrLn stderr $+                 printf "%s: Validation error: %.2f%%"+                        prefix ((1 - testScore ) * 100)+             return res+       let runEpoch :: Int+                    -> ( Concrete (XParams width r)+                       , StateAdam (XParams width r) )+                    -> IO (Concrete (XParams width r))+           runEpoch n (params2, _) | n > epochs = return params2+           runEpoch n paramsStateAdam@(!_, !_) = do+             unless (sNatValue width < 10) $+               hPutStrLn stderr $ printf "\n%s: [Epoch %d]" prefix n+             let trainDataShuffled = shuffle (mkStdGen $ n + 5) trainData+                 chunks = take maxBatches+                          $ zip [1 ..]+                          $ chunksOf totalBatchSize trainDataShuffled+             res <- foldM runBatch paramsStateAdam chunks+             runEpoch (succ n) res+       res <- runEpoch 1 (targetInit, initialStateAdam ftk)+       let testErrorFinal = 1 - ftest @lenTestData testDataS res+       assertEqualUpToEpsilon 1e-1 expected testErrorFinal++{-# SPECIALIZE mnistTestCaseRNNSO+  :: String+  -> Int -> Int -> SNat width -> SNat batch_size -> Int -> Double+  -> TestTree #-}++tensorADValMnistTestsRNNSO :: TestTree+tensorADValMnistTestsRNNSO = testGroup "RNNS Once MNIST tests"+  [ mnistTestCaseRNNSO "RNNSO 1 epoch, 1 batch" 1 1 (SNat @128) (SNat @150) 5000+                       (0.6026 :: Double)+  , mnistTestCaseRNNSO "RNNSO artificial 1 2 3 4 5" 2 3 (SNat @4) (SNat @5) 50+                       (0.8933333 :: Float)+  , mnistTestCaseRNNSO "RNNSO artificial 5 4 3 2 1" 5 4 (SNat @3) (SNat @2) 49+                       (0.9336734693877551 :: Double)+  , mnistTestCaseRNNSO "RNNSO 1 epoch, 0 batch" 1 0 (SNat @128) (SNat @150) 50+                       (1.0 :: Float)+  ]
+ test/simplified/TestRevFwdFold.hs view
@@ -0,0 +1,2884 @@+{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}+{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}+-- | Test with folds and nested derivatives.+module TestRevFwdFold+  ( testTrees+  ) where++import Prelude++import Data.Proxy (Proxy (Proxy))+import GHC.TypeLits (KnownNat, type (+))+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape++import HordeAd+import HordeAd.Core.AstFreshId (resetVarCounter)+import HordeAd.Core.Ops (tbuild1, treplicate)+import HordeAd.Core.OpsConcrete ()++import CrossTesting+import EqEpsilon++testTrees :: [TestTree]+testTrees =+  [ testCase "4SRrev" testFooRgrad+  , testCase "4SRrev2" testFooRrev2+  , testCase "4SRrevPP1" testFooRrevPP1+  , testCase "4SRrevPP2" testFooRrevPP2+  , testCase "4SRrev3" testFooRrev3+  , testCase "4S0Rrev" testSin0Rgrad+  , testCase "4S0RrevPP1" testSin0RrevPP1+  , testCase "4S0RrevPP2" testSin0RrevPP2+  , testCase "4S0Rrev3" testSin0Rrev3+  , testCase "4S0Rrev4" testSin0Rrev4+  , testCase "4S0RrevPP4" testSin0RrevPP4+  , testCase "4S0Rrev5" testSin0Rrev5+  , testCase "4S0RrevPP5" testSin0RrevPP5+  , testCase "4S0Rrev3'" testSin0Rrev3'+  , testCase "4S0Rrev4'" testSin0Rrev4'+  , testCase "4S0Rrev5'" testSin0Rrev5'+  , testCase "4S0Rfwd" testSin0Rjvp+  , testCase "4S0RfwdPP0" testSin0RfwdPP0+  , testCase "4S0RfwdPP1" testSin0RfwdPP1+  , testCase "4S0RfwdPP1FullUnsimp" testSin0RfwdPP1FullUnsimp+  , testCase "4S0RfwdPP1Full" testSin0RfwdPP1Full+  , testCase "4S0Rfwd3" testSin0Rfwd3+  , testCase "4S0Rfwd4" testSin0Rfwd4+  , testCase "4S0RfwdPP4P" testSin0RfwdPP4P+  , testCase "4S0RfwdPP4Dual" testSin0RfwdPP4Dual+  , testCase "4S0Rfwd5" testSin0Rfwd5+  , testCase "4S0RfwdPP5" testSin0RfwdPP5+  , testCase "4S0Rfwd3'" testSin0Rfwd3'+  , testCase "4S0Rfwd4'" testSin0Rfwd4'+  , testCase "4S0Rfwd5'" testSin0Rfwd5'+  , testCase "4S0Rrev5S" testSin0Rrev5S+  , testCase "4S0RrevPP5S" testSin0RrevPP5S+  , testCase "4S0Fold0" testSin0Fold0+  , testCase "4S0Fold0ForComparison" testSin0Fold0ForComparison+  , testCase "4S0Fold1" testSin0Fold1+  , testCase "4S0FoldB1" testSin0FoldB1+  , testCase "4S0FoldB1PP" testSin0FoldB1PP+  , testCase "4S0FoldB2" testSin0FoldB2+  , testCase "4S0FoldB3" testSin0FoldB3+  , testCase "4S0FoldB4" testSin0FoldB4+  , testCase "4S0Fold2" testSin0Fold2+  , testCase "4S0FoldForComparison" testSin0FoldForComparison+  , testCase "4S0Fold3" testSin0Fold3+  , testCase "4S0Fold4" testSin0Fold4+  , testCase "4S0Fold5" testSin0Fold5+  , testCase "4S0Fold6" testSin0Fold6+  , testCase "4S0Fold7" testSin0Fold7+  , testCase "4S0Fold8" testSin0Fold8+  , testCase "4S0Fold0S" testSin0Fold0S+  , testCase "4S0Fold1S" testSin0Fold1S+  , testCase "4S0Fold2S" testSin0Fold2S+  , testCase "4S0FoldForComparisonS" testSin0FoldForComparisonS+  , testCase "4S0Fold3S" testSin0Fold3S+  , testCase "4S0Fold4S" testSin0Fold4S+  , testCase "4S0Fold5S" testSin0Fold5S+  , testCase "4S0Fold6S" testSin0Fold6S+  , testCase "4S0Fold7S" testSin0Fold7S+  , testCase "4S0Fold8S" testSin0Fold8S+  , testCase "4S0Fold8rev" testSin0Fold8grad+  , testCase "4S0Fold8rev2" testSin0Fold8rev2+  , testCase "4S0Fold8Srev" testSin0Fold8Sgrad+  , testCase "4S0Fold8Srev2" testSin0Fold8Srev2+  , testCase "4S0Fold182Srev" testSin0Fold182Sgrad+  , testCase "4S0Fold182SrevPP" testSin0Fold182SrevPP+  , testCase "4S0Fold18Srev" testSin0Fold18Sgrad+  , testCase "4S0Fold8fwd" testSin0Fold8jvp+  , testCase "4S0Fold8fwd2" testSin0Fold8fwd2+  , testCase "4S0Fold8Sfwd" testSin0Fold8Sjvp+  , testCase "4S0Fold8Sfwd2" testSin0Fold8Sfwd2+  , testCase "4S0Fold5Sfwd" testSin0Fold5Sjvp+  , testCase "4S0Fold5Sfwds" testSin0Fold5Sfwds+  , testCase "4S0Scan0" testSin0Scan0+  , testCase "4S0Scan1" testSin0Scan1+  , testCase "4S0Scan1ForComparison" testSin0Scan1ForComparison+  , testCase "4S0Scan2" testSin0Scan2+  , testCase "4S0Scan3" testSin0Scan3+  , testCase "4S0Scan4" testSin0Scan4+  , testCase "4S0Scan5" testSin0Scan5+  , testCase "4S0Scan6" testSin0Scan6+  , testCase "4S0Scan7" testSin0Scan7+  , testCase "4S0Scan8" testSin0Scan8+  , testCase "4S0Scan8rev" testSin0Scan8grad+  , testCase "4S0Scan8rev2" testSin0Scan8rev2+  , testCase "4S0Scan8Srev2" testSin0Scan8Srev2+  , testCase "4S0Scan1RevPP1" testSin0Scan1RevPP1+  , testCase "4S0Scan1RevPPForComparison" testSin0Scan1RevPPForComparison+  , testCase "4S0ScanFwdPP" testSin0ScanFwdPP+  , testCase "4S0ScanFwdPPFull" testSin0ScanFwdPPFull+  , testCase "4S0Scan1Rev2PP1" testSin0Scan1Rev2PP1+  , testCase "4S0Scan1Rev2PPA" testSin0Scan1Rev2PPA+  , testCase "4S0Scan1Rev2PPForComparison" testSin0Scan1Rev2PPForComparison+  , testCase "4S0Scan1Rev2" testSin0Scan1Rev2+  , testCase "4S0Scan1Rev2ForComparison" testSin0Scan1Rev2ForComparison+  , testCase "4S0Scan1Rev3PP0" testSin0Scan1Rev3PP0+  , testCase "4S0Scan1Rev3PPForComparison" testSin0Scan1Rev3PPForComparison+  , testCase "4S0ScanFwd3PP" testSin0ScanFwd3PP+  , testCase "4S0Scan1Rev3" testSin0Scan1Rev3+  , testCase "4S0Scan1Rev3ForComparison" testSin0Scan1Rev3ForComparison+  , testCase "4S0Scan0fwd" testSin0Scan0jvp+  , testCase "4S0Scan1fwd" testSin0Scan1jvp+  , testCase "4S0Scan1FwdForComparison" testSin0Scan1FwdForComparison+  , testCase "4S0Scan8fwd" testSin0Scan8jvp+  , testCase "4S0Scan8fwd2" testSin0Scan8fwd2+  , testCase "4SUnitriangular0PP" testUnitriangular0PP+  , testCase "4SUnitriangular1PP" testUnitriangular1PP+  , testCase "4SUnitriangular2PP" testUnitriangular2PP+  , testCase "4S0rmapAccumRD0S" testSin0rmapAccumRD0S+  , testCase "4S0rmapAccumRD00SC" testSin0rmapAccumRD00SC+  , testCase "4S0rmapAccumRD00S0" testSin0rmapAccumRD00S0+  , testCase "4S0rmapAccumRD00S" testSin0rmapAccumRD00S+  , testCase "4S0rmapAccumRD00S7" testSin0rmapAccumRD00S7+  , testCase "4S0rmapAccumRD00SCacc0" testSin0rmapAccumRD00SCacc0+  , testCase "4S0rmapAccumRD00SCacc" testSin0rmapAccumRD00SCacc+  , testCase "4S0rmapAccumRD00Sacc0" testSin0rmapAccumRD00Sacc0+  , testCase "4S0rmapAccumRD00Sacc" testSin0rmapAccumRD00Sacc+  , testCase "4S0rmapAccumRD00SCall0" testSin0rmapAccumRD00SCall0+  , testCase "4S0rmapAccumRD00SCall" testSin0rmapAccumRD00SCall+  , testCase "4S0rmapAccumRD00Sall0" testSin0rmapAccumRD00Sall0+  , testCase "4S0rmapAccumRD00Sall" testSin0rmapAccumRD00Sall+  , testCase "4S0rmapAccumRD0R" testSin0rmapAccumRD0R+  , testCase "4S0rmapAccumRD01SN" testSin0rmapAccumRD01SN+  , testCase "4S0rmapAccumRD01SN3" testSin0rmapAccumRD01SN3+  , testCase "4S0rmapAccumRD01SN5" testSin0rmapAccumRD01SN5+  , testCase "4S0rmapAccumRD01SN51" testSin0rmapAccumRD01SN51+  , testCase "4S0rmapAccumRD01SN531a" testSin0rmapAccumRD01SN531a+  , testCase "4S0rmapAccumRD01SN531b0" testSin0rmapAccumRD01SN531b0+--  , testCase "4S0rmapAccumRD01SN531b0PP" testSin0rmapAccumRD01SN531b0PP+  , testCase "4S0rmapAccumRD01SN531b0PPj" testSin0rmapAccumRD01SN531b0PPj+  , testCase "4S0rmapAccumRD01SN531bRPPj" testSin0rmapAccumRD01SN531bRPPj+  , testCase "4S0rmapAccumRD01SN531c" testSin0rmapAccumRD01SN531c+  , testCase "4S0rmapAccumRD01SN531Slice" testSin0rmapAccumRD01SN531Slice+  , testCase "4S0rmapAccumRD01SN55" testSin0rmapAccumRD01SN55+  , testCase "4S0rmapAccumRD01SN55acc" testSin0rmapAccumRD01SN55acc+  , testCase "4S0rmapAccumRD01SN58" testSin0rmapAccumRD01SN58+  , testCase "4S0rmapAccumRD01SN7" testSin0rmapAccumRD01SN7+  , testCase "4S0ScanD51" testSin0ScanD51+  , testCase "4S0ScanD8" testSin0ScanD8+  , testCase "4S0ScanD8MapAccum" testSin0ScanD8MapAccum+  , testCase "4S0ScanD8rev" testSin0ScanD8grad+  , testCase "4S0ScanD8rev4" testSin0ScanD8rev4+  , testCase "4S0ScanD1RevPP" testSin0ScanD1RevPP+  , testCase "4S0ScanDFwdPP" testSin0ScanDFwdPP+  , testCase "4S0ScanD1Rev2PP" testSin0ScanD1Rev2PP+  , testCase "4S0ScanDFwd2PP" testSin0ScanDFwd2PP+  , testCase "4S0ScanDFwd3PP" testSin0ScanDFwd3PP+  , testCase "4S0ScanD1fwd" testSin0ScanD1jvp+  , testCase "4S0ScanD8fwd" testSin0ScanD8jvp+  , testCase "4S0ScanD8fwdMapAccum" testSin0ScanD8fwdMapAccum+  , testCase "4S0ScanD8fwd2" testSin0ScanD8fwd2+  , testCase "4S0FoldNestedS1" testSin0FoldNestedS1+  , testCase "4S0FoldNestedS1PP" testSin0FoldNestedS1PP+  , testCase "4S0FoldNestedR1PP" testSin0FoldNestedR1PP+  , testCase "4S0FoldNestedR0LengthPPs" testSin0FoldNestedR0LengthPPs+  , testCase "4S0FoldNestedR1LengthPPs" testSin0FoldNestedR1LengthPPs+  , testCase "4S0FoldNestedR2LengthPPs" testSin0FoldNestedR2LengthPPs+  , testCase "4S0FoldNestedR3LengthPPs" testSin0FoldNestedR3LengthPPs+-- takes too long:    , testCase "4S0FoldNestedR4LengthPPs" testSin0FoldNestedR4LengthPPs+-- takes too long:    , testCase "4S0FoldNestedR5LengthPPs" testSin0FoldNestedR5LengthPPs+  , testCase "4S0FoldNestedR2LengthPPsDummy7" testSin0FoldNestedR2LengthPPsDummy7+  , testCase "4S0FoldNestedR2Dummy7" testSin0FoldNestedR2Dummy7+  , testCase "4S0FoldNestedR2Tan" testSin0FoldNestedR2Tan+  , testCase "4S0FoldNestedS1FwdFwd0" testSin0FoldNestedS1FwdFwd0+  , testCase "4S0FoldNestedS1FwdFwd" testSin0FoldNestedS1FwdFwd+  , testCase "4S0FoldNestedS1RevRev" testSin0FoldNestedS1RevRev+  , testCase "4S0FoldNestedS2" testSin0FoldNestedS2+  , testCase "4S0FoldNestedS3" testSin0FoldNestedS3+  , testCase "4S0FoldNestedS4" testSin0FoldNestedS4+  , testCase "4S0FoldNestedS5" testSin0FoldNestedS5+  , testCase "4S0FoldNestedS5rev" testSin0FoldNestedS5grad+  , testCase "4S0FoldNestedS5fwd" testSin0FoldNestedS5jvp+  , testCase "4S0FoldNestedSi" testSin0FoldNestedSi+  , testCase "4S0FoldNestedR1" testSin0FoldNestedR1+  , testCase "4S0FoldNestedR1RevFwd" testSin0FoldNestedR1RevFwd+  , testCase "4S0FoldNestedR2" testSin0FoldNestedR2+  , testCase "4S0FoldNestedR2RevFwd" testSin0FoldNestedR2RevFwd+  , testCase "4S0FoldNestedR3" testSin0FoldNestedR3+  , testCase "4S0FoldNestedR4" testSin0FoldNestedR4+  , testCase "4S0FoldNestedR41" testSin0FoldNestedR41+  , testCase "4S0FoldNestedR40" testSin0FoldNestedR40+  , testCase "4S0FoldNestedR400" testSin0FoldNestedR400+  , testCase "4S0FoldNestedRi" testSin0FoldNestedRi+  , testCase "4S0FoldNestedR22" testSin0FoldNestedR22+  , testCase "4S0FoldNestedR21" testSin0FoldNestedR21+  , testCase "4S0FoldNestedR21PP" testSin0FoldNestedR21PP+  , testCase "4S0revhV" testSin0revhV+  , testCase "4S0revhVPP" testSin0revhVPP+  , testCase "4S0revhV4" testSin0revhV4+  , testCase "4S0revhV5" testSin0revhV5+  , testCase "4S0revhV6" testSin0revhV6+  , testCase "4S0revhV7" testSin0revhV7+  ]++foo :: RealFloatH a => (a, a, a) -> a+foo (x, y, z) =+  let w = x * sin y+  in atan2H z w + z * w++fooRgrad :: forall g a.+           (ADReady g, GoodScalar a, Differentiable a, ADTensorScalar a ~ a)+        => (a, a, a) -> (g (TKR 0 a), g (TKR 0 a), g (TKR 0 a))+fooRgrad (x, y, z) =+  let f :: forall f. ADReady f => f (TKProduct (TKProduct (TKR 0 a) (TKR 0 a)) (TKR 0 a)) -> f (TKR 0 a)+      f v = foo (tproject1 (tproject1 v), tproject2 (tproject1 v), tproject2 v)+      shapes = FTKProduct (FTKProduct (FTKR ZSR FTKScalar) (FTKR ZSR FTKScalar)) (FTKR ZSR FTKScalar)+      domsOf = kgrad (kfromR . f) shapes+                    (tpair (tpair (rconcrete $ Nested.rscalar x)+                                  (rconcrete $ Nested.rscalar y))+                           (rconcrete $ Nested.rscalar z))+  in ( tlet domsOf (\v -> tproject1 (tproject1 v))+     , tlet domsOf (\v -> tproject2 (tproject1 v))+     , tlet domsOf (\v -> tproject2 v) )++testFooRgrad :: Assertion+testFooRgrad = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 2.4396285219055063, rscalar (-1.953374825727421), rscalar 0.9654825811012627)+    (fooRgrad @Concrete @Double (1.1, 2.2, 3.3))++testFooRrev2 :: Assertion+testFooRrev2 = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 2.4396284, rscalar (-1.9533751), rscalar 0.96548253)+    (fooRgrad @Concrete @Float (1.1, 2.2, 3.3))++testFooRrevPP1 :: Assertion+testFooRrevPP1 = do+  resetVarCounter+  let (a1, _, _) = fooRgrad @(AstTensor AstMethodLet PrimalSpan) @Double (1.1, 2.2, 3.3)+  printAstPretty a1+    @?= "rfromS (let x15 = let x10 = sin (sscalar 2.2) ; x11 = sscalar 1.1 * x10 ; x12 = recip (sscalar 10.889999999999999 + x11 * x11) ; x13 = sin (sscalar 2.2) ; x14 = sscalar (-3.3) * x12 in tpair (tpair (x10 * x14 + sscalar 3.3 * x13) (sscalar 1.1 * (cos (sscalar 2.2) * x14) + sscalar 3.63 * cos (sscalar 2.2))) (x11 * x12 + sscalar 1.1 * x13) in tproject1 (tproject1 x15))"++testFooRrevPP2 :: Assertion+testFooRrevPP2 = do+  let (a1, _, _) = fooRgrad @(AstTensor AstMethodLet PrimalSpan) @Double (1.1, 2.2, 3.3)+  printAstSimple (simplifyInlineContract a1)+    @?= "rfromS (sscalar 2.668038132604647 + sscalar (-2.668038132604647) * recip (sscalar 11.680936386336942))"++testFooRrev3 :: Assertion+testFooRrev3 = do+  let f :: ADVal Concrete (TKR 0 Double) -> ADVal Concrete (TKR 0 Double)+      f (D a _) =+        let (a1, _, _) = fooRgrad @(ADVal Concrete) @Double+                                 (Nested.runScalar (unConcrete a), 2.2, 3.3)+        in a1+  assertEqualUpToEpsilon 1e-10+    (rscalar 0)+    (cgrad (kfromR . f) (rscalar 1.1))++testSin0Rgrad :: Assertion+testSin0Rgrad = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 0.4535961214255773)+    (rrev1 @Concrete @Double @0 @0 sin (rscalar 1.1))++testSin0RrevPP1 :: Assertion+testSin0RrevPP1 = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin (rscalar 1.1)+  printAstPretty a1+    @?= "rfromS (cos (sscalar 1.1))"++testSin0RrevPP2 :: Assertion+testSin0RrevPP2 = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin (rscalar 1.1)+  printAstSimple a1+    @?= "rfromS (cos (sscalar 1.1))"++testSin0Rrev3 :: Assertion+testSin0Rrev3 = do+  let f = rrev1 @(ADVal Concrete) @Double @0 @0 sin+  assertEqualUpToEpsilon 1e-10+    (rscalar (-0.8912073600614354))+    (cgrad (kfromR . f) (rscalar 1.1))++testSin0Rrev4 :: Assertion+testSin0Rrev4 = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 0.8988770945225438)+    ((rrev1 sin . rrev1 @Concrete @Double @0 @0 sin) (rscalar 1.1))++testSin0RrevPP4 :: Assertion+testSin0RrevPP4 = do+  let a1 = (rrev1 sin . rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin) (rscalar 1.1)+  printAstPretty (simplifyInline a1)+    @?= "rfromS (cos (cos (sscalar 1.1)))"++testSin0Rrev5 :: Assertion+testSin0Rrev5 = do+  assertEqualUpToEpsilon 1e-10+    (rscalar (-0.8912073600614354))+    (rrev1 @Concrete @Double @0 @0 (rrev1 sin) (rscalar 1.1))++testSin0RrevPP5 :: Assertion+testSin0RrevPP5 = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 (rrev1 sin) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sscalar (-0.8912073600614354))"++testSin0Rrev3' :: Assertion+testSin0Rrev3' = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar (-0.8912073600614354) :: Concrete (TKR 0 Double))+    (rev' (rrev1 sin) (rscalar 1.1))++testSin0Rrev4' :: Assertion+testSin0Rrev4' = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.39052780643689855 :: Concrete (TKR 0 Double))+    (rev' (rrev1 sin . rrev1 sin) (rscalar 1.1))++testSin0Rrev5' :: Assertion+testSin0Rrev5' = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar (-0.4535961214255773) :: Concrete (TKR 0 Double))+    (rev' (rrev1 (rrev1 sin)) (rscalar 1.1))++testSin0Rjvp :: Assertion+testSin0Rjvp = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 0.4535961214255773)+    (rfwd1 @Concrete @Double @0 @0 sin (rscalar 1.1))++testSin0RfwdPP0 :: Assertion+testSin0RfwdPP0 = do+  resetVarCounter+  let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double)+      a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin (rscalar 1.1)+  printAstPretty a1+    @?= "rfromS (cos (sscalar 1.1))"++testSin0RfwdPP1 :: Assertion+testSin0RfwdPP1 = do+  resetVarCounter+  let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double)+      a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin (rscalar 1.1)+  printAstPretty (simplifyInline a1)+    @?= "rfromS (cos (sscalar 1.1))"++testSin0RfwdPP1FullUnsimp :: Assertion+testSin0RfwdPP1FullUnsimp = do+  resetVarCounter+  let a1 :: AstTensor AstMethodLet FullSpan (TKR 0 Double)+      a1 = rfwd1 @(AstTensor AstMethodLet FullSpan) @Double @0 @0 sin (rscalar 1.1)+  printAstPretty a1+    @?= "rfromS (sscalar 1.0 * cos (sscalar 1.1))"++testSin0RfwdPP1Full :: Assertion+testSin0RfwdPP1Full = do+  resetVarCounter+  let a1 :: AstTensor AstMethodLet FullSpan (TKR 0 Double)+      a1 = rfwd1 @(AstTensor AstMethodLet FullSpan) @Double @0 @0 sin (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sscalar 0.4535961214255773)"++testSin0Rfwd3 :: Assertion+testSin0Rfwd3 = do+  let f :: ADVal Concrete (TKR 0 Double) -> ADVal Concrete (TKR 0 Double)+      f = rfwd1 @(ADVal Concrete) @Double @0 @0 sin+  assertEqualUpToEpsilon 1e-10+    (rscalar (-0.9803280960675791))+    (cjvp f (rscalar 1.1) (rscalar 1.1))++testSin0Rfwd4 :: Assertion+testSin0Rfwd4 = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 0.8988770945225438)+    ((rfwd1 sin . rfwd1 @Concrete @Double @0 @0 sin) (rscalar 1.1))++testSin0RfwdPP4P :: Assertion+testSin0RfwdPP4P = do+  let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double)+      a1 = (rfwd1 sin . rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 sin) (rscalar 1.1)+  printAstPretty (simplifyInline a1)+    @?= "rfromS (cos (cos (sscalar 1.1)))"++testSin0RfwdPP4Dual :: Assertion+testSin0RfwdPP4Dual = do+  let a1 :: AstTensor AstMethodLet DualSpan (TKR 0 Double)+      a1 = (rfwd1 sin . rfwd1 @(AstTensor AstMethodLet DualSpan) @Double @0 @0 sin) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sdualPart (sscalar 0.0) * cos (sdualPart (sscalar 0.0) * cos (sdualPart (sscalar 0.0))))"++testSin0Rfwd5 :: Assertion+testSin0Rfwd5 = do+  assertEqualUpToEpsilon 1e-10+    (rscalar (-0.8912073600614354))+    (rfwd1 @Concrete @Double @0 @0 (rfwd1 sin) (rscalar 1.1))++testSin0RfwdPP5 :: Assertion+testSin0RfwdPP5 = do+  let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 0 Double)+      a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0 (rfwd1 sin) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sscalar (-0.8912073600614354))"++testSin0Rfwd3' :: Assertion+testSin0Rfwd3' = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar (-0.8912073600614354) :: Concrete (TKR 0 Double))+    (rev' (rfwd1 sin) (rscalar 1.1))++testSin0Rfwd4' :: Assertion+testSin0Rfwd4' = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.39052780643689855 :: Concrete (TKR 0 Double))+    (rev' (rfwd1 sin . rfwd1 sin) (rscalar 1.1))++testSin0Rfwd5' :: Assertion+testSin0Rfwd5' = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar (-0.4535961214255773) :: Concrete (TKR 0 Double))+    (rev' (rfwd1 (rfwd1 sin)) (rscalar 1.1))++testSin0Rrev5S :: Assertion+testSin0Rrev5S = do+  assertEqualUpToEpsilon 1e-10+    (srepl (-0.8912073600614354))+    (srev1 @Concrete @Double @'[] @'[] (srev1 sin) (srepl 1.1))++testSin0RrevPP5S :: Assertion+testSin0RrevPP5S = do+  resetVarCounter+  let a1 = srev1 @(AstTensor AstMethodLet PrimalSpan) @Double @'[] @'[] (srev1 sin) (srepl 1.1)+  printAstPretty (simplifyInline a1)+    @?= "negate (sin (sscalar 1.1))"++testSin0Fold0 :: Assertion+testSin0Fold0 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1.0 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f x0 = rfold (\x _a -> sin x)+                            x0 (rrepl @_ @Double (0 :$: ZSR) 0)+           in f) (rscalar 1.1))++testSin0Fold0ForComparison :: Assertion+testSin0Fold0ForComparison = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1.0 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. f (TKR 0 Double) -> f (TKR 0 Double)+               f = id+           in f) (rscalar 1.1))++testSin0Fold1 :: Assertion+testSin0Fold1 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.4535961214255773 :: Concrete (TKR 0 Double))+    (rev' (\x0 -> rfold (\x _a -> sin x)+                        x0 (rrepl @1 @Double [1] 42)) (rscalar 1.1))++testSin0FoldB1 :: Assertion+testSin0FoldB1 = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 0 :: Concrete (TKR 0 Double))+    (rrev1 (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+                f x0 = rfold (\_x _a -> rscalar 7)+                         (rscalar 5) (rreplicate 1 x0)+            in f) (rscalar 1.1))++testSin0FoldB1PP :: Assertion+testSin0FoldB1PP = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan)+             (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+                  f x0 = rfold (\_x _a -> rscalar 7)+                           (rscalar 5) (rreplicate 1 x0)+              in f) (rscalar 1.1)+  printAstPretty a1+    @?= "rsum (tproject2 (tmapAccumRDer (SNat @1) <lambda> <lambda> <lambda> (tconvert (ConvCmp (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp ConvXS (Conv0X STKScalar))) (STKScalar) 1.0) (tpair (sconcrete (sfromListLinear [1] [Z1])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> (sscalar 5.0) (sconcrete (sfromListLinear [1] [1.1]))))) (sconcrete (sfromListLinear [1] [1.1]))))))"++testSin0FoldB2 :: Assertion+testSin0FoldB2 = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 0 :: Concrete (TKR 0 Double))+    (grad (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f x0 = rfold (\_x _a -> rscalar 7)+                        (rscalar 5) (rreplicate 1 x0)+           in kfromR . f) (rscalar 1.1))++testSin0FoldB3 :: Assertion+testSin0FoldB3 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f x0 = rfold (\_x _a -> rscalar 7)+                        (rscalar 5) (rreplicate 1 x0)+           in f) (rscalar 1.1))++testSin0FoldB4 :: Assertion+testSin0FoldB4 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f x0 = rfold (\_x _a -> rscalar 7)+                        x0 (rrepl @1 @Double [1] 42)+           in f) (rscalar 1.1))++testSin0Fold2 :: Assertion+testSin0Fold2 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.12389721944941383 :: Concrete (TKR 0 Double))+    (rev' (\x0 -> rfold (\x _a -> sin x)+                        x0 (rrepl @1 @Double [5] 42)) (rscalar 1.1))++testSin0FoldForComparison :: Assertion+testSin0FoldForComparison = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.12389721944941383 :: Concrete (TKR 0 Double))+    (rev' (sin . sin . sin . sin . sin) (rscalar 1.1))++testSin0Fold3 :: Assertion+testSin0Fold3 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.4535961214255773 :: Concrete (TKR 0 Double))+    (rev' (\a0 -> rfold (\_x a -> sin a)+                        (rscalar 84) (rreplicate 3 a0)) (rscalar 1.1))++testSin0Fold4 :: Assertion+testSin0Fold4 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar  (-0.7053476446727861) :: Concrete (TKR 0 Double))+    (rev' (\a0 -> rfold (\x a -> atan2H (sin x) (sin a))+                        (rscalar 2 * a0) (rreplicate 3 a0)) (rscalar 1.1))++testSin0Fold5 :: Assertion+testSin0Fold5 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1.2992412552109085 :: Concrete (TKR 0 Double))+    (rev' (\a0 -> rfold (\x a -> rsum+                                 $ atan2H (sin $ rreplicate 5 x)+                                          (rsum $ sin $ rsum+                                           $ rtr $ rreplicate 7 a))+                        (rscalar 2 * a0)+                        (rreplicate 3 (rreplicate 2 (rreplicate 5 a0)))) (rscalar 1.1))++testSin0Fold6 :: Assertion+testSin0Fold6 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 6 :: Concrete (TKR 0 Double))+    (rev' (\a0 -> rfold (\x a -> rtr+                                 $ rtr x + rreplicate 1 (rreplicate 2 a))+                        (rreplicate 2 (rreplicate 1 a0))+                        (rreplicate 2 a0)) (rscalar 1.1))++testSin0Fold7 :: Assertion+testSin0Fold7 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 250 :: Concrete (TKR 0 Double))+    (rev' (\a0 -> rfold (\x _a -> rtr $ rreplicate 5 $ rsum (rtr x))+                        (rreplicate 2 (rreplicate 5 a0))+                        (rreplicate 2 a0)) (rscalar 1.1))++testSin0Fold8 :: Assertion+testSin0Fold8 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar (-2.200311410593445) :: Concrete (TKR 0 Double))+    (rev' (\a0 -> rfold (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                        (rreplicate 3 a0)) (rscalar 1.1))++testSin0Fold0S :: Assertion+testSin0Fold0S = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1.0 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f x0 = sfold @0 @'[] @'[] @(TKScalar Double) @(TKScalar Double)+                            (\x _a -> sin x)+                            x0 (srepl 0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold1S :: Assertion+testSin0Fold1S = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.4535961214255773 :: Concrete (TKR 0 Double))+    (rev' ((let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+                f x0 = sfold (let g :: forall f2. ADReady f2+                                   => f2 (TKS '[] Double) -> f2 (TKS '[] Double)+                                   -> f2 (TKS '[] Double)+                                  g x _a = sin x+                              in g)+                        x0 (srepl @'[1] 42)+            in rfromS . f . sfromR)) (rscalar 1.1))++testSin0Fold2S :: Assertion+testSin0Fold2S = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.12389721944941383 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f x0 = sfold (\x _a -> sin x)+                        x0 (srepl @'[5] @Double 42)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0FoldForComparisonS :: Assertion+testSin0FoldForComparisonS = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.12389721944941383 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f = sin . sin . sin . sin . sin+          in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold3S :: Assertion+testSin0Fold3S = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.4535961214255773 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\_x a -> sin a)+                        (srepl 84) (sreplicate @3 @_ @_ @f a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold4S :: Assertion+testSin0Fold4S = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar  (-0.7053476446727861) :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\x a -> atan2H (sin x) (sin a))+                        (srepl 2 * a0) (sreplicate @3 @_ @_ @f a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold5S :: Assertion+testSin0Fold5S = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1.2992412552109085 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (let g :: forall f2. ADReady f2+                                   => f2 (TKS '[] Double) -> f2 (TKS '[2, 5] Double)+                                   -> f2 (TKS '[] Double)+                                 g x a = ssum+                                   $ atan2H (sin $ sreplicate @5 @_ @_ @f2 x)+                                            (ssum $ sin $ ssum+                                             $ str $ sreplicate @7 @_ @_ @f2 a)+                             in g)+                        (srepl 2 * a0)+                        (sreplicate @3 @_ @_ @f+                                    (sreplicate @2 @_ @_ @f+                                                (sreplicate @5 @_ @_ @f a0)))+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold6S :: Assertion+testSin0Fold6S = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 6 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[2, 1] Double)+               f a0 = sfold @2 @'[2, 1] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\x a -> str+                                 $ str x + sreplicate @1+                                                      (sreplicate @2 a))+                        (sreplicate @2 (sreplicate @1 a0))+                        (sreplicate @2 a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold7S :: Assertion+testSin0Fold7S = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 250 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[2, 5] Double)+               f a0 = sfold @2 @'[2, 5] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\x _a -> str $ sreplicate @5 $ ssum (str x))+                        (sreplicate @2 (sreplicate @5 a0))+                        (sreplicate @2 a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold8S :: Assertion+testSin0Fold8S = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar (-2.200311410593445) :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[2, 5] Double)+               f a0 = sfold @3 @'[2, 5] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\x a -> str $ sreplicate @5+                                 $ atan2H (ssum (str $ sin x))+                                          (sreplicate @2+                                           $ sin (ssum $ sreplicate @7 a)))+                        (sreplicate @2 (sreplicate @5 (srepl 2 * a0)))+                        (sreplicate @3 a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold8grad :: Assertion+testSin0Fold8grad = do+  assertEqualUpToEpsilon 1e-10+    (rscalar (-2.200311410593445) :: Concrete (TKR 0 Double))+    (rrev1 @Concrete @Double @0 @2+       (\a0 -> rfold (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                        (rreplicate 3 a0)) (rscalar 1.1))++testSin0Fold8rev2 :: Assertion+testSin0Fold8rev2 = do+  let h = rrev1 @(ADVal Concrete) @Double @0 @2+        (\a0 -> rfold (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                        (rreplicate 3 a0))+  assertEqualUpToEpsilon 1e-10+    (rscalar 98.72666469795736)+    (cgrad (kfromR . h) (rscalar 1.1))++testSin0Fold8Sgrad :: Assertion+testSin0Fold8Sgrad = do+  assertEqualUpToEpsilon 1e-10+    (rscalar (-2.200311410593445) :: Concrete (TKR 0 Double))+    (rrev1 (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[2, 5] Double)+                f a0 = sfold @3 @'[2, 5] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\x a -> str $ sreplicate @5+                                 $ atan2H (ssum (str $ sin x))+                                          (sreplicate @2+                                           $ sin (ssum $ sreplicate @7 a)))+                        (sreplicate @2 (sreplicate @5 (srepl 2 * a0)))+                        (sreplicate @3 a0)+            in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold8Srev2 :: Assertion+testSin0Fold8Srev2 = do+  let h = srev1 @(ADVal Concrete)+                (let f :: forall f. ADReady f+                       => f (TKS '[] Double) -> f (TKS '[2, 5] Double)+                     f a0 = sfold @3 @'[2, 5] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\x a -> str $ sreplicate @5+                                 $ atan2H (ssum (str $ sin x))+                                          (sreplicate @2+                                           $ sin (ssum $ sreplicate @7 a)))+                        (sreplicate @2 (sreplicate @5 (sscalar 2 * a0)))+                        (sreplicate @3 a0)+                 in f)+  assertEqualUpToEpsilon 1e-10+    (Concrete $ Nested.sscalar 6.182232283434464e-2)  -- seems quite unstable+    (cgrad (kfromS . h) (srepl 0.0001))++testSin0Fold182Sgrad :: Assertion+testSin0Fold182Sgrad = do+  assertEqualUpToEpsilon 1e-10+    (rscalar (-0.4409160296923509) :: Concrete (TKR 0 Double))+    (rrev1 (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[5] Double)+                f a0 = sfold @1 @'[5] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\_x a -> atan2H (sreplicate @5 a)+                                         (sreplicate @5+                                          $ sin (ssum $ sreplicate @7 a)))+                        (sreplicate @5 a0)+                        (sreplicate @1 a0)+            in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold182SrevPP :: Assertion+testSin0Fold182SrevPP = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan)+           (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[5] Double)+                f a0 = sfold @1 @'[5] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\_x a -> atan2H (sreplicate @5 a)+                                         (sreplicate @5+                                          $ sin (ssum $ sreplicate @7 a)))+                        (sreplicate @5 a0)+                        (sreplicate @1 a0)+            in rfromS . f . sfromR) (rscalar 1.1)+  printAstPretty a1+    @?= "rfromS (let v6 = tmapAccumRDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sreplicate [5] 1.0)) (tpair (sconcrete (sfromListLinear [1] [Z1])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sreplicate [5] 1.1)) (sconcrete (sfromListLinear [1] [1.1]))))) (sconcrete (sfromListLinear [1] [1.1])))) in ssum @5 (tproject1 v6) + tproject2 v6 !$ [0])"++testSin0Fold18Sgrad :: Assertion+testSin0Fold18Sgrad = do+  assertEqualUpToEpsilon 1e-10+    (rscalar (-2.4026418024701366) :: Concrete (TKR 0 Double))+    (rrev1 (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[2, 5] Double)+                f a0 = sfold @2 @'[2, 5] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\x a -> str $ sreplicate @5+                                 $ atan2H (ssum (str $ sin x))+                                          (sreplicate @2+                                           $ sin (ssum $ sreplicate @7 a)))+                        (sreplicate @2 (sreplicate @5 (srepl 2 * a0)))+                        (sreplicate @2 a0)+            in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold8jvp :: Assertion+testSin0Fold8jvp = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2, 5] (replicate 10 (-0.2200311410593445)))+    (rfwd1 @Concrete @Double @0 @2+       (\a0 -> rfold (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                        (rreplicate 3 a0)) (rscalar 1.1))++testSin0Fold8fwd2 :: Assertion+testSin0Fold8fwd2 = do+  let h :: ADVal Concrete (TKR 0 Double) -> ADVal Concrete (TKR 2 Double)+      h = rfwd1 @(ADVal Concrete) @Double @0 @2+        (\a0 -> rfold (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                        (rreplicate 3 a0))+  assertEqualUpToEpsilon 1e-10+    (rscalar 98.72666469795735)+    (cgrad (kfromR . rsum0 . h) (rscalar 1.1))++testSin0Fold8Sjvp :: Assertion+testSin0Fold8Sjvp = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2, 5] (replicate 10 (-0.2200311410593445)))+    (rfwd1 @Concrete+           (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[2, 5] Double)+                f a0 = sfold @3  @'[2, 5] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\x a -> str $ sreplicate @5+                                 $ atan2H (ssum (str $ sin x))+                                          (sreplicate @2+                                           $ sin (ssum $ sreplicate @7 a)))+                        (sreplicate @2 (sreplicate @5 (srepl 2 * a0)))+                        (sreplicate @3 a0)+            in rfromS . f . sfromR) (rscalar 1.1))++testSin0Fold8Sfwd2 :: Assertion+testSin0Fold8Sfwd2 = do+  let h :: ADVal Concrete (TKR 0 Double) -> ADVal Concrete (TKR 2 Double)+      h = rfwd1 @(ADVal Concrete)+                (let f :: forall f. ADReady f+                       => f (TKS '[] Double) -> f (TKS '[2, 5] Double)+                     f a0 = sfold @3 @'[2, 5] @'[] @(TKScalar Double) @(TKScalar Double)+                        (\x a -> str $ sreplicate @5+                                 $ atan2H (ssum (str $ sin x))+                                          (sreplicate @2+                                           $ sin (ssum $ sreplicate @7 a)))+                        (sreplicate @2 (sreplicate @5 (srepl 2 * a0)))+                        (sreplicate @3 a0)+                 in rfromS . f . sfromR)+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2, 5] (replicate 10 10.859933116775313))+    (cjvp h (rscalar 1.1) (rscalar 1.1))++testSin0Fold5Sjvp :: Assertion+testSin0Fold5Sjvp = do+  assertEqualUpToEpsilon 1e-10+    (rscalar 1.4291653807319993)+    (cjvp (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (let g :: forall f2. ADReady f2+                                   => f2 (TKS '[] Double) -> f2 (TKS '[2, 5] Double)+                                   -> f2 (TKS '[] Double)+                                 g x a = ssum+                                   $ atan2H (sin $ sreplicate @5 @_ @_ @f2 x)+                                            (ssum $ sin $ ssum+                                             $ str $ sreplicate @7 @_ @_ @f2 a)+                             in g)+                        (srepl 2 * a0)+                        (sreplicate @3 @_ @_ @f+                                    (sreplicate @2 @_ @_ @f+                                                (sreplicate @5 @_ @_ @f a0)))+           in rfromS . f . sfromR) (rscalar 1.1) (rscalar 1.1))++testSin0Fold5Sfwds :: Assertion+testSin0Fold5Sfwds = do+  assertEqualUpToEpsilon 1e-10+    (srepl 1.4291653807319993)+    (cjvp (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (let g :: forall f2. ADReady f2+                                   => f2 (TKS '[] Double) -> f2 (TKS '[2, 5] Double)+                                   -> f2 (TKS '[] Double)+                                 g x a = ssum+                                   $ atan2H (sin $ sreplicate @5 @_ @_ @f2 x)+                                            (ssum $ sin $ ssum+                                             $ str $ sreplicate @7 @_ @_ @f2 a)+                             in g)+                        (srepl 2 * a0)+                        (sreplicate @3 @_ @_ @f+                                    (sreplicate @2 @_ @_ @f+                                                (sreplicate @5 @_ @_ @f a0)))+           in f) (srepl 1.1) (srepl 1.1))++testSin0Scan0 :: Assertion+testSin0Scan0 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1)+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 1 Double)+               f x0 = rscan (\x _a -> sin x)+                            x0 (rrepl @_ @Double (0 :$: ZSR) 0)+           in f) (rscalar 1.1))++testSin0Scan1 :: Assertion+testSin0Scan1 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1,1,1] [1.4535961214255773] :: Concrete (TKR 5 Double))+    (rev' (\x0 -> rscan (\x _a -> sin x)+                        x0 (rrepl @1 @Double [1] 42))+          (ringestData [1,1,1,1,1] [1.1]))++testSin0Scan1ForComparison :: Assertion+testSin0Scan1ForComparison = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1,1,1] [1.4535961214255773] :: Concrete (TKR 5 Double))+    (rev' (\x0 -> rfromList [x0, sin x0])+          (ringestData [1,1,1,1,1] [1.1]))++testSin0Scan2 :: Assertion+testSin0Scan2 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1,1,1] [2.2207726343670955] :: Concrete (TKR 5 Double))+    (rev' (\x0 -> rscan (\x _a -> sin x)+                        x0 (rrepl @1 @Double [5] 42))+          (ringestData [1,1,1,1,1] [1.1]))++testSin0Scan3 :: Assertion+testSin0Scan3 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1,1,1] [1.360788364276732] :: Concrete (TKR 5 Double))+    (rev' (\a0 -> rscan (\_x a -> sin a)+                        (rreplicate0N [1,1,1,1,1] (rscalar 84))+                        (rreplicate 3 a0)) (ringestData [1,1,1,1,1] [1.1]))++testSin0Scan4 :: Assertion+testSin0Scan4 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1,1,1] [-0.4458209450295252] :: Concrete (TKR 5 Double))+    (rev' (\a0 -> rscan (\x a -> atan2H (sin x) (sin a))+                        (rreplicate0N [1,1,1,1,1] (rscalar 2) * a0)+                        (rreplicate 3 a0)) (ringestData [1,1,1,1,1] [1.1]))++testSin0Scan5 :: Assertion+testSin0Scan5 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1,1] [4.126141830000979] :: Concrete (TKR 4 Double))+    (rev' (\a0 -> rscan (\x a -> rsum+                                 $ atan2H (sin $ rreplicate 5 x)+                                         (rsum $ sin $ rsum+                                          $ rtr $ rreplicate 7 a))+                        (rreplicate0N [1,1,1,1] (rscalar 2) * a0)+                        (rreplicate 3 (rreplicate 2 (rreplicate 5 a0))))+          (ringestData [1,1,1,1] [1.1]))++testSin0Scan6 :: Assertion+testSin0Scan6 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1] [12] :: Concrete (TKR 2 Double))+    (rev' (\a0 -> rscan (\x a -> rtr+                                 $ rtr x + rreplicate 1 (rreplicate 2 a))+                        (rreplicate 2 (rreplicate 1 a0))+                        (rreplicate 2 a0)) (ringestData [1,1] [1.1]))++testSin0Scan7 :: Assertion+testSin0Scan7 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1] [310] :: Concrete (TKR 2 Double))+    (rev' (\a0 -> rscan (\x _a -> rtr $ rreplicate 5 $ rsum (rtr x))+                        (rreplicate 2 (rreplicate 5 a0))+                        (rreplicate 2 a0)) (ringestData [1,1] [1.1]))++testSin0Scan8 :: Assertion+testSin0Scan8 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1] [9.532987357352765] :: Concrete (TKR 3 Double))+    (rev' (\a0 -> rscan (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rreplicate0N [1,1,1] (rscalar 2) * a0)))+                        (rreplicate 3 a0)) (ringestData [1,1,1] [1.1]))++testSin0Scan8grad :: Assertion+testSin0Scan8grad = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [] [9.53298735735276])+    (rrev1 @Concrete @Double @0 @3+       (\a0 -> rscan (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                        (rreplicate 3 a0)) (rscalar 1.1))++testSin0Scan8rev2 :: Assertion+testSin0Scan8rev2 = do+  let h = rrev1 @(ADVal Concrete) @Double @0 @3+        (\a0 -> rscan (\x a -> rtr $ rreplicate 5+                               $ atan2H (rsum (rtr $ sin x))+                                        (rreplicate 2+                                         $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                        (rreplicate 3 a0))+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [] [285.9579482947575])+    (cgrad (kfromR . h) (rscalar 1.1))++testSin0Scan8Srev2 :: Assertion+testSin0Scan8Srev2 = do+  let h = srev1 @(ADVal Concrete) @Double @'[]+        (\a0 -> sscan (\x a -> str $ sreplicate @5+                               $ atan2H (ssum (str $ sin x))+                                        (sreplicate @2+                                         $ sin (ssum $ sreplicate @7 a)))+                        (sreplicate @2 (sreplicate @5 (sscalar 2 * a0)))+                        (sreplicate @3 a0))+  assertEqualUpToEpsilon 1e-10+    (sconcrete $ Nested.sfromListPrimLinear [] [285.9579482947575])+    (cgrad (kfromS . h) (sscalar 1.1))++testSin0Scan1RevPP1 :: Assertion+testSin0Scan1RevPP1 = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscan (\x _a -> sin x) x0+                           (rrepl @1 @Double [2] 42)) (rscalar 1.1)+  printAstPrettyButNested (simplifyInlineContract a1)+    @?= "rfromS (sscalar 1.0 + tproject1 (tmapAccumRDer (SNat @2) (\\x6 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair (cos (sfromR (tproject1 (tproject2 (tproject2 x6)))) * (tproject1 (tproject2 x6) + tproject1 x6)) (sscalar 0.0))) (\\x13 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair ((sfromR (tproject1 (tproject2 (tproject2 (tproject1 x13)))) * negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x13))))))) * (tproject1 (tproject2 (tproject2 x13)) + tproject1 (tproject2 x13)) + (tproject1 (tproject2 (tproject1 x13)) + tproject1 (tproject1 x13)) * cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x13)))))) (sscalar 0.0))) (\\x22 -> tconvert (ConvCmp (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (ConvT2 ConvId (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX)))) (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)))))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)))) (let x30 = cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x22))))) * tproject1 (tproject1 x22) in tpair x30 (tpair x30 (tpair (negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x22)))))) * ((tproject1 (tproject2 (tproject2 x22)) + tproject1 (tproject2 x22)) * tproject1 (tproject1 x22))) (sscalar 0.0))))) (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) (\\x31 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar))) (let x37 = sin (tproject1 x31) in tpair x37 (tpair (tproject1 x31) x37))) (\\x38 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar))) (let x49 = tproject1 (tproject1 x38) * cos (tproject1 (tproject2 x38)) in tpair x49 (tpair (tproject1 (tproject1 x38)) x49))) (\\x50 -> tpair (cos (tproject1 (tproject2 x50)) * (sfromR (tproject2 (tproject2 (tproject1 x50))) + tproject1 (tproject1 x50)) + sfromR (tproject1 (tproject2 (tproject1 x50)))) (sscalar 0.0)) (sscalar 1.1) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))"++testSin0Scan1RevPPForComparison :: Assertion+testSin0Scan1RevPPForComparison = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rfromList [sin (sin x0), sin x0, x0]) (rscalar 1.1)+  printAstPretty (simplifyInline a1)+    @?= "rfromS (sscalar 1.0 + (cos (sscalar 1.1) * cos (sin (sscalar 1.1)) + cos (sscalar 1.1)))"++testSin0ScanFwdPP :: Assertion+testSin0ScanFwdPP = do+  resetVarCounter+  let a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscan (\x _a -> sin x) x0+                           (rrepl @1 @Double [2] 42)) (rscalar 1.1)+  printAstPrettyButNested (simplifyInlineContract a1)+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) (\\x7 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (let x17 = tproject1 x7 * cos (sfromR (tproject1 (tproject2 (tproject2 x7)))) in tpair x17 x17)) (\\x18 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (let x28 = tproject1 (tproject1 x18) * cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x18))))) + (sfromR (tproject1 (tproject2 (tproject2 (tproject1 x18)))) * negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x18))))))) * tproject1 (tproject2 x18) in tpair x28 x28)) (\\x29 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (ConvT2 ConvId (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX)))))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)))) (let x38 = sfromR (tproject2 (tproject1 x29)) + tproject1 (tproject1 x29) in tpair (cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x29))))) * x38) (tpair (sscalar 0.0) (tpair (negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x29)))))) * (tproject1 (tproject2 x29) * x38)) (sscalar 0.0))))) (sscalar 1.0) (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) (\\x39 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar))) (let x45 = sin (tproject1 x39) in tpair x45 (tpair (tproject1 x39) x45))) (\\x46 -> tconvert (ConvCmp (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) ConvId) (ConvCmp (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) ConvId)))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar))) (let x50 = tproject1 (tproject1 x46) * cos (tproject1 (tproject2 x46)) in tpair x50 (tpair (tproject1 (tproject1 x46)) x50))) (\\x51 -> tpair (cos (tproject1 (tproject2 x51)) * (sfromR (tproject2 (tproject2 (tproject1 x51))) + tproject1 (tproject1 x51)) + sfromR (tproject1 (tproject2 (tproject1 x51)))) (sscalar 0.0)) (sscalar 1.1) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))))"++testSin0ScanFwdPPFull :: Assertion+testSin0ScanFwdPPFull = do+  resetVarCounter+  let a1 = rfwd1 @(AstTensor AstMethodLet FullSpan) @Double @0 @1+                 (\x0 -> rscan (\x _a -> sin x) x0+                           (rrepl @1 @Double [2] 42)) (rscalar 1.1)+  printAstPrettyButNested (simplifyInlineContract a1)+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) (\\x7 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (let x17 = sfromR (tproject1 x7) * cos (sfromR (tproject1 (tproject2 (tproject2 x7)))) in tpair x17 x17)) (\\x18 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX)) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (let x28 = sfromR (tproject1 (tproject1 x18)) * cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x18))))) + (sfromR (tproject1 (tproject2 (tproject2 (tproject1 x18)))) * negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x18))))))) * sfromR (tproject1 (tproject2 x18)) in tpair x28 x28)) (\\x29 -> tconvert (ConvCmp (ConvT2 ConvId (ConvT2 (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX) (ConvT2 ConvId (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX)))) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp (ConvXR STKScalar) ConvSX))))) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)))) (let x38 = tproject2 (tconvert (ConvT2 (ConvCmp ConvXS (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvRX)) (ConvCmp ConvXS (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvRX))) (STKProduct (STKR (SNat @0) STKScalar) (STKR (SNat @0) STKScalar)) (tproject1 x29)) + tproject1 (tconvert (ConvT2 (ConvCmp ConvXS (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvRX)) (ConvCmp ConvXS (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvRX))) (STKProduct (STKR (SNat @0) STKScalar) (STKR (SNat @0) STKScalar)) (tproject1 x29)) in tpair (cos (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x29))))) * x38) (tpair (sscalar 0.0) (tpair (negate (sin (sfromR (tproject1 (tproject2 (tproject2 (tproject2 x29)))))) * (sfromR (tproject1 (tproject2 x29)) * x38)) (sscalar 0.0))))) (rfromS (sscalar 1.0)) (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) (\\x39 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKProduct (STKR (SNat @0) STKScalar) (STKS [] STKScalar))) (let x45 = sin (sfromR (tproject1 x39)) in tpair x45 (tpair (tproject1 x39) x45))) (\\x46 -> tconvert (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvT2 ConvId (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKProduct (STKR (SNat @0) STKScalar) (STKS [] STKScalar))) (let x50 = sfromR (tproject1 (tproject1 x46)) * cos (sfromR (tproject1 (tproject2 x46))) in tpair x50 (tpair (tproject1 (tproject1 x46)) x50))) (\\x51 -> tconvert (ConvCmp (ConvT2 ConvId (ConvCmp (ConvXS' (FTKS [] FTKScalar)) ConvRX)) (ConvT2 (ConvCmp (ConvXR STKScalar) ConvSX) (ConvCmp (ConvXR STKScalar) ConvSX))) (STKProduct (STKS [] STKScalar) (STKS [] STKScalar)) (tpair (cos (sfromR (tproject1 (tproject2 x51))) * (sfromR (tproject2 (tproject2 (tproject1 x51))) + sfromR (tproject1 (tproject1 x51))) + sfromR (tproject1 (tproject2 (tproject1 x51)))) (sscalar 0.0))) (rfromS (sscalar 1.1)) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))))"++testSin0Scan1Rev2PP1 :: Assertion+testSin0Scan1Rev2PP1 = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscan (\x a -> sin x - a) x0+                           (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7]))) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sscalar 1.0 + tproject1 (tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.0,7.0]))))) (sconcrete (sfromListLinear [2] [5.0,7.0]))))))"++testSin0Scan1Rev2PPA :: Assertion+testSin0Scan1Rev2PPA = do+  resetVarCounter+  let art = revArtifactAdapt+                 UseIncomingCotangent+                 (\x0 -> rscan @_ @_ @_ @(TKScalar Double) (\x a -> sin x - a) x0+                           (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7])))+                 (FTKR ZSR FTKScalar)+  printArtifactPretty (simplifyArtifact art)+    @?= "\\dret x1 -> rfromS (sfromR dret !$ [0] + tproject1 (tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sslice (SNat @1) (SNat @2) (sfromR dret)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> x1 (sconcrete (sfromListLinear [2] [5.0,7.0]))))) (sconcrete (sfromListLinear [2] [5.0,7.0]))))))"++testSin0Scan1Rev2PPForComparison :: Assertion+testSin0Scan1Rev2PPForComparison = do+  resetVarCounter+  let art = revArtifactAdapt+                 UseIncomingCotangent+                 (\x0 -> rfromList [sin (sin x0 - rscalar 5) - rscalar 7, sin x0 - rscalar 5, x0])+                 (FTKR ZSR FTKScalar)+  printArtifactPretty @_ @(TKR 1 Double) (simplifyArtifact art)+    @?= "\\dret x1 -> rfromS (cos (sfromR x1) * (cos (sscalar (-5.0) + sin (sfromR x1)) * sfromR dret !$ [0]) + (cos (sfromR x1) * sfromR dret !$ [1] + sfromR dret !$ [2]))"++testSin0Scan1Rev2 :: Assertion+testSin0Scan1Rev2 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [] [1.1961317861865948] :: Concrete (TKR 0 Double))+    (rev' (\x0 -> rscan (\x a -> sin x - a) x0+                    (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7]))) (rscalar 1.1))++testSin0Scan1Rev2ForComparison :: Assertion+testSin0Scan1Rev2ForComparison = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [] [1.1961317861865948] :: Concrete (TKR 0 Double))+    (rev' (\x0 -> rfromList [sin (sin x0 - rscalar 5) - rscalar 7, sin x0 - rscalar 5, x0]) (rscalar 1.1))++testSin0Scan1Rev3PP0 :: Assertion+testSin0Scan1Rev3PP0 = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscan (\x a -> sin x - a) x0+                           (rfromList [x0 * rscalar 5, x0 * rscalar 7])) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (let v11 = tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001])))) in sscalar 1.0 + (sscalar 5.0 * tproject2 v11 !$ [0] + (sscalar 7.0 * tproject2 v11 !$ [1] + tproject1 v11)))"+++testSin0Scan1Rev3PPForComparison :: Assertion+testSin0Scan1Rev3PPForComparison = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rfromList [sin (sin x0 - x0 * rscalar 5) - x0 * rscalar 7, sin x0 - x0 * rscalar 5, x0]) (rscalar 1.1)+  printAstPretty (simplifyInline a1)+    @?= "rfromS (let x9 = cos (sscalar (-5.5) + sin (sscalar 1.1)) in sscalar (-11.0) + (cos (sscalar 1.1) * x9 + (sscalar (-5.0) * x9 + cos (sscalar 1.1))))"++testSin0ScanFwd3PP :: Assertion+testSin0ScanFwd3PP = do+  resetVarCounter+  let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Double)+      a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscan (\x a -> sin x - a) x0+                           (rfromList [x0 * rscalar 5, x0 * rscalar 7])) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sfromListLinear [2] [5.0,7.0])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))))))"++testSin0Scan1Rev3 :: Assertion+testSin0Scan1Rev3 = do+  assertEqualUpToEpsilon' 1e-5+    (ringestData [] [-10.076255083995068] :: Concrete (TKR 0 Double))+    (rev' (\x0 -> rscan (\x a -> sin x - a) x0+                           (rfromList [x0 * rscalar 5, x0 * rscalar 7])) (rscalar 1.1))++testSin0Scan1Rev3ForComparison :: Assertion+testSin0Scan1Rev3ForComparison = do+  assertEqualUpToEpsilon' 1e-5+    (ringestData [] [-10.076255083995068] :: Concrete (TKR 0 Double))+    (rev' (\x0 -> rfromList [sin (sin x0 - x0 * rscalar 5) - x0 * rscalar 7, sin x0 - x0 * rscalar 5, x0]) (rscalar 1.1))++testSin0Scan0jvp :: Assertion+testSin0Scan0jvp = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [1] [1.0])+    (rfwd1 @Concrete @Double @0 @1+    (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 1 Double)+         f x0 = rscan (\x _a -> sin x)+                      x0 (rrepl @_ @Double (0 :$: ZSR) 0)+     in f) (rscalar 1.1))++testSin0Scan1jvp :: Assertion+testSin0Scan1jvp = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2] [1.0,0.4535961214255773])+    (rfwd1 @Concrete @Double @0 @1+    (\x0 -> rscan (\x _a -> sin x)+                  x0 (rrepl @1 @Double [1] 42))+          (rscalar 1.1))++testSin0Scan1FwdForComparison :: Assertion+testSin0Scan1FwdForComparison = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2] [1.0,0.4535961214255773])+    (rfwd1 @Concrete @Double @0 @1+    (\x0 -> rfromList [x0, sin x0]) (rscalar 1.1))++testSin0Scan8jvp :: Assertion+testSin0Scan8jvp = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [4,2,5] [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445])+    (rfwd1 @Concrete @Double @0 @3+       (\a0 -> rscan (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                        (rreplicate 3 a0)) (rscalar 1.1))++testSin0Scan8fwd2 :: Assertion+testSin0Scan8fwd2 = do+  let h :: ADVal Concrete (TKR 0 Double) -> ADVal Concrete (TKR 3 Double)+      h = rfwd1 @(ADVal Concrete) @Double @0 @3+        (\a0 -> rscan (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                        (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                        (rreplicate 3 a0))+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [] [285.95794829475744])+    (cgrad (kfromR . rsum0 . h) (rscalar 1.1))++testUnitriangular0PP :: Assertion+testUnitriangular0PP = do+  resetVarCounter+  let k = 1000000+      a1 = rbuild1 @1 @(TKScalar Double) @(AstTensor AstMethodLet PrimalSpan) k+           $ \i -> rbuild1 k+           $ \j -> ifH (i <=. j) (rscalar 0) (rscalar 1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sgather (sconcrete (sfromListLinear [2] [0.0,1.0])) (\\[i3, i2] -> [ifH (0 <=. i2 + negate i3) 0 1]))"++unitriangular1 :: (KnownNat k, GoodScalar rk, ADReady target)+               => Int -> IShR k -> target (TKR (2 + k) rk)+unitriangular1 k sh =+  rbuild1 k $ \i ->+    rbuild1 k $ \j ->+      ifH (i <=. j) (rreplicate0N sh (rscalar 0)) (rreplicate0N sh (rscalar 1))++testUnitriangular1PP :: Assertion+testUnitriangular1PP = do+  resetVarCounter+  let sh = 2 :$: 3 :$: 6 :$: ZSR+      k = 10+      a1 = unitriangular1 @3 @Double @(AstTensor AstMethodLet PrimalSpan) k sh+  printAstPretty (simplifyInline a1)+    @?= "rfromS (sgather (sconcrete (sfromListLinear [2,2,3,6] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])) (\\[i3, i2] -> [ifH (0 <=. i2 + negate i3) 0 1]))"++unitriangular2 :: (KnownNat k, GoodScalar rk, ADReady target)+               => Int -> IShR k -> target (TKR (2 + k) rk)+unitriangular2 k sh =+  rgather @_ @_ @1 (k :$: k :$: sh)+          (rfromList [ rreplicate0N sh (rscalar 0)+                     , rreplicate0N sh (rscalar 1) ])+          (\(i :.: j :.: ZIR) -> ifH (i <. j) 0 1 :.: ZIR)++testUnitriangular2PP :: Assertion+testUnitriangular2PP = do+  resetVarCounter+  let sh = 2 :$: 3 :$: 6 :$: ZSR+      k = 10+      a1 = unitriangular2 @3 @Double @(AstTensor AstMethodLet PrimalSpan) k sh+  printAstPretty (simplifyInline a1)+    @?= "rfromS (sgather (sconcrete (sfromListLinear [2,2,3,6] [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])) (\\[i1, i2] -> [ifH (0 <=. i1 + negate i2) 1 0]))"++testSin0rmapAccumRD0S :: Assertion+testSin0rmapAccumRD0S = do+  assertEqualUpToEpsilon 1e-10+    (srepl 1)+    (grad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f x0 = tproject1 $ tmapAccumR (Proxy @f) (SNat @0)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g (TKS '[] Double)+                                 -> g (TKProduct (TKS '[] Double) (TKS '[] Double))+                               g x _a = tpair (sin x) (sin x)+                           in g)+                          x0+                          (srepl 0)+            in kfromS . f) (srepl 1.1))++testSin0rmapAccumRD00SC :: Assertion+testSin0rmapAccumRD00SC = do+  assertEqualUpToEpsilon 1e-10+    (srepl 1)+    (cgrad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+                f x0 = tproject1 $ tmapAccumL (Proxy @f) (SNat @0)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g (TKS '[] Double)+                                 -> g (TKProduct (TKS '[] Double) (TKS '[] Double))+                               g x _a = tpair (sin x) (sin x)+                           in g)+                          x0+                          (srepl 0)+             in kfromS . f) (srepl 1.1))++testSin0rmapAccumRD00S0 :: Assertion+testSin0rmapAccumRD00S0 = do+  assertEqualUpToEpsilon 1e-10+    (srepl 1)+    (grad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f x0 = tproject1 $ tmapAccumL (Proxy @f) (SNat @0)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g TKUnit+                                 -> g (TKProduct (TKS '[] Double) (TKS '[] Double))+                               g x _a = tpair (sin x) (sin x)+                           in g)+                          x0+                          (treplicate (SNat @0) stkUnit tunit)+            in kfromS . f) (srepl 1.1))++testSin0rmapAccumRD00S :: Assertion+testSin0rmapAccumRD00S = do+  assertEqualUpToEpsilon 1e-10+    (srepl 8.621412119476068e-2)+    (grad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f x0 = tproject1 $ tmapAccumR (Proxy @f) (SNat @7)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g TKUnit+                                 -> g (TKProduct (TKS '[] Double) (TKS '[] Double))+                               g x _a = tpair (sin x) (sin x)+                           in g)+                          x0+                          (treplicate (SNat @7) stkUnit tunit)+            in kfromS . f) (srepl 1.1))++testSin0rmapAccumRD00S7 :: Assertion+testSin0rmapAccumRD00S7 = do+  assertEqualUpToEpsilon 1e-10+    (srepl 1.4091291405664697)+    (grad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[7] Double)+               f x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @7)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g TKUnit+                                 -> g (TKProduct (TKS '[] Double) (TKS '[] Double))+                               g x _a = tpair (sin x) (sin x)+                           in g)+                          x0+                          (treplicate (SNat @7) stkUnit tunit)+            in kfromS . ssum0 . f) (srepl 1.1))++testSin0rmapAccumRD00SCacc0 :: Assertion+testSin0rmapAccumRD00SCacc0 = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (cvjp (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)+               f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @0)+                          ftkUnit+                          ftkUnit+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g (TKS '[] Double)+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (srepl 0)+            in f) (srepl 1.1) (sfromList0N []))++{- TODO: crashes due to a zero dimension+testSin0rmapAccumRD00SCacc0 :: Assertion+testSin0rmapAccumRD00SCacc0 = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (cgrad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)+                f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @0)+                          ftkUnit+                          ftkUnit+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g (TKS '[] Double)+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (srepl 0)+             in kfromS . ssum0 . f) (srepl 1.1))+-}++testSin0rmapAccumRD00SCacc :: Assertion+testSin0rmapAccumRD00SCacc = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (cgrad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[7] Z1)+                f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @7)+                          ftkUnit+                          ftkUnit+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g (TKS '[] Double)+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (srepl 0)+             in kfromS . ssum0 . f) (srepl 1.1))++testSin0rmapAccumRD00Sacc0 :: Assertion+testSin0rmapAccumRD00Sacc0 = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (grad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)+               f _x0 = tproject2 $ tmapAccumL (Proxy @f) (SNat @0)+                          ftkUnit+                          ftkUnit+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g (TKS '[] Double)+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (srepl 0)+             in kfromS . ssum0 . f) (srepl 1.1))++testSin0rmapAccumRD00Sacc :: Assertion+testSin0rmapAccumRD00Sacc = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (grad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[7] Z1)+               f _x0 = tproject2 $ tmapAccumL (Proxy @f) (SNat @7)+                          ftkUnit+                          ftkUnit+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g (TKS '[] Double)+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (srepl 0)+             in kfromS . ssum0 . f) (srepl 1.1))++testSin0rmapAccumRD00SCall0 :: Assertion+testSin0rmapAccumRD00SCall0 = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (cvjp (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)+               f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @0)+                          ftkUnit+                          ftkUnit+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g TKUnit+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (treplicate (SNat @0) stkUnit tunit)+            in f) (srepl 1.1) (sfromList0N []))++{- TODO: crashes due to a zero dimension+testSin0rmapAccumRD00SCall00 :: Assertion+testSin0rmapAccumRD00SCall00 = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (cgrad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)+                f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @0)+                          ftkUnit+                          ftkUnit+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g TKUnit+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (treplicate (SNat @0) stkUnit tunit)+             in kfromS . ssum0 . f) (srepl 1.1))+-}+testSin0rmapAccumRD00SCall :: Assertion+testSin0rmapAccumRD00SCall = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (cgrad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[7] Z1)+                f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @7)+                          ftkUnit+                          ftkUnit+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g TKUnit+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (treplicate (SNat @7) stkUnit tunit)+            in kfromS . ssum0 . f) (srepl 1.1))++testSin0rmapAccumRD00Sall0 :: Assertion+testSin0rmapAccumRD00Sall0 = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (grad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[0] Z1)+               f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @0)+                          ftkUnit+                          ftkUnit+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g TKUnit+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (treplicate (SNat @0) stkUnit tunit)+             in kfromS . ssum0 . f) (srepl 1.1))++testSin0rmapAccumRD00Sall :: Assertion+testSin0rmapAccumRD00Sall = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0)+    (grad (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[7] Z1)+               f _x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @7)+                          ftkUnit+                          ftkUnit+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g TKUnit -> g TKUnit+                                 -> g (TKProduct TKUnit TKUnit)+                               g x _a = tpair x tunit+                           in g)+                          tunit+                          (treplicate (SNat @7) stkUnit tunit)+             in kfromS . ssum0 . f) (srepl 1.1))++testSin0rmapAccumRD0R :: Assertion+testSin0rmapAccumRD0R = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1)+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f x0 = tproject1 $ tmapAccumR (Proxy @f) (SNat @0)+                          (FTKR ZSR FTKScalar)+                          (FTKR ZSR FTKScalar)+                          (FTKR ZSR FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g (TKR 0 Double) -> g (TKR 0 Double)+                                 -> g (TKProduct (TKR 0 Double) (TKR 0 Double))+                               g x _a = tpair (sin x) (sin x)+                           in g)+                          x0+                          (rrepl [0] 0)+            in f) (rscalar 1.1))++testSin0rmapAccumRD01SN :: Assertion+testSin0rmapAccumRD01SN = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.4535961214255773)+    (rev' (let f :: forall f. ADReady f+                 => f (TKS '[] Double) -> f (TKS '[1] Double)+               f x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @1)+                          (FTKProduct (FTKS ZSS FTKScalar) (FTKS ZSS FTKScalar))+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g (TKProduct (TKS '[] Double)+                                                 (TKS '[] Double))+                                 -> g (TKS '[] Double)+                                 -> g (TKProduct (TKProduct (TKS '[] Double)+                                                            (TKS '[] Double))+                                                 (TKS '[] Double))+                               g xh _a = let x = tproject2 xh+                                         in tpair (tpair (sin x) (sin x)) (sin x)+                           in g)+                          (tpair (srepl 3) x0)+                          (srepl 0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0rmapAccumRD01SN3 :: Assertion+testSin0rmapAccumRD01SN3 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.4535961214255773)+    (rev' (let f :: forall f. ADReady f+                 => f (TKS '[] Double) -> f (TKS '[1, 3] Double)+               f x0 = tproject2 $ tmapAccumR (Proxy @f) (SNat @1)+                          (FTKS ZSS FTKScalar)+                          (FTKS (SNat @3 :$$ ZSS) FTKScalar)+                          (FTKS (SNat @2 :$$ ZSS) FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double)+                                 -> g (TKS '[2] Double)+                                 -> g (TKProduct (TKS '[] Double)+                                                 (TKS '[3] Double))+                               g x _a = tpair (sin x)+                                              (sreplicate @3 (sin x / srepl 3))+                           in g)+                          x0+                          (srepl 0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0rmapAccumRD01SN5 :: Assertion+testSin0rmapAccumRD01SN5 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.4535961214255773)+    (rev' (let f :: forall f. ADReady f+                 => f (TKS '[] Double) -> f (TKS '[1, 3] Double)+               f x0 = tproject2 $ tproject2 $ tmapAccumR (Proxy @f) (SNat @1)+                          (FTKS ZSS FTKScalar)+                          (FTKProduct (FTKS (SNat @3 :$$ ZSS) FTKScalar) (FTKS (SNat @3 :$$ ZSS) FTKScalar))+                          (FTKProduct (FTKProduct (FTKS (SNat @2 :$$ ZSS) FTKScalar) (FTKS (SNat @2 :$$ ZSS) FTKScalar))+                                      (FTKProduct (FTKS (SNat @2 :$$ ZSS) FTKScalar) (FTKS (SNat @2 :$$ ZSS) FTKScalar)))+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double)+                                 -> g (TKProduct (TKProduct (TKS '[2] Double)+                                                            (TKS '[2] Double))+                                                 (TKProduct (TKS '[2] Double)+                                                            (TKS '[2] Double)))+                                 -> g (TKProduct (TKS '[] Double)+                                                 (TKProduct (TKS '[3] Double)+                                                            (TKS '[3] Double)))+                               g x a =+                                 tpair (sin x+                                           - smaxIndex+                                               @2 @'[]+                                               (tproject2 $ tproject1 a))+                                   (tpair (sreplicate @3+                                             (sindex0 @'[2]+                                                       (tproject1 $ tproject2 a) [1]+                                              / sin x / srepl 3))+                                          (sreplicate @3+                                             (ssum @2 (tproject2 $ tproject1 a)+                                              + sin x / srepl 3)))+                           in g)+                          x0+                          (tpair (tpair (srepl 0) (srepl 0))+                                 (tpair (srepl 0) (srepl 0)))+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0rmapAccumRD01SN51 :: Assertion+testSin0rmapAccumRD01SN51 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar (-69.90586521651421))+    (rev' (let f :: forall f. ADReady f+                 => f (TKS '[] Double) -> f (TKS '[] Double)+               f x0 = (\res -> ssum @6 (tproject1 $ tproject1 res)+                               + ssum0 @'[6, 5, 4, 3]+                                   (tproject2 res))+                      $ tbuild1 @f (SNat @6) knownSTK $ \j ->+                         tmapAccumR (Proxy @f) (SNat @5)+                          (FTKProduct (FTKS ZSS FTKScalar)+                                      (FTKS (SNat @3 :$$ ZSS) FTKScalar))+                          (FTKS (SNat @4 :$$ SNat @3 :$$ ZSS) FTKScalar)+                          (FTKProduct (FTKProduct (FTKS (SNat @2 :$$ ZSS) FTKScalar) (FTKS (SNat @3 :$$ ZSS) FTKScalar))+                                      (FTKProduct (FTKS (SNat @2 :$$ ZSS) FTKScalar) (FTKS (SNat @2 :$$ ZSS) FTKScalar)))+                          (let g :: forall g. ADReady g+                                 => g (TKProduct (TKS '[] Double)+                                                 (TKS '[3] Double))+                                 -> g (TKProduct (TKProduct (TKS '[2] Double)+                                                            (TKS '[3] Double))+                                                 (TKProduct (TKS '[2] Double)+                                                            (TKS '[2] Double)))+                                 -> g (TKProduct (TKProduct (TKS '[] Double)+                                                            (TKS '[3] Double))+                                                 (TKS '[4, 3] Double))+                               g xh a =+                                 let x = tproject1 xh+                                     x1 = tproject2 xh+                                 in tpair (tpair+                                            (sin x+                                           - smaxIndex+                                               @2 @'[]+                                               (tproject2 $ tproject2 a))+                                            (sreplicate @3+                                             (sindex0 @'[2]+                                                       (tproject1 $ tproject2 a) [1]+                                              / sin x / srepl 3)))+                                         (sbuild1 @4 $ \i ->+                                             (tproject2 $ tproject1 a)+                                             - sin x1 / sreplicate @3+                                                     (srepl 1 + sfromIndex0 i))+                           in g)+                          (tpair (x0 / (srepl 1 + sfromIndex0 j))+                                 (sreplicate @3 x0))+                          (tpair (tpair (srepl 1) (sreplicate0N @'[5, 3]+                                               (sfromIndex0 j)))+                                 (tpair (srepl 3) (srepl 4)))+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0rmapAccumRD01SN531a :: Assertion+testSin0rmapAccumRD01SN531a = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3]+       [1.8478609886246988,-22.194216099801963,-40.72162125038692])+    (rev' (let f :: forall f. ADReady f+                 => f (TKS '[3] Double) -> f (TKS '[2, 2, 2, 3] Double)+               f x0 = (\res -> srepl 2 - sreplicate @2 (tproject1 $ tproject1 res)+                               - tproject2 res)+                      $ tbuild1 @f (SNat @2) knownSTK $ \i ->+                       (tbuild1 @f (SNat @2) knownSTK $ \j ->+                         tmapAccumR (Proxy @f) (SNat @2)+                          (FTKProduct (FTKS (SNat @3 :$$ ZSS) FTKScalar)+                                      (FTKS (SNat @6 :$$ ZSS) FTKScalar))+                          (FTKS (SNat @3 :$$ ZSS) FTKScalar)+                          (FTKProduct (FTKS (SNat @1 :$$ ZSS) FTKScalar)+                                      (FTKS (SNat @3 :$$ ZSS) FTKScalar))+                          (let g :: forall g. ADReady g+                                 => g (TKProduct (TKS '[3] Double)+                                                 (TKS '[6] Double))+                                 -> g (TKProduct (TKS '[1] Double)+                                                 (TKS '[3] Double))+                                 -> g (TKProduct (TKProduct (TKS '[3] Double)+                                                            (TKS '[6] Double))+                                                 (TKS '[3] Double))+                               g xh a =+                                 let x = tproject1 xh+                                     x2 = tproject2 xh+                                 in tpair (tpair+                                            (sfromList+                                             [srepl 0.01, ssum @6 x2, srepl 0.3]+                                           - sin x - tproject2 a)+                                            (srepl 1 - x2+                                           - sreplicate @6+                                                 (ssum (sin x - tproject2 a))))+                                         (srepl 1 - sreplicate @3+                                             (ssum @1 (tproject1 a))+                                           - sin x / srepl 3+                                           - sreplicate @3+                                             (sindex0 @'[3]+                                                       (tproject2 a) [1]+                                             - smaxIndex+                                                 @3 @'[]+                                                 (sin x / srepl 3)))+                           in g)+                          (tpair (x0 / (srepl 1 + sreplicate @3 (sfromIndex0 j)))+                                 (sreplicate @6 (sfromIndex0 i)+                                          - sflatten (sappend x0 x0)))+                          (tpair (sfromList [srepl (-0.1), sreshape @'[] @'[1] $ sfromIndex0 j])+                                 ((sfromList0N+                                           [sscalar 0.4, sscalar (-0.01), sscalar (-0.3), sfromIndex0 i, sscalar 0.5, sscalar 1.3]))))+           in rfromS . f . sfromR) (ringestData [3] [1.1, 2, 3.14]))++testSin0rmapAccumRD01SN531b0 :: Assertion+testSin0rmapAccumRD01SN531b0 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 4)+    (rev' (let f :: forall f. ADReady f+                 => f (TKR 0 Double)+                 -> f (TKR 2 Double)+               f x0 = rfromS $ tproject1+                      $ tbuild1 @f (SNat @2) knownSTK $ \_i ->+                       (tbuild1 @f (SNat @2) knownSTK $ \_j ->+                         tmapAccumR (Proxy @f) (SNat @0)+                          (FTKS ZSS FTKScalar)+                          ftkUnit+                          (FTKR ZSR FTKScalar)+                          (let h :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g (TKR 0 Double)+                                 -> g (TKProduct (TKS '[] Double) TKUnit)+                               h xh _a = tpair xh tunit+                           in h)+                          (sfromR x0)+                          (rconcrete $ Nested.rfromListPrimLinear [0] []))+           in f) (rscalar 1.1))++-- Different result with -O1:+_testSin0rmapAccumRD01SN531b0PP :: Assertion+_testSin0rmapAccumRD01SN531b0PP = do+  resetVarCounter+  let f :: forall f. ADReady f+                 => f (TKR 0 Double)+                 -> f (TKR 2 Double)+      f x0 = rfromS $ tproject1+                      $ tbuild1 @f (SNat @2) knownSTK $ \_i ->+                       (tbuild1 @f (SNat @2) knownSTK $ \_j ->+                         tmapAccumR (Proxy @f) (SNat @0)+                          (FTKS ZSS FTKScalar)+                          ftkUnit+                          (FTKR ZSR FTKScalar)+                          (let h :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g (TKR 0 Double)+                                 -> g (TKProduct (TKS '[] Double) TKUnit)+                               h xh _a = tpair xh tunit+                           in h)+                          (sfromR x0)+                          (rconcrete $ Nested.rfromListPrimLinear [0] []))+  printAstPrettyButNested+    (simplifyInlineContract+     $ f @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))+    @?= "rfromS (sreplicate @2 (sreplicate @2 (tproject1 (tmapAccumRDer (SNat @0) (\\x7 -> tpair (tproject1 x7) Z1) (\\x8 -> tpair (tproject1 (tproject1 x8)) Z1) (\\x11 -> tpair (tproject1 (tproject1 x11)) (sscalar 0.0)) (sscalar 1.1) (sconcrete (sfromListLinear [0] []))))))"++testSin0rmapAccumRD01SN531b0PPj :: Assertion+testSin0rmapAccumRD01SN531b0PPj = do+  resetVarCounter+  let f :: forall f. ADReady f+        => f (TKR 0 Double) -> f (TKR 2 Double)+      f x0 = tlet (+                       (tbuild1 @f (SNat @2) knownSTK $ \i ->+                       (tbuild1 @f (SNat @2) knownSTK $ \j ->+                       (tmapAccumR (Proxy @f) (SNat @0)+                          (FTKS ZSS FTKScalar)+                          ftkUnit+                          (FTKR ZSR FTKScalar)+                          (let h :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g (TKR 0 Double)+                                 -> g (TKProduct (TKS '[] Double) TKUnit)+                               h xh _a = tpair xh tunit+                           in h)+                          (sfromIndex0 (i + j) + sfromR x0)+                          (rconcrete $ Nested.rfromListPrimLinear [0] [])))))+                        $ \ !d -> rfromS $ tproject1 d+  printAstPretty+    (simplifyInlineContract+     $ f @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))+    @?= "rfromS (tproject1 (tmapAccumRDer (SNat @0) <lambda> <lambda> <lambda> (sconcrete (sfromListLinear [2,2] [1.1,2.1,2.1,3.1])) (sconcrete (sfromListLinear [0,2,2] []))))"++testSin0rmapAccumRD01SN531bRPPj :: Assertion+testSin0rmapAccumRD01SN531bRPPj = do+  resetVarCounter+  let f :: forall f. ADReady f+        => f (TKR 0 Double) -> f (TKR 2 Double)+      f x0 = tlet (+                       (tbuild1 @f (SNat @2) knownSTK $ \i ->+                       (tbuild1 @f (SNat @2) knownSTK $ \j ->+                       (tmapAccumR (Proxy @f) (SNat @1)+                          (FTKR ZSR FTKScalar)+                          ftkUnit+                          (FTKR ZSR FTKScalar)+                          (let h :: forall g. ADReady g+                                 => g (TKR 0 Double) -> g (TKR 0 Double)+                                 -> g (TKProduct (TKR 0 Double) TKUnit)+                               h xh _a = tpair xh tunit+                           in h)+                          (rfromIndex0 (i + j) + x0)+                          (rconcrete $ Nested.rfromListPrimLinear [0] [])))))+                        $ \ !d -> tproject1 d+  printAstPretty+    (simplifyInlineContract+     $ f @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))+    @?= "rfromS (tproject1 (tmapAccumRDer (SNat @1) <lambda> <lambda> <lambda> (sconcrete (sfromListLinear [2,2] [1.1,2.1,2.1,3.1])) (sconcrete (sfromListLinear [0,2,2] []))))"++testSin0rmapAccumRD01SN531c :: Assertion+testSin0rmapAccumRD01SN531c = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar (-1.8866871148429984))+    (rev' (let f :: forall f. ADReady f+                 => f (TKS '[] Double) -> f (TKS '[2, 2, 2] Double)+               f x0 = (\res -> srepl 2 - sreplicate @2 (tproject1 res)+                               - tproject2 res)+                      $ tbuild1 @f (SNat @2) knownSTK $ \i ->+                       (tbuild1 @f (SNat @2) knownSTK $ \j ->+                       (tmapAccumR (Proxy @f) (SNat @2)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g (TKS '[] Double)+                                 -> g (TKProduct (TKS '[] Double) (TKS '[] Double))+                               g x a =+                                tpair (sin x - a)+                                      (srepl 1 - sin x / srepl 3 - a)+                           in g)+                          (x0 / (srepl 1 + sfromIndex0 j))+                          (sfromList0N [sscalar 0.4, sfromIndex0 i])))+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0rmapAccumRD01SN531Slice :: Assertion+testSin0rmapAccumRD01SN531Slice = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 4)+    (rev' (let f :: forall f. ADReady f+                 => f (TKS '[] Double) -> f (TKS '[2, 2] Double)+               f x0 = tproject1+                      $ tbuild1 @f (SNat @2) knownSTK $ \_i ->+                       (tbuild1 @f (SNat @2) knownSTK $ \_j ->+                       (tmapAccumR (Proxy @f) (SNat @1)+                          (FTKS ZSS FTKScalar)+                          ftkUnit+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double) -> g TKUnit+                                 -> g (TKProduct (TKS '[] Double) TKUnit)+                               g x _a =+                                tpair x tunit+                           in g)+                          x0+                          (treplicate (SNat @1) stkUnit tunit)))+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0rmapAccumRD01SN55 :: Assertion+testSin0rmapAccumRD01SN55 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 4.1200926532396815)+    (rev' (let f :: forall f. ADReady f+                 => f (TKS '[] Double) -> f (TKS '[5, 3] Double)+               f x0 = (\res -> sreplicate @5 (tproject1 res)+                               * (tproject1 $ tproject2 res)+                               + (tproject2 $ tproject2 res))+                      $ tmapAccumL (Proxy @f) (SNat @5)+                          (FTKS (SNat @3 :$$ ZSS) FTKScalar)+                          (FTKProduct (FTKS (SNat @3 :$$ ZSS) FTKScalar)+                                      (FTKS (SNat @3 :$$ ZSS) FTKScalar))+                          ftkUnit+                          (let g :: forall g. ADReady g+                                 => g (TKS '[3] Double)+                                 -> g TKUnit+                                 -> g (TKProduct (TKS '[3] Double)+                                                 (TKProduct (TKS '[3] Double)+                                                            (TKS '[3] Double)))+                               g x _a =+                                 tpair (sin x - x)+                                       (tpair (sreplicate @3+                                             (sindex0 @'[3] x [1]+                                              - smaxIndex+                                                  @3 @'[]+                                                  (x / sin x / srepl 3)))+                                          (sreplicate @3+                                             (ssum @3 x)+                                           + sin x / srepl 3))+                           in g)+                          (sreplicate @3 x0)+                          (treplicate (SNat @5) stkUnit tunit)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0rmapAccumRD01SN55acc :: Assertion+testSin0rmapAccumRD01SN55acc = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [3] [-21.0,-42.0,-21.0])+    (rev' (let f :: forall f. ADReady f+                 => f (TKS '[3] Double) -> f (TKS '[2, 3] Double)+               f x0 = (\res -> srepl 2 - str (sreplicate @3+                                         $ ssum @7+                                         $ str (tproject2 $ tproject1 $ tproject2 res))+                               - (tproject2 $ tproject2 res))+                      $ tmapAccumR (Proxy @f) (SNat @2)+                          ftkUnit+                          (FTKProduct (FTKProduct (FTKS (SNat @3 :$$ ZSS) FTKScalar)+                                                  (FTKS (SNat @7 :$$ ZSS) FTKScalar))+                                      (FTKS (SNat @3 :$$ ZSS) FTKScalar))+                          (FTKProduct (FTKS (SNat @1 :$$ ZSS) FTKScalar)+                                      (FTKS (SNat @3 :$$ ZSS) FTKScalar))+                          (let g :: forall g. ADReady g+                                 => g TKUnit+                                 -> g (TKProduct (TKS '[1] Double)+                                                 (TKS '[3] Double))+                                 -> g (TKProduct TKUnit+                                         (TKProduct (TKProduct (TKS '[3] Double)+                                                              (TKS '[7] Double))+                                                    (TKS '[3] Double)))+                               g _xh a =+                                let x = sreplicate @3 @_ @_ @g (sscalar 2)+                                in tpair tunit+                                     (tpair+                                        (tpair+                                           (singestData [0.1, 0.2, 0.3]+                                            - sin x - tproject2 a)+                                           (srepl 1 - sreplicate @7+                                                 (ssum+                                                  $ sin x - tproject2 a)))+                                        (sreplicate0N (sscalar 1)+                                         - sreplicate @3+                                             (ssum @1 (tproject1 a))+                                           - sin x / sreplicate0N (sscalar 3)+                                           - sreplicate @3+                                             (sindex0 @'[3]+                                                       (tproject2 a) [1]+                                             - smaxIndex+                                                 @3 @'[]+                                                 (sin x / sreplicate0N (sscalar 3)))))+                           in g)+                          tunit+                          (tpair (singestData [-0.1, 0.23])+                                 (sfromList0N+                                    [sindex0 x0 [1], sscalar (-0.01), sscalar (-0.3), ssum x0, sscalar 0.5, sscalar 1.3]))+           in rfromS . f . sfromR) (ringestData [3] [1.1, 2, 3.14]))++testSin0rmapAccumRD01SN58 :: Assertion+testSin0rmapAccumRD01SN58 = do+  assertEqualUpToEpsilon 1e-10+    (sconcrete $ Nested.sfromListPrimLinear @_ @'[5] knownShS [0,0,0,0,1.1])+    (cjvp (let f :: forall f. ADReady f+                 => f (TKS '[] Double) -> f (TKS '[5] Double)+               f x0 = tproject2+                      $ tmapAccumR (Proxy @f) (SNat @5)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          (FTKS ZSS FTKScalar)+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double)+                                 -> g (TKS '[] Double)+                                 -> g (TKProduct (TKS '[] Double) (TKS '[] Double))+                               g x _a =+                                 tpair (sscalar 1) x+                           in g)+                          x0+                          (srepl 0)+           in f) (srepl 1.1) (srepl 1.1))++testSin0rmapAccumRD01SN7 :: Assertion+testSin0rmapAccumRD01SN7 = do+  assertEqualUpToEpsilon 1e-10+    (srepl 0.4535961214255773)+    (cgrad (let f :: forall f. ADReady f+                  => f (TKS '[] Double) -> f (TKProduct (TKS '[] Double)+                                                 (TKProduct (TKS '[1, 3] Double)+                                                            (TKS '[1, 3] Double)))+                f x0 = tmapAccumR (Proxy @f) (SNat @1)+                          (FTKS ZSS FTKScalar)+                          (FTKProduct (FTKS (SNat @3 :$$ ZSS) FTKScalar)+                                      (FTKS (SNat @3 :$$ ZSS) FTKScalar))+                          (FTKProduct (FTKProduct (FTKS (SNat @2 :$$ ZSS) FTKScalar) (FTKS (SNat @2 :$$ ZSS) FTKScalar))+                                      (FTKProduct (FTKS (SNat @2 :$$ ZSS) FTKScalar) (FTKS (SNat @2 :$$ ZSS) FTKScalar)))+                          (let g :: forall g. ADReady g+                                 => g (TKS '[] Double)+                                 -> g (TKProduct (TKProduct (TKS '[2] Double)+                                                            (TKS '[2] Double))+                                                 (TKProduct (TKS '[2] Double)+                                                            (TKS '[2] Double)))+                                 -> g (TKProduct (TKS '[] Double)+                                                 (TKProduct (TKS '[3] Double)+                                                            (TKS '[3] Double)))+                               g x a =+                                  tpair (sin x+                                           ** smaxIndex+                                                @2 @'[]+                                                (tproject2 $ tproject1 a))+                                    (tpair (sreplicate @3+                                             (sin x / srepl 6+                                              + sindex0 @'[2]+                                                        (tproject1 $ tproject2 a) [1]+                                                / sin x / srepl 3))+                                       (sreplicate @3+                                             (ssum @2 (tproject2 $ tproject1 a)+                                              + sin x / srepl 6)))+                           in g)+                          x0+                          (tpair+                             (tpair (sreplicate0N $ sscalar 0)+                                    (sreplicate0N $ sscalar 0))+                             (tpair (sreplicate0N $ sscalar 0)+                                    (sreplicate0N $ sscalar 0)))+            in tdot0Target (FTKProduct (FTKS ZSS FTKScalar)+                             (FTKProduct (FTKS (SNat @1 :$$ SNat @3 :$$ ZSS) FTKScalar)+                                         (FTKS (SNat @1 :$$ SNat @3 :$$ ZSS) FTKScalar))) (treplTarget 1 (FTKProduct (FTKS ZSS FTKScalar)+                             (FTKProduct (FTKS (SNat @1 :$$ SNat @3 :$$ ZSS) FTKScalar)+                                         (FTKS (SNat @1 :$$ SNat @3 :$$ ZSS) FTKScalar)))) . f @(ADVal Concrete)) (sscalar 1.1))++rscanZip :: forall rn n rn2 n2 target.+            (GoodScalar rn, KnownNat n, GoodScalar rn2, ADReady target)+         => (forall f. ADReady f+             => f (TKR n rn) -> f (TKR n2 rn2) -> f (TKR n rn))+         -> FullShapeTK (TKR n2 rn2)+         -> target (TKR n rn)+         -> target (TKR (1 + n2) rn2)+         -> target (TKR (1 + n) rn)+rscanZip f eftk acc0 es =+  let width = rwidth es+      ftk = tftk knownSTK acc0+  in withSNat width $ \snat ->+    tlet+      (tmapAccumL Proxy snat ftk ftk eftk+         (let g :: forall f. ADReady f+                => f (TKR n rn) -> f (TKR n2 rn2)+                -> f (TKProduct (TKR n rn) (TKR n rn))+              g acc e = tlet (f acc e) (\res -> tpair res res)+          in g)+         acc0 es)+      (\res -> rappend (rfromList [acc0]) (tproject2 res))++testSin0ScanD51 :: Assertion+testSin0ScanD51 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1,1] [319.68688158967257] :: Concrete (TKR 4 Double))+    (rev' (\a0 -> rscanZip (\x a -> rsum+                                 $ atan2H (sin $ rreplicate 5 x)+                                         (rsum $ sin $ rsum+                                          $ rtr $ rreplicate 7+                                          $ rreplicate 2 $ rreplicate 5+                                          $ rsum $ rsum a))+                         (FTKR (8 :$: 3 :$: 1 :$: 1 :$: 1 :$: 1 :$: ZSR) FTKScalar)+                         (rreplicate0N [1,1,1,1] (rscalar 2) * a0)+                         (rreplicate 3 (rreplicate 8 (rreplicate 3 a0)))+                         )+          (ringestData [1,1,1,1] [1.1]))++testSin0ScanD8 :: Assertion+testSin0ScanD8 = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1] [9.532987357352765] :: Concrete (TKR 3 Double))+    (rev' (\a0 -> rscanZip (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum (rreplicate 7 a))))+                       (FTKR (1 :$: 1 :$: 1 :$: ZSR) FTKScalar)+                       (rreplicate 2 (rreplicate 5+                                        (rreplicate0N [1,1,1] (rscalar 2) * a0)))+                       (rreplicate 3 a0))+          (ringestData [1,1,1] [1.1]))++testSin0ScanD8MapAccum :: Assertion+testSin0ScanD8MapAccum = do+  assertEqualUpToEpsilon' 1e-10+    (ringestData [1,1,1] [9.532987357352765] :: Concrete (TKR 3 Double))+    (rev'+       (\a0 -> tproject2+               $ tmapAccumR Proxy (SNat @4)+                   (FTKR (2 :$: 5 :$: 1 :$: 1 :$: 1 :$: ZSR) FTKScalar)+                   (FTKR (2 :$: 5 :$: 1 :$: 1 :$: 1 :$: ZSR) FTKScalar)+                   (FTKR (1 :$: 1 :$: 1 :$: ZSR) FTKScalar)+                   (let g :: forall g. ADReady g+                          => g (TKR 5 Double) -> g (TKR 3 Double)+                          -> g (TKProduct (TKR 5 Double) (TKR 5 Double))+                        g x a =+                          tpair (rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                                x+                    in g)+                      (rreplicate 2 (rreplicate 5+                                      (rreplicate0N [1,1,1] (rscalar 2) * a0)))+                      (rreplicate 4 a0))+       (ringestData [1,1,1] [1.1]))++testSin0ScanD8grad :: Assertion+testSin0ScanD8grad = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [] [9.53298735735276])+    (rrev1 @Concrete @Double @0 @3+        (\a0 -> rscanZip (\x a -> rtr $ rreplicate 5+                                  $ atan2H (rsum (rtr $ sin x))+                                           (rreplicate 2+                                            $ sin (rsum (rreplicate 7 a))))+                         (FTKR ZSR FTKScalar)+                         (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                         (rreplicate 3 a0)) (rscalar 1.1))++testSin0ScanD8rev4 :: Assertion+testSin0ScanD8rev4 = do+  let h :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      h = rrev1 @f @Double @0 @3+        (\a0 -> rscanZip (\x a -> rtr $ rreplicate 5+                                  $ atan2H (rsum (rtr $ sin x))+                                           (rreplicate 2+                                            $ sin (rsum (rreplicate 7 a))))+                         (FTKR ZSR FTKScalar)+                         (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                         (rreplicate 3 a0))+  assertEqualUpToEpsilon' 1e-10+    (ringestData [] [285.95794829475744])+    (rev' h (rscalar 1.1))++testSin0ScanD1RevPP :: Assertion+testSin0ScanD1RevPP = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscanZip (\x _a -> sin x)+                           (FTKR ZSR FTKScalar)+                           x0 (rrepl @1 @Double [2] 42)) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sscalar 1.0 + tproject1 (tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))"++testSin0ScanDFwdPP :: Assertion+testSin0ScanDFwdPP = do+  resetVarCounter+  let a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscanZip (\x _a -> sin x)+                           (FTKR ZSR FTKScalar)+                           x0 (rrepl @1 @Double [2] 42)) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sreplicate [2] 42.0))))) (sconcrete (sreplicate [2] 42.0))))))))"++testSin0ScanD1Rev2PP :: Assertion+testSin0ScanD1Rev2PP = do+  resetVarCounter+  let a1 = rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscanZip (\x a -> sin x - a)+                         (FTKR ZSR FTKScalar)+                         x0 (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7]))) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sscalar 1.0 + tproject1 (tmapAccumRDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 0.0) (tpair (sconcrete (sreplicate [2] 1.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.0,7.0]))))) (sconcrete (sfromListLinear [2] [5.0,7.0]))))))"++testSin0ScanDFwd2PP :: Assertion+testSin0ScanDFwd2PP = do+  resetVarCounter+  let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Double)+      a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscanZip (\x a -> sin x - a)+                         (FTKR ZSR FTKScalar)+                         x0 (rconcrete (Nested.rfromListPrimLinear @Double @1 [2] [5, 7]))) (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sreplicate [2] 0.0)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.0,7.0]))))) (sconcrete (sfromListLinear [2] [5.0,7.0]))))))))"++testSin0ScanDFwd3PP :: Assertion+testSin0ScanDFwd3PP = do+  resetVarCounter+  let a1 :: AstTensor AstMethodLet PrimalSpan (TKR 1 Double)+      a1 = rfwd1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @1+                 (\x0 -> rscanZip (\x a -> sin x - a)+                                (FTKR ZSR FTKScalar)+                                x0 (rfromList [x0 * rscalar 5, x0 * rscalar 7]))+                 (rscalar 1.1)+  printAstPretty (simplifyInlineContract a1)+    @?= "rfromS (sappend (sconcrete (sfromListLinear [1] [1.0])) (sfromR (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sfromListLinear [2] [5.0,7.0])) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @2) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))) (sconcrete (sfromListLinear [2] [5.5,7.700000000000001]))))))))"++testSin0ScanD1jvp :: Assertion+testSin0ScanD1jvp = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [2] [1.0,0.4535961214255773])+    (rfwd1 @Concrete @Double @0 @1+    (\x0 -> rscanZip (\x _a -> sin x)+                   (FTKR ZSR FTKScalar)+                   x0 (rrepl @1 @Double [1] 42))+          (rscalar 1.1))++testSin0ScanD8jvp :: Assertion+testSin0ScanD8jvp = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [4,2,5] [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445])+    (rfwd1 @Concrete @Double @0 @3+       (\a0 -> rscanZip (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum (rreplicate 7 a))))+                      (FTKR ZSR FTKScalar)+                      (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                      (rreplicate 3 a0)) (rscalar 1.1))++testSin0ScanD8fwdMapAccum :: Assertion+testSin0ScanD8fwdMapAccum = do+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [4,2,5] [2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.5864059429583657,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.24026418024701368,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445,-0.2200311410593445])+    (rfwd1 @Concrete @Double @0 @3 @Double+       (\a0 -> rreverse $ tproject2+               $ tmapAccumR Proxy (SNat @4)+                   (FTKR (2 :$: 5 :$: ZSR) FTKScalar)+                   (FTKR (2 :$: 5 :$: ZSR) FTKScalar)+                   (FTKR ZSR FTKScalar)+                   (let g :: forall g. ADReady g+                          => g (TKR 2 Double) -> g (TKR 0 Double)+                          -> g (TKProduct (TKR 2 Double) (TKR 2 Double))+                        g x a =+                         tpair (rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum $ rreplicate 7 a)))+                               x+                    in g)+                      (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                      (rreplicate 4 a0)) (rscalar 1.1))++testSin0ScanD8fwd2 :: Assertion+testSin0ScanD8fwd2 = do+  let h :: ADVal Concrete (TKR 0 Double) -> ADVal Concrete (TKR 3 Double)+      h = rfwd1 @(ADVal Concrete) @Double @0 @3+        (\a0 -> rscanZip (\x a -> rtr $ rreplicate 5+                                 $ atan2H (rsum (rtr $ sin x))+                                         (rreplicate 2+                                          $ sin (rsum (rreplicate 7 a))))+                       (FTKR ZSR FTKScalar)+                       (rreplicate 2 (rreplicate 5 (rscalar 2 * a0)))+                       (rreplicate 3 a0))+  assertEqualUpToEpsilon 1e-10+    (rconcrete $ Nested.rfromListPrimLinear [] [285.95794829475744])+    (cgrad (kfromR . rsum0 . h) (rscalar 1.1))++testSin0FoldNestedS1 :: Assertion+testSin0FoldNestedS1 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2.0504979297616553e-43 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\x a ->+                        sfold (\x2 a2 -> srepl 0.7 * x2 * a2)+                              a (sreplicate @7 x))+                            a0 (sreplicate @3 a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0FoldNestedS1PP :: Assertion+testSin0FoldNestedS1PP = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+      f z = sfold (\x a ->+               sfold (\x2 a2 -> x2 + tan a2)+                     a (sreplicate @22 x))+                  z (sreplicate @11 z)+      g :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+      g = kgrad (kfromS . f) (FTKS ZSS FTKScalar)+  printAstPretty+    (g @(AstTensor AstMethodLet PrimalSpan) (sscalar 1.1))+    @?= "let v6 = tmapAccumRDer (SNat @11) <lambda> <lambda> <lambda> (sscalar 1.0) (tpair (sconcrete (sreplicate [11] Z1)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @11) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sreplicate [11] 1.1))))) (sconcrete (sreplicate [11] 1.1)))) in ssum @11 (tproject2 v6) + tproject1 v6"++testSin0FoldNestedR1PP :: Assertion+testSin0FoldNestedR1PP = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      f z = rfold (\x a ->+               rfold (\x2 a2 -> x2 + tan a2)+                     a (rreplicate 22 x))+                  z (rreplicate 11 z)+      g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)+  printAstPretty+    (g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1))+    @?= "rfromS (let v6 = tmapAccumRDer (SNat @11) <lambda> <lambda> <lambda> (tconvert (ConvCmp (ConvCmp (ConvXR STKScalar) (ConvCmp (ConvXX' (FTKX [] FTKScalar)) ConvSX)) (ConvCmp ConvXS (Conv0X STKScalar))) (STKScalar) 1.0) (tpair (sconcrete (sreplicate [11] Z1)) (tpair (tproject1 (tproject2 (tmapAccumLDer (SNat @11) <lambda> <lambda> <lambda> (sscalar 1.1) (sconcrete (sreplicate [11] 1.1))))) (sconcrete (sreplicate [11] 1.1)))) in ssum @11 (sfromR (tproject2 v6)) + sfromR (tproject1 v6))"++testSin0FoldNestedR0LengthPPs :: Assertion+testSin0FoldNestedR0LengthPPs = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      f z = rfold (\x a -> x + tan a)+                  z (rreplicate 2 z)+      g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)+  length+    (printAstSimple+      (simplifyInlineContract+       $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))+    @?= 4454++testSin0FoldNestedR1LengthPPs :: Assertion+testSin0FoldNestedR1LengthPPs = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      f z = rfold (\x a ->+               rfold (\x2 a2 -> x2 + tan a2)+                     a (rreplicate 2 x))+                  z (rreplicate 2 z)+      g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)+  length+    (printAstSimple+      (simplifyInlineContract+       $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))+    @?= 42039++testSin0FoldNestedR2LengthPPs :: Assertion+testSin0FoldNestedR2LengthPPs = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      f z = rfold (\x a ->+               rfold (\x2 a2 ->+                 rfold (\x3 a3 -> x3 + tan a3)+                       a2 (rreplicate 2 x2))+                     a (rreplicate 2 x))+                  z (rreplicate 2 z)+      g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)+  length+    (printAstSimple+       (simplifyInlineContract+        $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))+    @?= 530520++testSin0FoldNestedR3LengthPPs :: Assertion+testSin0FoldNestedR3LengthPPs = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      f z = rfold (\x a ->+               rfold (\x2 a2 ->+                 rfold (\x3 a3 ->+                   rfold (\x4 a4 -> x4 + tan a4)+                         a3 (rreplicate 2 x3))+                       a2 (rreplicate 2 x2))+                     a (rreplicate 2 x))+                  z (rreplicate 2 z)+      g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)+  length+    (printAstSimple+       (simplifyInlineContract+        $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))+    @?= 8035414++-- Takes 70s, probably due to something (simplification?) forcing all derivs.+_testSin0FoldNestedR4LengthPPs :: Assertion+_testSin0FoldNestedR4LengthPPs = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      f z = rfold (\x a ->+               rfold (\x2 a2 ->+                 rfold (\x3 a3 ->+                   rfold (\x4 a4 ->+                     rfold (\x5 a5 -> x5 + tan a5)+                           a4 (rreplicate 2 x4))+                         a3 (rreplicate 2 x3))+                       a2 (rreplicate 2 x2))+                     a (rreplicate 2 x))+                  z (rreplicate 2 z)+      g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)+  length+    (printAstSimple+       (simplifyInlineContract+        $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))+    @?= 0++_testSin0FoldNestedR5LengthPPs :: Assertion+_testSin0FoldNestedR5LengthPPs = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      f z = rfold (\x a ->+               rfold (\x2 a2 ->+                 rfold (\x3 a3 ->+                   rfold (\x4 a4 ->+                     rfold (\x5 a5 ->+                       rfold (\x6 a6 -> x6 + tan a6)+                             a5 (rreplicate 2 x5))+                           a4 (rreplicate 2 x4))+                         a3 (rreplicate 2 x3))+                       a2 (rreplicate 2 x2))+                     a (rreplicate 2 x))+                  z (rreplicate 2 z)+      g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)+  length+    (printAstSimple+       (simplifyInlineContract+        $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))+    @?= 0++testSin0FoldNestedR2LengthPPsDummy7 :: Assertion+testSin0FoldNestedR2LengthPPsDummy7 = do+  resetVarCounter+  let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      f z = rfold (\x a ->+               rfold (\x2 a2 ->+                 rfold (\_x3 _a3 -> rscalar 7)+                   -- the 7 causes Dummy RepM values+                   -- with the more precise typing of folds+                       a2 (rreplicate 2 x2))+                     a (rreplicate 2 x))+                  z (rreplicate 2 z)+      g :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+      g = kgrad (kfromR . f) (FTKR ZSR FTKScalar)+  length+    (printAstSimple+       (simplifyInlineContract+        $ g @(AstTensor AstMethodLet PrimalSpan) (rscalar 1.1)))+    @?= 232921++testSin0FoldNestedR2Dummy7 :: Assertion+testSin0FoldNestedR2Dummy7 = do+ resetVarCounter+ assertEqualUpToEpsilon' 1e-10+  (rscalar 0 :: Concrete (TKR 0 Double))+  (rev'+   (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+        f z = rfold (\x a ->+               rfold (\x2 a2 ->+                 rfold (\_x3 _a3 -> rscalar 7)+                   -- the 7 causes Dummy RepM values+                   -- with the more precise typing of folds+                       a2 (rreplicate 2 x2))+                     a (rreplicate 2 x))+                  z (rreplicate 2 z)+    in f) (rscalar 0.0001))++testSin0FoldNestedR2Tan :: Assertion+testSin0FoldNestedR2Tan = do+ assertEqualUpToEpsilon' 1e-10+  (rscalar 25.000016360009603 :: Concrete (TKR 0 Double))+  (rev'+   (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+        f z = rfold (\x a ->+                 rfold (\x2 a2 ->+                   rfold (\x3 a3 -> x3 + tan a3)+                         a2 (rreplicate 2 x2))+                       a (rreplicate 2 x))+                    z (rreplicate 2 z)+    in f) (rscalar 0.0001))++testSin0FoldNestedS1FwdFwd0 :: Assertion+testSin0FoldNestedS1FwdFwd0 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2.0504979297616553e-43 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\x a ->+                        sfold (\x2 a2 -> srepl 0.7 * x2 * a2)+                              a (sreplicate @7 x))+                            a0 (sreplicate @3 a0)+           in rfromS . sfwd1 f . sfromR) (rscalar 1.1))++testSin0FoldNestedS1FwdFwd :: Assertion+testSin0FoldNestedS1FwdFwd = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2.0504979297616553e-43 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\x a ->+                        sfold (\x2 a2 ->+                                 x2 * sfwd1 (sfwd1 (\b2 -> srepl 0.7 * b2)) a2)+                              a (sreplicate @7 x))+                            a0 (sreplicate @3 a0)+           in rfwd1 $ rfromS . sfwd1 f . sfromR) (rscalar 1.1))++testSin0FoldNestedS1RevRev :: Assertion+testSin0FoldNestedS1RevRev = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2.0504979297616553e-43 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\x a ->+                        sfold (\x2 a2 ->+                                 x2 * srev1 (srev1 (\b2 -> srepl 0.7 * b2)) a2)+                              a (sreplicate @7 x))+                            a0 (sreplicate @3 a0)+           in rrev1 $ rfromS . srev1 f . sfromR) (rscalar 1.1))++testSin0FoldNestedS2 :: Assertion+testSin0FoldNestedS2 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 3.175389686661287e-207 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\x a ->+                        sfold (\x2 a2 ->+                          sfold (\x3 a3 -> srepl 0.7 * x3 * a3)+                                a2 (sreplicate @4 x2))+                              a (sreplicate @3 x))+                            a0 (sreplicate @2 a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0FoldNestedS3 :: Assertion+testSin0FoldNestedS3 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 7.320500000000004e-4 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\x a ->+                        sfold (\x2 a2 ->+                          sfold (\x3 a3 ->+                            sfold (\x4 a4 -> srepl 0.1 * x4 * a4)+                                  a3 (sreplicate @1 x3))+                                a2 (sreplicate @2 x2))+                              a (sreplicate @1 x))+                            a0 (sreplicate @2 a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0FoldNestedS4 :: Assertion+testSin0FoldNestedS4 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1.2400927000000009e-5 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\x a ->+                        sfold (\x2 a2 ->+                          sfold (\x3 a3 ->+                            sfold (\x4 a4 ->+                              sfold (\x5 a5 -> srepl 0.1 * x5 * a5)+                                    a4 (sreplicate @2 x4))+                                  a3 (sreplicate @1 x3))+                                a2 (sreplicate @1 x2))+                              a (sreplicate @2 x))+                            a0 (sreplicate @1 a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0FoldNestedS5 :: Assertion+testSin0FoldNestedS5 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.22000000000000003 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+               f a0 = sfold (\x a ->+                        sfold (\x2 a2 ->+                          sfold (\x3 a3 ->+                            sfold (\x4 a4 ->+                              sfold (\x5 a5 ->+                                sfold (\x6 a6 -> sscalar 0.1 * x6 * a6)+                                      a5 (sreplicate @1 x5))+                                    a4 (sreplicate @1 x4))+                                  a3 (sreplicate @1 x3))+                                a2 (sreplicate @1 x2))+                              a (sreplicate @1 x))+                            a0 (sreplicate @1 a0)++           in rfromS . f . sfromR) (rscalar 1.1))++testSin0FoldNestedS5grad :: Assertion+testSin0FoldNestedS5grad = do+  let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+      f a0 = sfold (\x a ->+                        sfold (\x2 a2 ->+                          sfold (\x3 a3 ->+                            sfold (\x4 a4 ->+                              sfold (\x5 a5 ->+                                sfold (\x6 a6 -> sscalar 0.1 * x6 * a6)+                                      a5 (sreplicate @1 x5))+                                    a4 (sreplicate @1 x4))+                                  a3 (sreplicate @1 x3))+                                a2 (sreplicate @1 x2))+                              a (sreplicate @1 x))+                            a0 (sreplicate @1 a0)+  assertEqualUpToEpsilon 1e-10+    (srepl 0.22000000000000003)+    (srev1 @Concrete @Double @'[] @'[] f (sscalar 1.1))++testSin0FoldNestedS5jvp :: Assertion+testSin0FoldNestedS5jvp = do+  let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[] Double)+      f a0 = sfold (\x a ->+                        sfold (\x2 a2 ->+                          sfold (\x3 a3 ->+                            sfold (\x4 a4 ->+                              sfold (\x5 a5 ->+                                sfold (\x6 a6 -> sscalar 0.1 * x6 * a6)+                                      a5 (sreplicate @1 x5))+                                    a4 (sreplicate @1 x4))+                                  a3 (sreplicate @1 x3))+                                a2 (sreplicate @1 x2))+                              a (sreplicate @1 x))+                            a0 (sreplicate @1 a0)+  assertEqualUpToEpsilon 1e-10+    (srepl 0.22000000000000003)+    (sfwd1 @Concrete @Double @'[] @'[] f (sscalar 1.1))++testSin0FoldNestedSi :: Assertion+testSin0FoldNestedSi = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar  (-0.20775612781643243) :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKS '[] Double) -> f (TKS '[3] Double)+               f a0 = sfold (\x a -> atan2H+                                       (sscan (+) (ssum x)+                                          (sscan (*) (srepl 2)+                                                 (sreplicate @1 a)))+                                       (sscan (\x1 a1 ->+                                          sfold (\x2 a2 ->+                                            sfold (\x3 a3 ->+                                                     srepl 0.001 * (x3 * a3 - x3))+                                                  a2 (sscan (+) x2+                                                        (sreplicate @3 a2)))+                                                x1 (sreplicate @1 a1))+                                              a (sscan (-) (srepl 0)+                                                   (sslice (SNat @0) (SNat @1) SNat x))))+                            (sreplicate @3 $ srepl 2 * a0) (sreplicate @2 a0)+           in rfromS . f . sfromR) (rscalar 1.1))++testSin0FoldNestedR1 :: Assertion+testSin0FoldNestedR1 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2.0504979297616553e-43 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 -> rscalar 0.7 * x2 * a2)+                              a (rreplicate 7 x))+                            a0 (rreplicate 3 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedR1RevFwd :: Assertion+testSin0FoldNestedR1RevFwd = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 3.175389686661287e-207 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 ->+                                 x2 * rfwd1 (rrev1 (\b2 -> rscalar 0.7 * b2)) a2)+                              a (rreplicate 4 x))+                            a0 (rreplicate 2 a0)+           in rrev1 $ rfwd1 f) (rscalar 1.1))++testSin0FoldNestedR2 :: Assertion+testSin0FoldNestedR2 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 3.175389686661287e-207 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 ->+                          rfold (\x3 a3 -> rscalar 0.7 * x3 * a3)+                                a2 (rreplicate 4 x2))+                              a (rreplicate 3 x))+                            a0 (rreplicate 2 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedR2RevFwd :: Assertion+testSin0FoldNestedR2RevFwd = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 3.175389686661287e-207 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 ->+                          rfold (\x3 a3 ->+                                   x3 * rrev1 (rfwd1 (rrev1 (\b3 ->+                                                               rscalar 0.7 * b3))) a3)+                                a2 (rreplicate 4 x2))+                              a (rreplicate 3 x))+                            a0 (rreplicate 2 a0)+           in rfwd1 $ rrev1 $ rfwd1 f) (rscalar 1.1))++testSin0FoldNestedR3 :: Assertion+testSin0FoldNestedR3 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 7.320500000000004e-4 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 ->+                          rfold (\x3 a3 ->+                            rfold (\x4 a4 -> rscalar 0.1 * x4 * a4)+                                  a3 (rreplicate 1 x3))+                                a2 (rreplicate 2 x2))+                              a (rreplicate 1 x))+                            a0 (rreplicate 2 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedR4 :: Assertion+testSin0FoldNestedR4 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1.2400927000000009e-5 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 ->+                          rfold (\x3 a3 ->+                            rfold (\x4 a4 ->+                              rfold (\x5 a5 -> rscalar 0.1 * x5 * a5)+                                    a4 (rreplicate 2 x4))+                                  a3 (rreplicate 1 x3))+                                a2 (rreplicate 1 x2))+                              a (rreplicate 2 x))+                            a0 (rreplicate 1 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedR41 :: Assertion+testSin0FoldNestedR41 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 0.22000000000000003 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 ->+                          rfold (\x3 a3 ->+                            rfold (\x4 a4 ->+                              rfold (\x5 a5 -> rscalar 0.1 * x5 * a5)+                                    a4 (rreplicate 1 x4))+                                  a3 (rreplicate 1 x3))+                                a2 (rreplicate 1 x2))+                              a (rreplicate 1 x))+                            a0 (rreplicate 1 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedR40 :: Assertion+testSin0FoldNestedR40 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1.0 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 ->+                          rfold (\x3 a3 ->+                            rfold (\x4 a4 ->+                              rfold (\x5 a5 -> rscalar 0.1 * x5 * a5)+                                    a4 (rreplicate 0 x4))+                                  a3 (rreplicate 0 x3))+                                a2 (rreplicate 0 x2))+                              a (rreplicate 0 x))+                            a0 (rreplicate 0 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedR400 :: Assertion+testSin0FoldNestedR400 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 1.0 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 ->+                          rfold (\x3 a3 ->+                            rfold (\x4 a4 ->+                              rfold (\_x5 _a5 -> 0)+                                    a4 (rreplicate 0 x4))+                                  a3 (rreplicate 0 x3))+                                a2 (rreplicate 0 x2))+                              a (rreplicate 0 x))+                            a0 (rreplicate 0 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedRi :: Assertion+testSin0FoldNestedRi = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar (-0.20775612781643243) :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 1 Double)+               f a0 = rfold (\x a -> atan2H+                                       (rscan (+) (rsum x)+                                          (rscan (*) (rscalar 2)+                                                 (rreplicate 1 a)))+                                       (rscan (\x1 a1 ->+                                          rfold (\x2 a2 ->+                                            rfold (\x3 a3 ->+                                                     rscalar 0.001 * (x3 * a3 - x3))+                                                  a2 (rscan (+) x2+                                                            (rreplicate 3 a2)))+                                                x1 (rreplicate 1 a1))+                                              a (rscan (-) (rscalar 0) (rslice 0 1 x))))+                            (rreplicate 3 $ rscalar 2 * a0) (rreplicate 2 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedR22 :: Assertion+testSin0FoldNestedR22 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 2.877421010384167e-5 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a ->+                        rfold (\x2 a2 ->+                          rfold (\x3 a3 -> rscalar 0.44 * x3 * a3)+                                a2 (rscan (\x4 a4 -> x4 + a4) x2+                                          (rreplicate 2 a2)))+                              (rfold (\x4 a4 -> x4 * a4) a+                                     (rreplicate 2 x)) (rreplicate 3 x))+                            a0 (rreplicate 2 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedR21 :: Assertion+testSin0FoldNestedR21 = do+  assertEqualUpToEpsilon' 1e-10+    (rscalar 7.667553331540788e-3 :: Concrete (TKR 0 Double))+    (rev' (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a -> tlet (x + a) $ \xpa ->+                          rfold (\x3 a3 -> rscalar 0.1 * x3 * a3)+                                (rfold (\x4 a4 -> x4 * a4) xpa+                                       (rreplicate 2 x))+                                (rscan (\x4 a4 -> x4 + a4) xpa+                                       (rreplicate 2 xpa)))+                            a0 (rreplicate 2 a0)+           in f) (rscalar 1.1))++testSin0FoldNestedR21PP :: Assertion+testSin0FoldNestedR21PP = do+  resetVarCounter+  let a1 =+        rrev1 @(AstTensor AstMethodLet PrimalSpan) @Double @0 @0+          (let f :: forall f. ADReady f => f (TKR 0 Double) -> f (TKR 0 Double)+               f a0 = rfold (\x a -> tlet (x + a) $ \xpa ->+                          rfold (\x3 a3 -> rscalar 0.1 * x3 * a3)+                                (rfold (\x4 a4 -> x4 * a4) xpa+                                       (rreplicate 2 x))+                                (rscan (\x4 a4 -> x4 + a4) xpa+                                       (rreplicate 2 xpa)))+                            a0 (rreplicate 2 a0)+           in f) (rscalar 1.1)+  length (printAstSimple (simplifyInlineContract a1))+    @?= 67598++testSin0revhV :: Assertion+testSin0revhV = do+  let f :: forall g. BaseTensor g+        => g (TKR 0 Double) -> g (TKR 0 Double)+      f = kgrad @_ @Double @g (kfromR . sin) (FTKR ZSR FTKScalar)+  assertEqualUpToEpsilon 1e-10+    (rscalar 0.4535961214255773)+    (f @Concrete (rscalar 1.1))++testSin0revhVPP :: Assertion+testSin0revhVPP = do+  resetVarCounter+  let f :: forall g. BaseTensor g+        => g (TKR 0 Double) -> g (TKR 0 Double)+      f = kgrad @_ @Double @g (kfromR . sin) (FTKR ZSR FTKScalar)+  printAstSimple (f @(AstTensor AstMethodLet PrimalSpan)+                                    (rscalar 1.1))+    @?= "rfromS (cos (sscalar 1.1))"++testSin0revhV4 :: Assertion+testSin0revhV4 = do+  let doms = FTKR ZSR FTKScalar+      doms3 = FTKR (3 :$: ZSR) FTKScalar+      f :: forall g. (BaseTensor g)+        => g (TKR 1 Double) -> g (TKR 1 Double)+      f x =+        rvjp @1 @_ @(TKScalar Double) @g+               (rscanZip const doms (rscalar 5))+               doms3 x (ringestData [4] [1, 2, 3, 4])+  assertEqualUpToEpsilon 1e-10+    (rfromList [rscalar 0, rscalar 0, rscalar 0])+    (cgrad (kfromR . rsum0 . f) (rreplicate 3 (rscalar 1.1)))++testSin0revhV5 :: Assertion+testSin0revhV5 = do+  let doms = FTKR ZSR FTKScalar+      doms3 = FTKS (SNat @3 :$$ ZSS) FTKScalar+      f :: forall g. (BaseTensor g)+        => g (TKS '[3] Double) -> g (TKS '[3] Double)+      f x =+        svjp @_ @_ @(TKScalar Double) @g+               (\v -> sfromR @_ @'[4] $ rscanZip const doms (rscalar 5) (rfromS v))+                doms3 x (singestData @'[4] [1, 2, 3, 4])+  assertEqualUpToEpsilon 1e-10+    (sfromList @3 [sscalar 0, sscalar 0, sscalar 0])+    (cgrad (kfromS . ssum0 . f) (sreplicate @3 (sscalar 1.1)))++testSin0revhV6 :: Assertion+testSin0revhV6 = do+  let doms = FTKR ZSR FTKScalar+      doms3 = FTKR (3 :$: ZSR) FTKScalar+      f :: forall g. (BaseTensor g)+        => g (TKR 1 Double) -> g (TKR 1 Double)+      f x =+        rvjp @1 @_ @(TKScalar Double) @g+               (rscanZip (\_ z -> z * z) doms (rscalar 5))+               doms3 x (ringestData [4] [1, 2, 3, 4])+  assertEqualUpToEpsilon 1e-10+    (ringestData [3] [4.0,6.0,8.0])+    (cgrad (kfromR . rsum0 . f) (rreplicate 3 (rscalar 1.1)))++testSin0revhV7 :: Assertion+testSin0revhV7 = do+  let doms = FTKR ZSR FTKScalar+      doms3 = FTKS (SNat @3 :$$ ZSS) FTKScalar+      f :: forall g. (BaseTensor g)+        => g (TKS '[3] Double) -> g (TKS '[3] Double)+      f x =+        svjp @_ @_ @(TKScalar Double) @g+               (\v -> sfromR @_ @'[4] $ rscanZip (\_ z -> z * z) doms (rscalar 5) (rfromS v))+                doms3 x (singestData @'[4] [1, 2, 3, 4])+  assertEqualUpToEpsilon 1e-10+    (singestData @'[3] [4.0,6.0,8.0])+    (cgrad (kfromS . ssum0 . f) (sreplicate @3 (sscalar 1.1)))
+ test/tool/CrossTesting.hs view
@@ -0,0 +1,509 @@+-- | Testing harness that differentiates a single objective function using+-- over a twenty different pipeline variants and cross-checks the results.+module CrossTesting+  ( rev', assertEqualUpToEpsilon'+  , t16, t16b, t48, t128, t128b, t128c+  , rrev1, rfwd1, srev1, sfwd1+  ) where++import Prelude++import GHC.Exts (IsList (..))+import GHC.TypeLits (KnownNat)+import System.IO.Unsafe (unsafePerformIO)+import Test.Tasty.HUnit hiding (assert)++import Data.Array.Nested.Ranked.Shape+import Data.Array.Nested.Shaped.Shape++import HordeAd.ADEngine+  ( IncomingCotangentHandling (..)+  , cjvp+  , jvp+  , revInterpretArtifact+  , revProduceArtifactWithoutInterpretation+  )+import HordeAd.AstEngine+import HordeAd.Core.Adaptor+import HordeAd.Core.Ast+import HordeAd.Core.AstEnv+import HordeAd.Core.AstFreshId+import HordeAd.Core.AstInterpret+import HordeAd.Core.AstTools+import HordeAd.Core.CarriersADVal+import HordeAd.Core.CarriersAst+import HordeAd.Core.CarriersConcrete+import HordeAd.Core.ConvertTensor (kfromR, kfromS)+import HordeAd.Core.Ops+import HordeAd.Core.OpsADVal+import HordeAd.Core.TensorKind+import HordeAd.Core.Types+import HordeAd.Core.Unwind+import HordeAd.OpsTensor++import EqEpsilon++crevMaybeBoth+  :: forall r m f src tgt.+     ( f ~ Concrete, X src ~ X (DValue src), KnownSTK (X src)+     , AdaptableTarget (ADVal Concrete) src+     , AdaptableTarget (ADVal Concrete) tgt+     , AdaptableTarget Concrete (DValue src)+     , tgt ~ ADVal f (TKR m r) )+  => (src -> tgt)+  -> DValue src+  -> (f (ADTensorKind (X src)), f (TKR m r))+{-# INLINE crevMaybeBoth #-}+crevMaybeBoth f vals =+  let g :: ADVal Concrete (X src) -> ADVal Concrete (TKR m r)+      g = toTarget . f . fromTarget+      valsH = toTarget vals+  in crevOnParams Nothing g (tftk knownSTK valsH) valsH++rev' :: forall r m n v a w.+        ( KnownNat m, KnownNat n, GoodScalar r+        , v ~ Concrete (TKR m r)+        , w ~ Concrete (ADTensorKind (TKR m r))+        , a ~ Concrete (ADTensorKind (TKR n r)) )+     => (forall f. ADReady f => f (TKR n r) -> f (TKR m r))+     -> Concrete (TKR n r)+     -> ( ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a+          , AstTensor AstMethodLet PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+          , v, v, v, v, v, v, v, v, v, v, v, v, v, v+          , a, a, a, a, a, a, a, a, a, a, a, a, a, a+          , Concrete (TKR n r), w, w, w+          , a, a, a, a, a )+        , ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a+          , AstTensor AstMethodLet PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+          , v, v, v, v, v, v, v, v, v, v, v, v, v, v+          , a, a, a, a, a, a, a, a, a, a, a, a, a, a+          , Concrete (TKR n r), w, w, w+          , a, a, a, a, a )+          )+{-# NOINLINE rev' #-}+rev' f vals = unsafePerformIO $ do+  setTotalSharing False+  !resNormalSharing <- rev1 f vals+  setTotalSharing True+  !resTotalSharing <- rev1 f vals+  setTotalSharing False+  return (resNormalSharing, resTotalSharing)++rev1 :: forall r m n v a w.+        ( KnownNat m, KnownNat n, GoodScalar r+        , v ~ Concrete (TKR m r)+        , w ~ Concrete (ADTensorKind (TKR m r))+        , a ~ Concrete (ADTensorKind (TKR n r)) )+     => (forall f. ADReady f => f (TKR n r) -> f (TKR m r))+     -> Concrete (TKR n r)+     -> IO ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a+           , AstTensor AstMethodLet PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+           , v, v, v, v, v, v, v, v, v, v, v, v, v, v+           , a, a, a, a, a, a, a, a, a, a, a, a, a, a+           , Concrete (TKR n r), w, w, w+           , a, a, a, a, a )+{-# NOINLINE rev1 #-}+rev1 f !vals = do+  let !value0 = f vals+      ftk = tftk knownSTK vals+      ftkz = adFTK ftk+      g :: ADVal Concrete (TKR n r)+        -> ADVal Concrete (TKR m r)+      g inputs = f $ fromTarget inputs+      !(!gradient1, !value1) = crevMaybeBoth g vals+      gradientRrev1 = rrev1 @Concrete @r @n @m f vals+      secondRrev1 = rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))+                            ftkz (rrev1 @_ @r @n @m f) vals+      g9 :: ADVal (AstRaw PrimalSpan) (TKR n r)+         -> ADVal (AstRaw PrimalSpan) (TKR m r)+      g9 inputs = f @(ADVal (AstRaw PrimalSpan))+                  $ fromTarget inputs+        -- fromTarget is fine, because primal of inputs is a variable,+        -- hence it's duplicable+      artifactsGradAst9 =+        fst $ revProduceArtifactWithoutInterpretation+                IgnoreIncomingCotangent g9 ftk+      !(!gradient9, !value9) = revInterpretArtifact7 artifactsGradAst9+      revInterpretArtifact7+        :: AstArtifactRev (TKR n r) (TKR m r)+        -> (Concrete (ADTensorKind (TKR n r)), Concrete (TKR m r))+      revInterpretArtifact7 a1 = revInterpretArtifact a1 vals Nothing+      hGeneral+        :: (ADReady fgen, ADReady f1)+        => (f1 (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))+        -> (AstTensor AstMethodLet FullSpan (TKR n r) -> f1 (TKR n r))+        -> (AstTensor AstMethodLet FullSpan (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))+        -> fgen (TKR n r)+        -> fgen (TKR m r)+      hGeneral fx1 fx2 gx inputs =+        let (var, ast) =+              funToAst (FTKR (rshape vals) FTKScalar) Nothing (fx1 . f . fx2)+            env = extendEnv var inputs emptyEnv+        in interpretAstFull env (gx ast)+      h :: ADReady f1+        => (f1 (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))+        -> (AstTensor AstMethodLet FullSpan (TKR n r) -> f1 (TKR n r))+        -> (AstTensor AstMethodLet FullSpan (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))+        -> ADVal Concrete (TKR n r)+        -> ADVal Concrete (TKR m r)+      h fx1 fx2 gx inputs =+        hGeneral @(ADVal Concrete) fx1 fx2 gx+                 (fromTarget inputs)+      !(!gradient2, !value2) =+        crevMaybeBoth (h id id id) vals+      !(!gradient3, !value3) =+        crevMaybeBoth (h id id simplifyInlineContract) vals+      !(!gradient2UnSimp, !value2UnSimp) =+        crevMaybeBoth (h unAstNoSimplify AstNoSimplify id) vals+      gradientRrev2UnSimp =+        rrev1 @Concrete @r @n @m @r+              (hGeneral unAstNoSimplify AstNoSimplify id) vals+      secondRrev2UnSimp =+        rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))+                ftkz (rrev1 @_ @r @n @m @r+                            (hGeneral unAstNoSimplify AstNoSimplify id)) vals+      !(!gradient3UnSimp, !value3UnSimp) =+        crevMaybeBoth (h (simplifyInline . unAstNoSimplify) AstNoSimplify simplifyInlineContract)+                      vals+      gradientRrev3UnSimp =+        rrev1 @Concrete @r @n @m @r+              (hGeneral unAstNoSimplify AstNoSimplify simplifyInlineContract) vals+      secondRrev3UnSimp =+        rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))+                ftkz (rrev1 @_ @r @n @m @r+                            (hGeneral unAstNoSimplify AstNoSimplify simplifyInlineContract)) vals+      !(!gradient4, !value4) =+        crevMaybeBoth (h unAstNoVectorize AstNoVectorize id)+                      vals+          -- use the AstNoVectorize instance that does no vectorization+          -- and then interpret the results as the Ast instance+      gradientRrev4 =+        rrev1 @Concrete @r @n @m @r+              (hGeneral unAstNoVectorize AstNoVectorize id) vals+      secondRrev4 =+        rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))+                ftkz (rrev1 @_ @r @n @m @r+                            (hGeneral unAstNoVectorize AstNoVectorize id)) vals+      !(!gradient5, !value5) =+        crevMaybeBoth (h unAstNoVectorize AstNoVectorize simplifyInlineContract)+                      vals+      gradientRrev5 =+        rrev1 @Concrete @r @n @m @r+              (hGeneral unAstNoVectorize AstNoVectorize simplifyInlineContract) vals+      secondRrev5 =+        rrevFTK @Concrete @(TKR n r) @(ADTensorKind (TKR n r))+                ftkz (rrev1 @_ @r @n @m @r+                            (hGeneral unAstNoVectorize AstNoVectorize simplifyInlineContract)) vals+      astVectSimp = simplifyInlineContract $ snd $ funToAst (FTKR (rshape vals) FTKScalar) Nothing f+      astSimp =+        simplifyInlineContract $ simplifyInlineContract $ snd  -- builds simplify with difficulty+        $ funToAst (FTKR (rshape vals) FTKScalar) Nothing (unAstNoVectorize . f . AstNoVectorize)+      -- Here comes the part with Ast gradients.+      hAst :: ADReady f1+           => (f1 (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))+           -> (AstTensor AstMethodLet FullSpan (TKR n r) -> f1 (TKR n r))+           -> (AstTensor AstMethodLet FullSpan (TKR m r) -> AstTensor AstMethodLet FullSpan (TKR m r))+           -> ADVal (AstRaw PrimalSpan) (TKR n r)+           -> ADVal (AstRaw PrimalSpan) (TKR m r)+      hAst fx1 fx2 gx inputs+        = hGeneral @(ADVal (AstRaw PrimalSpan))+                   fx1 fx2 gx (fromTarget inputs)+        -- fromTarget is fine, because primal of inputs is a variable,+        -- hence it's duplicable+      artifactsGradAst =+        fst $ revProduceArtifactWithoutInterpretation+                IgnoreIncomingCotangent (hAst id id id) ftk+      !(!gradient2Ast, !value2Ast) =+        revInterpretArtifact7 artifactsGradAst+      !(!gradient2AstS, !value2AstS) =+        revInterpretArtifact7 (simplifyArtifact artifactsGradAst)+      artifactsGradAstT =+        fst $ revProduceArtifactWithoutInterpretation+                UseIncomingCotangent (hAst simplifyInline id id) ftk+      !(!gradient2AstST, !value2AstST) =+        revInterpretArtifact7 (simplifyArtifact artifactsGradAstT)+      artifactsSimpleAst =+        fst $ revProduceArtifactWithoutInterpretation+                IgnoreIncomingCotangent (hAst id id simplifyInlineContract) ftk+      !(!gradient3Ast, !value3Ast) =+        revInterpretArtifact7 artifactsSimpleAst+      !(!gradient3AstS, !value3AstS) =+        revInterpretArtifact7 (simplifyArtifact artifactsSimpleAst)+      artifactsGradAstUnSimp =+        fst $ revProduceArtifactWithoutInterpretation+                IgnoreIncomingCotangent (hAst unAstNoSimplify AstNoSimplify id) ftk+      !(!gradient2AstUnSimp, !value2AstUnSimp) =+        revInterpretArtifact7 artifactsGradAstUnSimp+      !(!gradient2AstSUnSimp, !value2AstSUnSimp) =+        revInterpretArtifact7 (simplifyArtifact artifactsGradAstUnSimp)+      artifactsSimpleAstUnSimp =+        fst $ revProduceArtifactWithoutInterpretation+                IgnoreIncomingCotangent (hAst unAstNoSimplify AstNoSimplify simplifyInlineContract)+                ftk+      !(!gradient3AstUnSimp, !value3AstUnSimp) =+        revInterpretArtifact7 artifactsSimpleAstUnSimp+      !(!gradient3AstSUnSimp, !value3AstSUnSimp) =+        revInterpretArtifact7 (simplifyArtifact artifactsSimpleAstUnSimp)+      artifactsPrimalAst =+        fst $ revProduceArtifactWithoutInterpretation+                IgnoreIncomingCotangent (hAst unAstNoVectorize AstNoVectorize id) ftk+      !(!gradient4Ast, !value4Ast) =+        revInterpretArtifact7 artifactsPrimalAst+      !(!gradient4AstS, !value4AstS) =+        revInterpretArtifact7 (simplifyArtifact artifactsPrimalAst)+      artifactsPSimpleAst =+        fst $ revProduceArtifactWithoutInterpretation+                IgnoreIncomingCotangent (hAst unAstNoVectorize AstNoVectorize simplifyInlineContract)+                ftk+      !(!gradient5Ast, !value5Ast) =+        revInterpretArtifact7 artifactsPSimpleAst+      -- Due to no vectorization this may result in huge terms,+      -- which then take forever to inline into (substitution into huge term)+      -- and to simplify, so we ignore this test for now.+      (gradient5AstS, value5AstS) =+        revInterpretArtifact7 (simplifyArtifact artifactsPSimpleAst)+      !cderivative = cjvp f vals vals+      !derivative = jvp f vals vals+      !derivativeRfwd1 = rfwd1ds @Concrete @r @n @m @r f vals+                         $ toADTensorKindShared ftk vals+  return+    ( value0, value1, value2, value3, value2UnSimp, value3UnSimp+    , value4, value5+    , gradient1, gradientRrev1, gradient2, gradient3+    , gradient2UnSimp, gradientRrev2UnSimp+    , gradient3UnSimp, gradientRrev3UnSimp+    , gradient4, gradientRrev4, gradient5, gradientRrev5+    , astVectSimp, astSimp+    , value9, value2Ast, value2AstS, value2AstST, value3Ast, value3AstS+    , value2AstUnSimp, value2AstSUnSimp, value3AstUnSimp, value3AstSUnSimp+    , value4Ast, value4AstS, value5Ast, value5AstS+    , gradient9, gradient2Ast, gradient2AstS, gradient2AstST+    , gradient3Ast, gradient3AstS+    , gradient2AstUnSimp, gradient2AstSUnSimp+    , gradient3AstUnSimp, gradient3AstSUnSimp+    , gradient4Ast, gradient4AstS, gradient5Ast, gradient5AstS+    , vals, cderivative, derivative, derivativeRfwd1+    , secondRrev1, secondRrev2UnSimp, secondRrev3UnSimp+    , secondRrev4, secondRrev5+    )++assertEqualUpToEpsilon'+    :: ( KnownNat n, KnownNat m+       , v ~ Concrete (TKR m r)+       , w ~ Concrete (ADTensorKind (TKR m r))+       , a ~ Concrete (ADTensorKind (TKR n r))+       , AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v+       , AssertEqualUpToEpsilon (ADTensorScalar r)+       , GoodScalar r, GoodScalar (ADTensorScalar r), HasCallStack)+    => Rational  -- ^ error margin (i.e., the epsilon)+    -> Concrete (TKR n r)  -- ^ expected reverse derivative value+    -> ( ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a+         , AstTensor AstMethodLet PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+         , v, v, v, v, v, v, v, v, v, v, v, v, v, v+         , a, a, a, a, a, a, a, a, a, a, a, a, a, a+         , Concrete (TKR n r), w, w, w+         , a, a, a, a, a )+       , ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a+         , AstTensor AstMethodLet PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+         , v, v, v, v, v, v, v, v, v, v, v, v, v, v+         , a, a, a, a, a, a, a, a, a, a, a, a, a, a+         , Concrete (TKR n r), w, w, w+         , a, a, a, a, a ) )+    -> Assertion+assertEqualUpToEpsilon' errMargin expected' (tup1, tup2) = do+  assertEqualUpToEpsilon1 errMargin expected' tup1+  assertEqualUpToEpsilon1 errMargin expected' tup2++assertEqualUpToEpsilon1+    :: ( KnownNat n, KnownNat m+       , v ~ Concrete (TKR m r)+       , w ~ Concrete (ADTensorKind (TKR m r))+       , a ~ Concrete (ADTensorKind (TKR n r))+       , AssertEqualUpToEpsilon a, AssertEqualUpToEpsilon v+       , AssertEqualUpToEpsilon (ADTensorScalar r)+       , GoodScalar r, GoodScalar (ADTensorScalar r), HasCallStack)+    => Rational  -- ^ error margin (i.e., the epsilon)+    -> Concrete (TKR n r)  -- ^ expected reverse derivative value+    -> ( v, v, v, v, v, v, v, v, a, a, a, a, a, a, a, a, a, a, a, a+       , AstTensor AstMethodLet PrimalSpan (TKR m r), AstTensor AstMethodLet PrimalSpan (TKR m r)+       , v, v, v, v, v, v, v, v, v, v, v, v, v, v+       , a, a, a, a, a, a, a, a, a, a, a, a, a, a+       , Concrete (TKR n r), w, w, w+       , a, a, a, a, a )+         -- ^ actual values+    -> Assertion+assertEqualUpToEpsilon1+    errMargin expected'+    ( value0, value1, value2, value3, value2UnSimp, value3UnSimp+    , value4, value5+    , gradient1, gradientRrev1, gradient2, gradient3+    , gradient2UnSimp, gradientRrev2UnSimp+    , gradient3UnSimp, gradientRrev3UnSimp+    , gradient4, gradientRrev4, gradient5, gradientRrev5+    , _astVectSimp, _astSimp+    , value9, value2Ast, value2AstS, value2AstST, value3Ast, value3AstS+    , value2AstUnSimp, value2AstSUnSimp, value3AstUnSimp, value3AstSUnSimp+    , value4Ast, value4AstS, value5Ast, _value5AstS+    , gradient9, gradient2Ast, gradient2AstS, gradient2AstST+    , gradient3Ast, gradient3AstS+    , gradient2AstUnSimp, gradient2AstSUnSimp+    , gradient3AstUnSimp, gradient3AstSUnSimp+    , gradient4Ast, gradient4AstS, gradient5Ast, _gradient5AstS+    , vals, cderivative, derivative, derivativeRfwd1+    , secondRrev1, secondRrev2UnSimp, secondRrev3UnSimp+    , secondRrev4, secondRrev5 ) = do+  let ftk = tftk knownSTK vals+      expected = toADTensorKindShared ftk expected'+  assertEqualUpToEpsilonWithMark "Val ADVal" errMargin value0 value1+  assertEqualUpToEpsilonWithMark "Val Vectorized" errMargin value0 value2+  assertEqualUpToEpsilonWithMark "Val Vect+Simp" errMargin value0 value3+  assertEqualUpToEpsilonWithMark "Val V UnS" errMargin value0 value2UnSimp+  assertEqualUpToEpsilonWithMark "Val V+S UnS" errMargin value0 value3UnSimp+  assertEqualUpToEpsilonWithMark "Val NotVect" errMargin value0 value4+  assertEqualUpToEpsilonWithMark "Val Simplified" errMargin value0 value5+  assertEqualUpToEpsilonWithMark "Grad ADVal" errMargin expected gradient1+  assertEqualUpToEpsilonWithMark "Grad ADVal rrev"+                                 errMargin expected gradientRrev1+  assertEqualUpToEpsilonWithMark "Grad Vectorized" errMargin expected gradient2+  assertEqualUpToEpsilonWithMark "Grad Vect+Simp" errMargin expected gradient3+  assertEqualUpToEpsilonWithMark "Grad V UnS" errMargin expected gradient2UnSimp+  assertEqualUpToEpsilonWithMark "Grad V UnS rrev2"+                                 errMargin expected gradientRrev2UnSimp+  assertEqualUpToEpsilonWithMark "Second V UnS rrev2"+                                 (max 1e-5 errMargin) secondRrev1 secondRrev2UnSimp+  assertEqualUpToEpsilonWithMark "Grad V+S UnS"+                                 errMargin expected gradient3UnSimp+  assertEqualUpToEpsilonWithMark "Grad V+S UnS rrev"+                                 errMargin expected gradientRrev3UnSimp+  assertEqualUpToEpsilonWithMark "Second V+S UnS rrev"+                                 (max 1e-5 errMargin) secondRrev1 secondRrev3UnSimp+  assertEqualUpToEpsilonWithMark "Grad NotVect" errMargin expected gradient4+  assertEqualUpToEpsilonWithMark "Grad NotVect rrev"+                                 errMargin expected gradientRrev4+  assertEqualUpToEpsilonWithMark "Second NotVect rrev"+                                 (max 1e-5 errMargin) secondRrev1 secondRrev4+  assertEqualUpToEpsilonWithMark "Grad Simplified" errMargin expected gradient5+  assertEqualUpToEpsilonWithMark "Grad Simplified rrev2"+                                 errMargin expected gradientRrev5+  assertEqualUpToEpsilonWithMark "Second Simplified rrev2"+                                 (max 1e-5 errMargin) secondRrev1 secondRrev5+  assertEqualUpToEpsilonWithMark "Val Ast Vectorized" errMargin value0 value2Ast+  assertEqualUpToEpsilonWithMark "Val Ast V S" errMargin value0 value2AstS+  assertEqualUpToEpsilonWithMark "Val Ast V ST" errMargin value0 value2AstST+  assertEqualUpToEpsilonWithMark "Val Ast Vect+Simp" errMargin value0 value3Ast+  assertEqualUpToEpsilonWithMark "Val Ast V+S S" errMargin value0 value3AstS+  assertEqualUpToEpsilonWithMark "Val Ast V UnS" errMargin value0+                                                 value2AstUnSimp+  assertEqualUpToEpsilonWithMark "Val Ast V S UnS" errMargin value0+                                                   value2AstSUnSimp+  assertEqualUpToEpsilonWithMark "Val Ast Vect+Simp UnS" errMargin value0+                                                         value3AstUnSimp+  assertEqualUpToEpsilonWithMark "Val Ast V+S S UnS" errMargin value0+                                                     value3AstSUnSimp+  assertEqualUpToEpsilonWithMark "Val Ast NotVect" errMargin value0 value4Ast+  assertEqualUpToEpsilonWithMark "Val Ast NotVect S" errMargin value0 value4AstS+  assertEqualUpToEpsilonWithMark "Val Ast Simplified" errMargin value0 value5Ast+--  assertEqualUpToEpsilonWithMark "Val Ast S S" errMargin value0 value5AstS+  assertEqualUpToEpsilonWithMark "Grad Ast Vectorized"+                                 errMargin expected gradient2Ast+  assertEqualUpToEpsilonWithMark "Grad Ast Vectorized S"+                                 errMargin expected gradient2AstS+  assertEqualUpToEpsilonWithMark "Grad Ast Vectorized ST"+                                 errMargin expected gradient2AstST+  assertEqualUpToEpsilonWithMark "Grad Ast Vect+Simp"+                                 errMargin expected gradient3Ast+  assertEqualUpToEpsilonWithMark "Grad Ast Vect+Simp S"+                                 errMargin expected gradient3AstS+  assertEqualUpToEpsilonWithMark "Grad Ast Vectorized UnS"+                                 errMargin expected gradient2AstUnSimp+  assertEqualUpToEpsilonWithMark "Grad Ast Vectorized S UnS"+                                 errMargin expected gradient2AstSUnSimp+  assertEqualUpToEpsilonWithMark "Grad Ast Vect+Simp UnS"+                                 errMargin expected gradient3AstUnSimp+  assertEqualUpToEpsilonWithMark "Grad Ast Vect+Simp S UnS"+                                 errMargin expected gradient3AstSUnSimp+  assertEqualUpToEpsilonWithMark "Grad Ast NotVect"+                                 errMargin expected gradient4Ast+  assertEqualUpToEpsilonWithMark "Grad Ast NotVect S"+                                 errMargin expected gradient4AstS+  assertEqualUpToEpsilonWithMark "Grad Ast Simplified"+                                 errMargin expected gradient5Ast+--  assertEqualUpToEpsilonWithMark "Grad Ast Simplified S"+--                                 errMargin expected gradient5AstS+  assertEqualUpToEpsilonWithMark "Val ADVal Ast" errMargin value0 value9+  assertEqualUpToEpsilonWithMark "Grad ADVal Ast" errMargin expected gradient9+  assertEqualUpToEpsilonWithMark "Derivatives" errMargin cderivative derivative+  assertEqualUpToEpsilonWithMark "Derivatives rjvp"+                                 errMargin cderivative derivativeRfwd1+  -- The formula for comparing derivative and gradient is due to @awf+  -- at https://github.com/Mikolaj/horde-ad/issues/15#issuecomment-1063251319+  -- and a similar property stated mathematically is in Lemma 1 in+  -- https://www.microsoft.com/en-us/research/uploads/prod/2021/08/higher-order-ad.pdf+  assertEqualUpToEpsilonWithMark "Reverse vs forward"+                                 1e-5 (rsum0 derivative) (rdot0 expected (toADTensorKindShared ftk vals))+  {- TODO: this most probably leaks gigabytes of strings from one test case+  -- to another in -O0 mode, leading to OOMs, so it's disabled for now.+  -- We could also try to stream the strings and compare on the fly.+  --+  -- No Eq instance, so let's compare the text.+  assertEqual "Idempotence of simplification of non-vectorized AST"+              (show astSimp)+              (show (simplifyInlineContract astSimp))+  assertEqual "Idempotence of simplification of vectorized AST"+              (show astVectSimp)+              (show (simplifyInlineContract astVectSimp))+  -}++t16 :: (GoodScalar r, Fractional r) => Concrete (TKR 5 r)+t16 = ringestData (fromList [2, 2, 1, 2, 2]) [5, 2, 6, 1, -2, 0.000001, 0.1, -0.2, 13.1, 9, 8, -4, 34, 2.99432, -33, 26]++t16b :: (GoodScalar r, Fractional r) => Concrete (TKR 4 r)+t16b = ringestData (fromList [2, 2, 2, 2]) [5, 2, 6, 1, -2, 0, 0.1, -0.2, 13.1, 9, 8, -4, 582934, 2.99432, -335, 26]++t48 :: (GoodScalar r, Fractional r) => Concrete (TKR 7 r)+t48 = ringestData (fromList [3, 1, 2, 2, 1, 2, 2]) [18.1,29.1,32.1,40.1,52.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,58.0,54.99432,97.1,52.8943200001, 5, 2, 6, 1, -2, 0.92, 0.1, -0.2, 13.1, 9, 8, -4, 34, 2.99432, -33, 26, 2, 2, 2, 2, -0.2,-0.2,-0.2,-0.2,25.0003,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003]++t128 :: (GoodScalar r, Fractional r) => Concrete (TKR 10 r)+t128 = ringestData (fromList [1, 2, 2, 1, 2, 2, 2, 2, 2, 1]) [29.1,32.1,40.1,29.0,53.99432,97.1,58.8943200001,18.1,29.1,32.1,40.1,32.0,53.99432,97.1,25.8943200001, 5, 2, 6, 1, -2, 97.1,58.8943200001,97.1,55.8943200001,97.1,58.8943200001,18.1,29.1,32.1,40.1,32.1,32.1,40.1,53.0,53.99432, -0.00001, 0.1, -0.2, 13.1, 9, 8, -4, 29, 2.99432, -335, 26, 2, 2, 2, 2, -0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,-0.2,-0.2,-0.2,-0.2,25.0003,25.0003,25.0003,25.0003,40.1,8.0,11.0,-3.0,25.89432,28.79432,-39.09999999999997,25.8,40.1,8.0,11.0,-3.0,25.89432,28.79432,-19.09999999999997,25.8, 8.1,29.1,32.1,40.1,32.1,40.1,292.0,53.99432,97.1,55.8943200001,97.1,85.8943200001,97.1,85.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,22.0,53.99432,97.1,82.8943200001,97.1,22.8943200001,97.1,58.8943200001,18.1,29.1,32.1,40.1,32.1,40.1,32.1,40.1,89.0,53.99432,97.1,56.8943200001,97.1,52.8943200001,97.1,55.8943200001]++t128b :: (GoodScalar r, Fractional r) => Concrete (TKR 4 r)+t128b = rreshape (4 :$: 2 :$: 4 :$: 4 :$: ZSR) t128++t128c :: (GoodScalar r, Fractional r) => Concrete (TKR 4 r)+t128c = rreshape (2 :$: 2 :$: 8 :$: 4 :$: ZSR) t128++rrev1 :: forall g r n m r3.+         (ADReady g, GoodScalar r, KnownNat n, GoodScalar r3, KnownNat m)+      => (forall f. ADReady f => f (TKR n r) -> f (TKR m r3)) -> g (TKR n r)+      -> g (ADTensorKind (TKR n r))+rrev1 f u = kgrad (kfromR . rsum0 . f) (tftk knownSTK u) u++rrevFTK :: forall g x z. (ADReady g, KnownSTK x)+        => FullShapeTK z -> (forall f. ADReady f => f x -> f z) -> g x+        -> g (ADTensorKind x)+rrevFTK ftk f u = kgrad (tsum0Target ftk . f) (tftk knownSTK u) u++rfwd1ds :: forall g r n m r3.+           (ADReady g, GoodScalar r, KnownNat n)+        => (forall f. ADReady f => f (TKR n r) -> f (TKR m r3)) -> g (TKR n r)+        -> g (ADTensorKind (TKR n r))+        -> g (ADTensorKind (TKR m r3))+rfwd1ds f u = rjvp f (tftk knownSTK u) u++rfwd1 :: forall g r n m r3.+         (ADReady g, GoodScalar r, GoodScalar (ADTensorScalar r), KnownNat n)+      => (forall f. ADReady f => f (TKR n r) -> f (TKR m r3)) -> g (TKR n r)+      -> g (ADTensorKind (TKR m r3))+rfwd1 f u = rfwd1ds f u (rrepl (rshape u) 1)++srev1 :: forall g r sh sh2 r3.+         (ADReady g, GoodScalar r, KnownShS sh, GoodScalar r3, KnownShS sh2)+      => (forall f. ADReady f => f (TKS sh r) -> f (TKS sh2 r3)) -> g (TKS sh r)+      -> g (ADTensorKind (TKS sh r))+srev1 f u = kgrad (kfromS. ssum0 . f) (tftk knownSTK u) u++sfwd1 :: forall g r sh sh2 r3.+         (ADReady g, GoodScalar r, GoodScalar (ADTensorScalar r), KnownShS sh)+      => (forall f. ADReady f => f (TKS sh r) -> f (TKS sh2 r3)) -> g (TKS sh r)+      -> g (ADTensorKind (TKS sh2 r3))+sfwd1 f u = sjvp f (tftk knownSTK u) u (srepl @_ @(ADTensorScalar r) 1)
+ test/tool/EqEpsilon.hs view
@@ -0,0 +1,352 @@+{-# LANGUAGE UndecidableInstances #-}+-- | Operations for comparing values up to a tolerance margin, to be used+-- for tests.+module EqEpsilon+  ( EqEpsilon, setEpsilonEq+  , AssertEqualUpToEpsilon(..)+  , assertEqualUpToEpsilonWithMark, assertEqualUpToEpsilon+  , assertCloseElem, assertClose, (@?~)+  ) where++import Prelude++import Data.Int (Int64)+import Data.IORef+import Foreign.C (CInt)+import System.IO.Unsafe (unsafePerformIO)+import Test.Tasty.HUnit+import Test.Tasty.Options++import Shared++newtype EqEpsilon = EqEpsilon Rational++instance IsOption EqEpsilon where+  defaultValue = EqEpsilon eqEpsilonDefault+  parseValue s = fmap (EqEpsilon . toRational)+                      ((safeRead :: String -> Maybe Double) s)+  optionName = return "eq-epsilon"+  optionHelp = return $ "Epsilon to use for floating point comparisons: abs(a-b) <= epsilon. Default: " ++ show (fromRational eqEpsilonDefault :: Double)++-- | Default value for eqEpsilonRef.+eqEpsilonDefault :: Rational+eqEpsilonDefault = 1e-6++-- | Global epsilon used to compare floating point values.+eqEpsilonRef :: IORef Rational+{-# NOINLINE eqEpsilonRef #-}+eqEpsilonRef = unsafePerformIO $ newIORef eqEpsilonDefault++-- | Global epsilon setter (to be called once).+setEpsilonEq :: EqEpsilon -> IO ()+setEpsilonEq (EqEpsilon x) = atomicWriteIORef eqEpsilonRef x+++-- * Helper functions++assert_list :: forall a. HasCallStack+            => (a -> a -> Assertion)+                 -- ^ The function used to make an assertion+                 -- on two elements (expected, actual)+            -> [a]  -- ^ The expected value+            -> [a]  -- ^ The actual value+            -> Assertion+assert_list make_assert expected actual =+  if lenE == lenA then+    go_assert_list expected actual+  else+    assertFailure $ "List too " ++ (if lenE < lenA then "long" else "short")+                    ++ ": expected " ++ show lenE ++ " elements, but got: "+                    ++ show lenA+  where+    lenE :: Int = length expected+    lenA :: Int = length actual++    go_assert_list :: [a]  -- The expected value+                   -> [a]  -- The actual value+                   -> Assertion+    go_assert_list [] [] = assertBool "" True+    go_assert_list [] (_:_) = assertFailure "More list elements than expected!"+    go_assert_list (_:_) [] = assertFailure "Less list elements than expected!"+    go_assert_list (head_exp:tail_exp) (head_act:tail_act) =+      make_assert head_exp head_act >> go_assert_list tail_exp tail_act++assert_shape+  :: forall a b. (HasShape a, Linearizable a b, HasCallStack)+  => (b -> b -> Assertion)  -- ^ The function used to make an assertion+                            -- on two elements (expected, actual)+  -> a                      -- ^ The expected value+  -> a                      -- ^ The actual value+  -> Assertion+assert_shape make_assert expected actual =+  if shapeE == shapeA then+    assert_list make_assert (linearize expected) (linearize actual)+  else+    assertFailure $ "Expected shape: " ++ show shapeE ++ ", but got: "+                    ++ show shapeA+  where+    shapeE = shapeL expected+    shapeA = shapeL actual+++-- * Generic comparisons with explicit error margin++-- | Asserts that the specified actual floating point value is close+-- to the expected value. The output message will contain the prefix,+-- the expected value, and the actual value.+--+-- If the prefix is the empty string (i.e., @\"\"@), then the prefix is omitted+-- and only the expected and actual values are output.+assert_close_eps :: (Real a, Show a, HasCallStack)+                 => String    -- ^ The message prefix+                 -> String    -- ^ The message suffix+                 -> Rational  -- ^ The error margin+                 -> a         -- ^ The expected value+                 -> a         -- ^ The actual value+                 -> Assertion+assert_close_eps preface epilogue eqEpsilon expected actual = do+  assertBool (message eqEpsilon)+             (realToFrac (abs (expected - actual)) <= eqEpsilon)+  where+    msg = "expected: " ++ show expected ++ "\n but got: " ++ show actual+    message errorMargin =+      (if null preface then "" else preface ++ "\n")+      ++ msg ++ "\n (maximum margin of error: "+      ++ show (realToFrac errorMargin :: Double) ++ ")"+      ++ (if null epilogue+             || (lowercase epilogue == lowercase preface)+             || (lowercase epilogue == lowercase msg)+          then ""+          else "\n" ++ epilogue)+++-- * AssertEqualUpToEpsilon class++class Show a => AssertEqualUpToEpsilon a where+  assertEqualUpToEpsilonWithMsg+    :: String    -- ^ message suffix+    -> Rational  -- ^ error margin (i.e., the epsilon)+    -> a         -- ^ expected value+    -> a         -- ^ actual value+    -> Assertion++assertEqualUpToEpsilonWithMark+  :: (AssertEqualUpToEpsilon a, HasCallStack)+  => String  -- ^ message suffix's prefix+  -> Rational  -- ^ error margin (i.e., the epsilon)+  -> a  -- ^ expected value+  -> a  -- ^ actual value+  -> Assertion+assertEqualUpToEpsilonWithMark mark error_margin expected actual =+  let prefix = if null mark then "" else "*In " ++ mark ++ "*\n"+  in assertEqualUpToEpsilonWithMsg+       (prefix ++ "Expected: " ++ show expected+               ++ "\n but got: " ++ show actual)+       error_margin+       expected actual++assertEqualUpToEpsilon+  :: (AssertEqualUpToEpsilon a, HasCallStack)+  => Rational  -- ^ error margin (i.e., the epsilon)+  -> a  -- ^ expected value+  -> a  -- ^ actual value+  -> Assertion+assertEqualUpToEpsilon = assertEqualUpToEpsilonWithMark ""++instance AssertEqualUpToEpsilon Double where+  assertEqualUpToEpsilonWithMsg = assert_close_eps ""++instance AssertEqualUpToEpsilon Float where+  assertEqualUpToEpsilonWithMsg = assert_close_eps ""++instance AssertEqualUpToEpsilon Int64 where+  assertEqualUpToEpsilonWithMsg = assert_close_eps ""++instance AssertEqualUpToEpsilon CInt where+  assertEqualUpToEpsilonWithMsg = assert_close_eps ""++instance (AssertEqualUpToEpsilon a,+          AssertEqualUpToEpsilon b)+         => AssertEqualUpToEpsilon (a,b) where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon (e1,e2) (a1,a2) =+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e1 a1 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e2 a2++instance (AssertEqualUpToEpsilon a,+          AssertEqualUpToEpsilon b,+          AssertEqualUpToEpsilon c)+         => AssertEqualUpToEpsilon (a,b,c) where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon (e1,e2,e3) (a1,a2,a3) =+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e1 a1 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e2 a2 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e3 a3++instance (AssertEqualUpToEpsilon a,+          AssertEqualUpToEpsilon b,+          AssertEqualUpToEpsilon c,+          AssertEqualUpToEpsilon d)+         => AssertEqualUpToEpsilon (a,b,c,d) where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon (e1,e2,e3,e4) (a1,a2,a3,a4) =+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e1 a1 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e2 a2 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e3 a3 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e4 a4++instance (AssertEqualUpToEpsilon a,+          AssertEqualUpToEpsilon b,+          AssertEqualUpToEpsilon c,+          AssertEqualUpToEpsilon d,+          AssertEqualUpToEpsilon e)+         => AssertEqualUpToEpsilon (a,b,c,d,e) where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon (e1,e2,e3,e4,e5)+                                              (a1,a2,a3,a4,a5) =+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e1 a1 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e2 a2 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e3 a3 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e4 a4 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e5 a5++instance (AssertEqualUpToEpsilon a,+          AssertEqualUpToEpsilon b,+          AssertEqualUpToEpsilon c,+          AssertEqualUpToEpsilon d,+          AssertEqualUpToEpsilon e,+          AssertEqualUpToEpsilon f)+         => AssertEqualUpToEpsilon (a,b,c,d,e,f) where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon (e1,e2,e3,e4,e5,e6)+                                              (a1,a2,a3,a4,a5,a6) =+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e1 a1 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e2 a2 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e3 a3 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e4 a4 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e5 a5 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e6 a6++instance (AssertEqualUpToEpsilon a,+          AssertEqualUpToEpsilon b,+          AssertEqualUpToEpsilon c,+          AssertEqualUpToEpsilon d,+          AssertEqualUpToEpsilon e,+          AssertEqualUpToEpsilon f,+          AssertEqualUpToEpsilon g)+         => AssertEqualUpToEpsilon (a,b,c,d,e,f,g) where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon (e1,e2,e3,e4,e5,e6,e7)+                                              (a1,a2,a3,a4,a5,a6,a7) =+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e1 a1 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e2 a2 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e3 a3 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e4 a4 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e5 a5 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e6 a6 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e7 a7++instance (AssertEqualUpToEpsilon a,+          AssertEqualUpToEpsilon b,+          AssertEqualUpToEpsilon c,+          AssertEqualUpToEpsilon d,+          AssertEqualUpToEpsilon e,+          AssertEqualUpToEpsilon f,+          AssertEqualUpToEpsilon g,+          AssertEqualUpToEpsilon h)+         => AssertEqualUpToEpsilon (a,b,c,d,e,f,g,h) where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon (e1,e2,e3,e4,e5,e6,e7,e8)+                                              (a1,a2,a3,a4,a5,a6,a7,a8) =+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e1 a1 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e2 a2 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e3 a3 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e4 a4 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e5 a5 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e6 a6 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e7 a7 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e8 a8++instance (AssertEqualUpToEpsilon a,+          AssertEqualUpToEpsilon b,+          AssertEqualUpToEpsilon c,+          AssertEqualUpToEpsilon d,+          AssertEqualUpToEpsilon e,+          AssertEqualUpToEpsilon f,+          AssertEqualUpToEpsilon g,+          AssertEqualUpToEpsilon h,+          AssertEqualUpToEpsilon i)+         => AssertEqualUpToEpsilon (a,b,c,d,e,f,g,h,i) where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon (e1,e2,e3,e4,e5,e6,e7,e8,e9)+                                              (a1,a2,a3,a4,a5,a6,a7,a8,a9) =+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e1 a1 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e2 a2 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e3 a3 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e4 a4 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e5 a5 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e6 a6 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e7 a7 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e8 a8 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e9 a9++instance (AssertEqualUpToEpsilon a,+          AssertEqualUpToEpsilon b,+          AssertEqualUpToEpsilon c,+          AssertEqualUpToEpsilon d,+          AssertEqualUpToEpsilon e,+          AssertEqualUpToEpsilon f,+          AssertEqualUpToEpsilon g,+          AssertEqualUpToEpsilon h,+          AssertEqualUpToEpsilon i,+          AssertEqualUpToEpsilon j)+         => AssertEqualUpToEpsilon (a,b,c,d,e,f,g,h,i,j) where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon (e1,e2,e3,e4,e5,e6,e7,e8,e9,e10)+                                              (a1,a2,a3,a4,a5,a6,a7,a8,a9,a10) =+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e1  a1 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e2  a2 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e3  a3 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e4  a4 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e5  a5 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e6  a6 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e7  a7 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e8  a8 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e9  a9 >>+    assertEqualUpToEpsilonWithMsg msg eqEpsilon e10 a10++instance {-# OVERLAPPABLE #-}+         (Show a, HasShape a, Linearizable a b, AssertEqualUpToEpsilon b)+         => AssertEqualUpToEpsilon a where+  assertEqualUpToEpsilonWithMsg msg eqEpsilon =+    assert_shape (assertEqualUpToEpsilonWithMsg msg eqEpsilon)+++-- * Generic comparisons without explicit error margin++-- | Asserts that the specified actual floating point value is close+-- to at least one of the expected values.+assertCloseElem :: forall a. (Real a, Fractional a, Show a, HasCallStack)+                => String   -- ^ The message prefix+                -> [a]      -- ^ The expected values+                -> a        -- ^ The actual value+                -> Assertion+assertCloseElem preface expected actual = do+  eqEpsilon <- readIORef eqEpsilonRef+  go_assert eqEpsilon expected+  where+    msg = (if null preface then "" else preface ++ "\n")+          ++ "wrong result: " ++ show actual+          ++ " is expected to be a member of " ++ show expected+    go_assert :: Rational -> [a] -> Assertion+    go_assert _ [] = assertFailure msg+    go_assert eqEps (h:t) =+      if abs (h-actual) <= fromRational eqEps+      then assert_close_eps msg "" (fromRational eqEps) h actual+      else go_assert eqEps t++assertClose :: (AssertEqualUpToEpsilon a, HasCallStack)+      => a  -- ^ The expected value+      -> a  -- ^ The actual value+      -> Assertion+assertClose expected actual = do+  eqEpsilon <- readIORef eqEpsilonRef+  assertEqualUpToEpsilon (fromRational eqEpsilon) expected actual++infix 1 @?~+(@?~) :: (AssertEqualUpToEpsilon a, HasCallStack)+      => a  -- ^ The actual value+      -> a  -- ^ The expected value+      -> Assertion+(@?~) = flip assertClose
+ test/tool/Shared.hs view
@@ -0,0 +1,89 @@+{-# LANGUAGE UndecidableInstances #-}+-- | Additional classes that help in comparing values in tests.+module Shared+  ( lowercase, HasShape (shapeL), Linearizable (linearize)+  ) where++import Prelude++import Data.Char qualified+import Data.Foldable qualified+import Data.Int (Int64)+import Data.Vector.Storable qualified as VS+import Foreign.C (CInt)+import GHC.Exts (IsList (..))+import GHC.TypeLits (KnownNat)++import Data.Array.Nested qualified as Nested+import Data.Array.Nested.Shaped.Shape++import HordeAd.Core.CarriersConcrete+import HordeAd.Core.Types++lowercase :: String -> String+lowercase = map Data.Char.toLower++-- | Things that have shape.+class HasShape a where+  shapeL :: a -> [Int]++instance (KnownNat n, Nested.PrimElt a) => HasShape (Nested.Ranked n a) where+  shapeL = toList . Nested.rshape++instance KnownShS sh => HasShape (Nested.Shaped sh a) where+  shapeL _ = toList $ knownShS @sh++instance HasShape (RepConcrete y) => HasShape (Concrete y) where+  shapeL = shapeL . unConcrete++instance HasShape Double where+  shapeL _ = []++instance HasShape Float where+  shapeL _ = []++instance HasShape Int64 where+  shapeL _ = []++instance HasShape CInt where+  shapeL _ = []++instance HasShape Z1 where+  shapeL _ = [0]++instance {-# OVERLAPPABLE #-} (Foldable t) => HasShape (t a) where+  shapeL = (: []) . length++-- | Things that can be linearized, i.e. converted to a list.+class Linearizable a b | a -> b where+  linearize :: a -> [b]++instance (VS.Storable a, Nested.PrimElt a)+         => Linearizable (Nested.Ranked n a) a where+  linearize = VS.toList . Nested.rtoVector++instance (VS.Storable a, Nested.PrimElt a)+         => Linearizable (Nested.Shaped sh a) a where+  linearize = VS.toList . Nested.stoVector++instance Linearizable (RepConcrete y) a+         => Linearizable (Concrete y) a where+  linearize = linearize . unConcrete++instance Linearizable Double Double where+  linearize x = [x]++instance Linearizable Float Float where+  linearize x = [x]++instance Linearizable Int64 Int64 where+  linearize x = [x]++instance Linearizable CInt CInt where+  linearize x = [x]++instance Linearizable Z1 Z1 where+  linearize _ = []++instance {-# OVERLAPPABLE #-} (Foldable t) => Linearizable (t a) a where+  linearize = Data.Foldable.toList