diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,165 @@ + GNU LESSER GENERAL PUBLIC LICENSE + Version 3, 29 June 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + + This version of the GNU Lesser General Public License incorporates +the terms and conditions of version 3 of the GNU General Public +License, supplemented by the additional permissions listed below. + + 0. Additional Definitions. + + As used herein, "this License" refers to version 3 of the GNU Lesser +General Public License, and the "GNU GPL" refers to version 3 of the GNU +General Public License. + + "The Library" refers to a covered work governed by this License, +other than an Application or a Combined Work as defined below. + + An "Application" is any work that makes use of an interface provided +by the Library, but which is not otherwise based on the Library. +Defining a subclass of a class defined by the Library is deemed a mode +of using an interface provided by the Library. + + A "Combined Work" is a work produced by combining or linking an +Application with the Library. The particular version of the Library +with which the Combined Work was made is also called the "Linked +Version". + + The "Minimal Corresponding Source" for a Combined Work means the +Corresponding Source for the Combined Work, excluding any source code +for portions of the Combined Work that, considered in isolation, are +based on the Application, and not on the Linked Version. + + The "Corresponding Application Code" for a Combined Work means the +object code and/or source code for the Application, including any data +and utility programs needed for reproducing the Combined Work from the +Application, but excluding the System Libraries of the Combined Work. + + 1. Exception to Section 3 of the GNU GPL. + + You may convey a covered work under sections 3 and 4 of this License +without being bound by section 3 of the GNU GPL. + + 2. Conveying Modified Versions. + + If you modify a copy of the Library, and, in your modifications, a +facility refers to a function or data to be supplied by an Application +that uses the facility (other than as an argument passed when the +facility is invoked), then you may convey a copy of the modified +version: + + a) under this License, provided that you make a good faith effort to + ensure that, in the event an Application does not supply the + function or data, the facility still operates, and performs + whatever part of its purpose remains meaningful, or + + b) under the GNU GPL, with none of the additional permissions of + this License applicable to that copy. + + 3. Object Code Incorporating Material from Library Header Files. + + The object code form of an Application may incorporate material from +a header file that is part of the Library. You may convey such object +code under terms of your choice, provided that, if the incorporated +material is not limited to numerical parameters, data structure +layouts and accessors, or small macros, inline functions and templates +(ten or fewer lines in length), you do both of the following: + + a) Give prominent notice with each copy of the object code that the + Library is used in it and that the Library and its use are + covered by this License. + + b) Accompany the object code with a copy of the GNU GPL and this license + document. + + 4. Combined Works. + + You may convey a Combined Work under terms of your choice that, +taken together, effectively do not restrict modification of the +portions of the Library contained in the Combined Work and reverse +engineering for debugging such modifications, if you also do each of +the following: + + a) Give prominent notice with each copy of the Combined Work that + the Library is used in it and that the Library and its use are + covered by this License. + + b) Accompany the Combined Work with a copy of the GNU GPL and this license + document. + + c) For a Combined Work that displays copyright notices during + execution, include the copyright notice for the Library among + these notices, as well as a reference directing the user to the + copies of the GNU GPL and this license document. + + d) Do one of the following: + + 0) Convey the Minimal Corresponding Source under the terms of this + License, and the Corresponding Application Code in a form + suitable for, and under terms that permit, the user to + recombine or relink the Application with a modified version of + the Linked Version to produce a modified Combined Work, in the + manner specified by section 6 of the GNU GPL for conveying + Corresponding Source. + + 1) Use a suitable shared library mechanism for linking with the + Library. A suitable mechanism is one that (a) uses at run time + a copy of the Library already present on the user's computer + system, and (b) will operate properly with a modified version + of the Library that is interface-compatible with the Linked + Version. + + e) Provide Installation Information, but only if you would otherwise + be required to provide such information under section 6 of the + GNU GPL, and only to the extent that such information is + necessary to install and execute a modified version of the + Combined Work produced by recombining or relinking the + Application with a modified version of the Linked Version. (If + you use option 4d0, the Installation Information must accompany + the Minimal Corresponding Source and Corresponding Application + Code. If you use option 4d1, you must provide the Installation + Information in the manner specified by section 6 of the GNU GPL + for conveying Corresponding Source.) + + 5. Combined Libraries. + + You may place library facilities that are a work based on the +Library side by side in a single library together with other library +facilities that are not Applications and are not covered by this +License, and convey such a combined library under terms of your +choice, if you do both of the following: + + a) Accompany the combined library with a copy of the same work based + on the Library, uncombined with any other library facilities, + conveyed under the terms of this License. + + b) Give prominent notice with the combined library that part of it + is a work based on the Library, and explaining where to find the + accompanying uncombined form of the same work. + + 6. Revised Versions of the GNU Lesser General Public License. + + The Free Software Foundation may publish revised and/or new versions +of the GNU Lesser General Public License from time to time. Such new +versions will be similar in spirit to the present version, but may +differ in detail to address new problems or concerns. + + Each version is given a distinguishing version number. If the +Library as you received it specifies that a certain numbered version +of the GNU Lesser General Public License "or any later version" +applies to it, you have the option of following the terms and +conditions either of that published version or of any later version +published by the Free Software Foundation. If the Library as you +received it does not specify a version number of the GNU Lesser +General Public License, you may choose any version of the GNU Lesser +General Public License ever published by the Free Software Foundation. + + If the Library as you received it specifies that a proxy can decide +whether future versions of the GNU Lesser General Public License shall +apply, that proxy's public statement of acceptance of any version is +permanent authorization for you to choose that version for the +Library. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/TypedFlow.hs b/TypedFlow.hs new file mode 100644 --- /dev/null +++ b/TypedFlow.hs @@ -0,0 +1,24 @@ +{-| +Module : TypedFlow +Description : Higher-Order Typed Binding to TensorFlow and Deep Learning Library +Copyright : (c) Jean-Philippe Bernardy, 2017 +License : LGPL-3 +Maintainer : jean-philippe.bernardy@gu.se +Stability : experimental + +This module re-exports all functions. +-} + +module TypedFlow + (module TypedFlow.Types + ,module TypedFlow.TF + ,module TypedFlow.Layers + ,module TypedFlow.Learn + ,module GHC.TypeLits) where + +import TypedFlow.TF +import TypedFlow.Types +import TypedFlow.Layers +import TypedFlow.Learn +import GHC.TypeLits + diff --git a/TypedFlow/Layers.hs b/TypedFlow/Layers.hs new file mode 100644 --- /dev/null +++ b/TypedFlow/Layers.hs @@ -0,0 +1,8 @@ + +module TypedFlow.Layers + (module TypedFlow.Layers.Core + ,module TypedFlow.Layers.RNN) where + +import TypedFlow.Layers.Core +import TypedFlow.Layers.RNN + diff --git a/TypedFlow/Layers/Core.hs b/TypedFlow/Layers/Core.hs new file mode 100644 --- /dev/null +++ b/TypedFlow/Layers/Core.hs @@ -0,0 +1,167 @@ +{-| +Module : TypedFlow.Layers.Core +Description : Core layers and combinators. +Copyright : (c) Jean-Philippe Bernardy, 2017 +License : LGPL-3 +Maintainer : jean-philippe.bernardy@gu.se +Stability : experimental +-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE TypeInType #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveFoldable #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UnicodeSyntax #-} +{-# LANGUAGE PatternSynonyms #-} + +module TypedFlow.Layers.Core + ( + -- * Dense + DenseP(..), dense, (#), + -- * Dropout + DropProb(..), mkDropout, mkDropouts, + -- * Embedding + EmbeddingP(..), embedding, + -- * Convolutional + ConvP(..), conv, maxPool2D) + +where + +import Prelude hiding (tanh,Num(..),Floating(..),floor) +import qualified Prelude +import GHC.TypeLits +-- import Text.PrettyPrint.Compact (float) +import TypedFlow.TF +import TypedFlow.Types +import Control.Monad.State (gets) +-- import Data.Type.Equality +-- import Data.Kind (Type,Constraint) +import Data.Monoid ((<>)) +--------------------- +-- Linear functions + + +-- type (a ⊸ b) = DenseP Float32 a b + +-- | A dense layer is a linear function form a to b: a transformation matrix and a bias. +data DenseP t a b = DenseP {denseWeights :: Tensor '[a,b] (Flt t) + ,denseBiases :: Tensor '[b] (Flt t)} + +----------------------- +-- Feed-forward layers + +-- | Parameters for the embedding layers +newtype EmbeddingP numObjects embeddingSize t = EmbeddingP (Tensor '[numObjects, embeddingSize] ('Typ 'Float t)) + +instance (KnownNat numObjects, KnownBits b, KnownNat embeddingSize) => KnownTensors (EmbeddingP numObjects embeddingSize b) where + travTensor f s (EmbeddingP p) = EmbeddingP <$> travTensor f s p + +instance (KnownNat numObjects, KnownBits b, KnownNat embeddingSize) => ParamWithDefault (EmbeddingP numObjects embeddingSize b) where + defaultInitializer = EmbeddingP (randomUniform (-0.05) 0.05) + +-- | embedding layer +embedding :: ∀ embeddingSize numObjects batchSize t. + EmbeddingP numObjects embeddingSize t -> Tensor '[batchSize] Int32 -> Tensor '[embeddingSize,batchSize] ('Typ 'Float t) +embedding (EmbeddingP param) input = gather @ '[embeddingSize] (transpose param) input + +instance (KnownNat a, KnownNat b, KnownBits t) => KnownTensors (DenseP t a b) where + travTensor f s (DenseP x y) = DenseP <$> travTensor f (s<>"_w") x <*> travTensor f (s<>"_bias") y + +instance (KnownNat n, KnownNat m, KnownBits b) => ParamWithDefault (DenseP b n m) where + defaultInitializer = DenseP glorotUniform (truncatedNormal 0.1) + +-- | Dense layer (Apply a linear function) +(#), dense :: ∀m n batchSize t. DenseP t n m -> Tensor '[n, batchSize] (Flt t) -> Tensor '[m, batchSize] (Flt t) +(DenseP weightMatrix bias) # v = (weightMatrix ∙ v) + bias + +dense = (#) + +-- | A drop probability. (This type is used to make sure one does not +-- confuse keep probability and drop probability) +data DropProb = DropProb Float + +-- | Generate a dropout function. The mask applied by the returned +-- function will be constant for any given call to mkDropout. This +-- behavior allows to use the same mask in the several steps of an +-- RNN. +mkDropout :: forall s t. KnownShape s => KnownBits t => DropProb -> Gen (Tensor s ('Typ 'Float t) -> Tensor s ('Typ 'Float t)) +mkDropout (DropProb dropProb) = do + let keepProb = 1.0 Prelude.- dropProb + isTraining <- gets genTrainingPlaceholder + mask <- assign (if_ isTraining + (floor (randomUniform keepProb (1 Prelude.+ keepProb)) ⊘ constant keepProb) + ones) + return (mask ⊙) + +newtype EndoTensor t s = EndoTensor (Tensor s t -> Tensor s t) + +-- | Generate a dropout function for an heterogeneous tensor vector. +mkDropouts :: KnownBits t => KnownLen shapes => All KnownShape shapes => DropProb -> Gen (HTV ('Typ 'Float t) shapes -> HTV ('Typ 'Float t) shapes) +mkDropouts d = appEndoTensor <$> mkDropouts' shapeSList where + mkDropouts' :: forall shapes t. KnownBits t => All KnownShape shapes => + SList shapes -> Gen (NP (EndoTensor ('Typ 'Float t)) shapes) + mkDropouts' LZ = return Unit + mkDropouts' (LS _ rest) = do + x <- mkDropout d + xs <- mkDropouts' rest + return (EndoTensor x :* xs) + + appEndoTensor :: NP (EndoTensor t) s -> HTV t s -> HTV t s + appEndoTensor Unit Unit = Unit + appEndoTensor (EndoTensor f :* fs) (F x :* xs) = F (f x) :* appEndoTensor fs xs + + +------------------------ +-- Convolutional layers + +data ConvP t outChannels inChannels filterSpatialShape + = ConvP (T ('[outChannels,inChannels] ++ filterSpatialShape) ('Typ 'Float t)) (T '[outChannels] ('Typ 'Float t)) + +instance (KnownNat outChannels,KnownNat inChannels, KnownShape filterSpatialShape, KnownBits t) => + ParamWithDefault (ConvP t outChannels inChannels filterSpatialShape) where + defaultInitializer = prodHomo @filterSpatialShape @'[outChannels] $ + knownAppend @filterSpatialShape @'[outChannels] $ + ConvP (transposeN' (reshape i)) (constant 0.1) + where i :: T '[inChannels,Product filterSpatialShape* outChannels] (Flt t) + i = knownProduct @filterSpatialShape glorotUniform + +instance (KnownNat outChannels,KnownNat inChannels, KnownShape filterSpatialShape, KnownBits t) => + KnownTensors (ConvP t outChannels inChannels filterSpatialShape) where + travTensor f s (ConvP x y) = ConvP <$> travTensor f (s<>"_filters") x <*> travTensor f (s <> "_biases") y + +-- | Size-preserving convolution layer +conv :: forall outChannels filterSpatialShape inChannels s t. + ((1 + Length filterSpatialShape) ~ Length s, + Length filterSpatialShape <= 3, + KnownLen filterSpatialShape) => -- the last dim of s is the batch size + ConvP t outChannels inChannels filterSpatialShape -> + T ('[inChannels] ++ s) ('Typ 'Float t) -> (T ('[outChannels] ++ s) ('Typ 'Float t)) +conv (ConvP filters bias) input = convolution input filters + bias + + +-- | 2 by 2 maxpool layer. +maxPool2D :: forall stridex (stridey::Nat) batch height width channels t. + (KnownNat stridex, KnownNat stridey) => + T '[channels,width*stridex,height*stridex,batch] (Flt t) -> T '[channels,width,height,batch] (Flt t) +maxPool2D (T value) = T (funcall "tf.nn.max_pool" [value + ,showShape @'[1,stridex,stridey,1] + ,showShape @'[1,stridex,stridey,1] + ,named "padding" (str "SAME") ]) + diff --git a/TypedFlow/Layers/RNN.hs b/TypedFlow/Layers/RNN.hs new file mode 100644 --- /dev/null +++ b/TypedFlow/Layers/RNN.hs @@ -0,0 +1,463 @@ +{-| +Module : TypedFlow.Layers.RNN +Description : RNN cells, layers and combinators. +Copyright : (c) Jean-Philippe Bernardy, 2017 +License : LGPL-3 +Maintainer : jean-philippe.bernardy@gu.se +Stability : experimental +-} + +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE TypeInType #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveFoldable #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UnicodeSyntax #-} +{-# LANGUAGE PatternSynonyms #-} + +module TypedFlow.Layers.RNN ( + -- * Types + RnnCell, RnnLayer, + -- * Combinators + stackRnnCells, (.-.), + stackRnnLayers, (.--.), + bothRnnLayers,(.++.), + withBypass, + onStates, + timeDistribute, timeDistribute', + -- * RNN Cells + cellInitializerBit, + LSTMP(..), + lstm, + GRUP(..), + gru, + -- * RNN unfolding functions + rnn, + rnnBackward, + rnnBackwardsWithCull, + rnnWithCull, + -- * Attention mechanisms + -- ** Scoring functions + AttentionScoring, + multiplicativeScoring, + AdditiveScoringP(..), additiveScoring, + -- ** Attention functions + AttentionFunction, + uniformAttn, + luongAttention, + -- ** Attention combinators + attentiveWithFeedback + ) + +where + +import Prelude hiding (tanh,Num(..),Floating(..),floor) +import GHC.TypeLits +-- import Text.PrettyPrint.Compact (float) +import TypedFlow.TF +import TypedFlow.Types +import TypedFlow.Layers.Core (DenseP(..),(#)) +-- import Data.Type.Equality +-- import Data.Kind (Type,Constraint) +import Data.Monoid ((<>)) + + +-- | A cell in an rnn. @state@ is the state propagated through time. +type RnnCell t states input output = (HTV (Flt t) states , input) -> Gen (HTV (Flt t) states , output) + +-- | A layer in an rnn. @n@ is the length of the time sequence. @state@ is the state propagated through time. +type RnnLayer b n state input t output u = HTV (Flt b) state -> Tensor (n ': input) t -> Gen (HTV (Flt b) state , Tensor (n ': output) u) + +-------------------------------------- +-- Combinators + + +-- | Compose two rnn layers. This is useful for example to combine +-- forward and backward layers. +(.--.),stackRnnLayers :: forall s1 s2 a t b u c v n bits. KnownLen s1 => + RnnLayer bits n s1 a t b u -> RnnLayer bits n s2 b u c v -> RnnLayer bits n (s1 ++ s2) a t c v +stackRnnLayers f g (hsplit @s1 -> (s0,s1)) x = do + (s0',y) <- f s0 x + (s1',z) <- g s1 y + return (happ s0' s1',z) + +infixr .--. +(.--.) = stackRnnLayers + + +-- | Compose two rnn layers in parallel. +bothRnnLayers,(.++.) :: forall s1 s2 a t b u c n bs bits. KnownLen s1 => + RnnLayer bits n s1 a t '[b,bs] u -> RnnLayer bits n s2 a t '[c,bs] u -> RnnLayer bits n (s1 ++ s2) a t '[b+c,bs] u +bothRnnLayers f g (hsplit @s1 -> (s0,s1)) x = do + (s0',y) <- f s0 x + (s1',z) <- g s1 x + return (happ s0' s1',concat1 y z) + + +infixr .++. +(.++.) = bothRnnLayers + +-- | Apply a function on the cell state(s) before running the cell itself. +onStates :: (HTV (Flt t) xs -> HTV (Flt t) xs) -> RnnCell t xs a b -> RnnCell t xs a b +onStates f cell (h,x) = do + cell (f h, x) + +-- | Stack two RNN cells (LHS is run first) +stackRnnCells, (.-.) :: forall s0 s1 a b c t. KnownLen s0 => RnnCell t s0 a b -> RnnCell t s1 b c -> RnnCell t (s0 ++ s1) a c +stackRnnCells l1 l2 (hsplit @s0 -> (s0,s1),x) = do + (s0',y) <- l1 (s0,x) + (s1',z) <- l2 (s1,y) + return ((happ s0' s1'),z) + +(.-.) = stackRnnCells + +-- | Run the cell, and forward the input to the output, by concatenation with the output of the cell. +withBypass :: RnnCell b s0 (T '[x,bs] t) (T '[y,bs] t) -> RnnCell b s0 (T '[x,bs] t) (T '[x+y,bs] t) +withBypass cell (s,x) = do + (s',y) <- cell (s,x) + return (s',concat0 x y) + +-------------------------------------- +-- Cells + +-- | Convert a pure function (feed-forward layer) to an RNN cell by +-- ignoring the RNN state. +timeDistribute :: (a -> b) -> RnnCell t '[] a b +timeDistribute pureLayer = timeDistribute' (return . pureLayer) + +-- | Convert a stateless generator into an RNN cell by ignoring the +-- RNN state. +timeDistribute' :: (a -> Gen b) -> RnnCell t '[] a b +timeDistribute' stateLess (Unit,a) = do + b <- stateLess a + return (Unit,b) + +-- | Standard RNN gate initializer. (The recurrent kernel is +-- orthogonal to avoid divergence; the input kernel is glorot) +cellInitializerBit :: ∀ n x t. (KnownNat n, KnownNat x, KnownBits t) => DenseP t (n + x) n +cellInitializerBit = DenseP (concat0 recurrentInitializer kernelInitializer) biasInitializer + where + recurrentInitializer :: Tensor '[n, n] ('Typ 'Float t) + recurrentInitializer = randomOrthogonal + kernelInitializer :: Tensor '[x, n] ('Typ 'Float t) + kernelInitializer = glorotUniform + biasInitializer = zeros + +-- | Parameter for an LSTM +data LSTMP t n x = LSTMP (DenseP t (n+x) n) (DenseP t (n+x) n) (DenseP t (n+x) n) (DenseP t (n+x) n) + +instance (KnownNat n, KnownNat x, KnownBits t) => KnownTensors (LSTMP t n x) where + travTensor f s (LSTMP x y z w) = LSTMP <$> travTensor f (s<>"_f") x <*> travTensor f (s<>"_i") y <*> travTensor f (s<>"_c") z <*> travTensor f (s<>"_o") w +instance (KnownNat n, KnownNat x, KnownBits t) => ParamWithDefault (LSTMP t n x) where + defaultInitializer = LSTMP forgetInit cellInitializerBit cellInitializerBit cellInitializerBit + where forgetInit = DenseP (denseWeights cellInitializerBit) ones + +-- | Standard LSTM +lstm :: ∀ n x bs t. LSTMP t n x -> + RnnCell t '[ '[n,bs], '[n,bs]] (Tensor '[x,bs] (Flt t)) (Tensor '[n,bs] (Flt t)) +lstm (LSTMP wf wi wc wo) (VecPair ht1 ct1, input) = do + hx <- assign (concat0 ht1 input) + let f = sigmoid (wf # hx) + i = sigmoid (wi # hx) + cTilda = tanh (wc # hx) + o = sigmoid (wo # hx) + c <- assign ((f ⊙ ct1) + (i ⊙ cTilda)) + h <- assign (o ⊙ tanh c) + return (VecPair h c, h) + +-- -- | LSTM for an attention model. The result of attention is combined using + to generate output (bad!) +-- attentiveLstmPlus :: forall x n bs t. KnownNat bs => +-- AttentionFunction t bs n n -> +-- LSTMP t n x -> +-- RnnCell t '[ '[n,bs], '[n,bs]] (Tensor '[x,bs] (Flt t)) (Tensor '[n,bs] (Flt t)) +-- attentiveLstmPlus att w x = do +-- (VecPair ht ct, _ht) <- lstm w x +-- a <- att ht +-- let ht' = ht ⊕ a -- alternatively add a dense layer to combine +-- return (VecPair ht' ct, a) + +-- | Parameter for a GRU +data GRUP t n x = GRUP (T [n+x,n] ('Typ 'Float t)) (T [n+x,n] ('Typ 'Float t)) (T [n+x,n] ('Typ 'Float t)) + +instance (KnownNat n, KnownNat x, KnownBits t) => KnownTensors (GRUP t n x) where + travTensor f s (GRUP x y z) = GRUP <$> travTensor f (s<>"_z") x <*> travTensor f (s<>"_r") y <*> travTensor f (s<>"_w") z +instance (KnownNat n, KnownNat x, KnownBits t) => ParamWithDefault (GRUP t n x) where + defaultInitializer = GRUP (denseWeights cellInitializerBit) (denseWeights cellInitializerBit) (denseWeights cellInitializerBit) + + +-- | Standard GRU cell +gru :: ∀ n x bs t. (KnownNat bs, KnownNat n, KnownBits t) => GRUP t n x -> + RnnCell t '[ '[n,bs] ] (Tensor '[x,bs] (Flt t)) (Tensor '[n,bs] (Flt t)) +gru (GRUP wz wr w) (VecSing ht1, xt) = do + hx <- assign (concat0 ht1 xt) + let zt = sigmoid (wz ∙ hx) + rt = sigmoid (wr ∙ hx) + hTilda = tanh (w ∙ (concat0 (rt ⊙ ht1) xt)) + ht <- assign ((ones ⊝ zt) ⊙ ht1 + zt ⊙ hTilda) + return (VecSing ht, ht) + +---------------------------------------------- +-- "Attention" layers + + +-- | An attention scoring function. This function should produce a +-- score (between 0 and 1) for each of the @nValues@ entries of size +-- @valueSize@. +type AttentionScoring t batchSize keySize valueSize nValues = + Tensor '[keySize,batchSize] ('Typ 'Float t) -> Tensor '[nValues,valueSize,batchSize] ('Typ 'Float t) -> Tensor '[nValues,batchSize] ('Typ 'Float t) + +-- | A function which attends to an external input. Typically a +-- function of this type is a closure which has the attended input in +-- its environment. +type AttentionFunction t batchSize keySize valueSize = + T '[keySize,batchSize] (Flt t) -> Gen (T '[valueSize,batchSize] (Flt t)) + +{- NICER, SLOW + +type AttentionScoring t batchSize keySize valueSize = + Tensor '[keySize,batchSize] ('Typ 'Float t) -> Tensor '[valueSize,batchSize] ('Typ 'Float t) -> Tensor '[batchSize] ('Typ 'Float t) + + +-- | @attnExample1 θ h st@ combines each element of the vector h with +-- s, and applies a dense layer with parameters θ. The "winning" +-- element of h (using softmax) is returned. +uniformAttn :: ∀ valueSize m keySize batchSize t. KnownNat m => KnownBits t => + T '[batchSize] Int32 -> + AttentionScoring t batchSize keySize valueSize -> + T '[m,valueSize,batchSize] (Flt t) -> AttentionFunction t batchSize keySize valueSize +uniformAttn lengths score hs_ ht = do + xx <- mapT (score ht) hs_ + let αt :: T '[m,batchSize] (Flt t) + αt = softmax0 (mask ⊙ xx) + ct :: T '[valueSize,batchSize] (Flt t) + ct = squeeze0 (matmul hs_ (expandDim0 αt)) + mask = cast (sequenceMask @m lengths) -- mask according to length + return ct + + + +-- | A multiplicative scoring function. See +-- https://github.com/tensorflow/nmt#background-on-the-attention-mechanism +-- commit 75aa22dfb159f10a1a5b4557777d9ff547c1975a +multiplicativeScoring :: forall valueSize keySize batchSize t. + T [keySize,valueSize] ('Typ 'Float t) -> AttentionScoring t batchSize keySize valueSize +multiplicativeScoring w dt h = h · ir + where ir :: T '[valueSize,batchSize] ('Typ 'Float t) + ir = w ∙ dt + + +additiveScoring :: AdditiveScoringP sz keySize valueSize t -> AttentionScoring t batchSize valueSize keySize +additiveScoring (AdditiveScoringP v w1 w2) dt h = squeeze0 (v ∙ tanh ((w1 ∙ h) ⊕ (w2 ∙ dt))) + +-} + +-- | @attnExample1 θ h st@ combines each element of the vector h with +-- s, and applies a dense layer with parameters θ. The "winning" +-- element of h (using softmax) is returned. +uniformAttn :: ∀ valueSize m keySize batchSize t. KnownNat m => KnownBits t + => AttentionScoring t batchSize keySize valueSize m -- ^ scoring function + -> T '[batchSize] Int32 -- ^ lengths of the inputs + -> T '[m,valueSize,batchSize] (Flt t) -- ^ inputs + -> AttentionFunction t batchSize keySize valueSize +uniformAttn score lengths hs_ ht = do + let αt :: T '[m,batchSize] (Flt t) + xx = score ht hs_ + αt = softmax0 (mask ⊙ xx) + ct :: T '[valueSize,batchSize] (Flt t) + ct = squeeze0 (matmul hs_ (expandDim0 αt)) + mask = cast (sequenceMask @m lengths) -- mask according to length + return ct + +-- | Add some attention to an RnnCell, and feed the attention vector to +-- the next iteration in the rnn. (This follows the diagram at +-- https://github.com/tensorflow/nmt#background-on-the-attention-mechanism +-- commit 75aa22dfb159f10a1a5b4557777d9ff547c1975a). +attentiveWithFeedback ::forall attSize cellSize inputSize bs w ss. + AttentionFunction w bs cellSize attSize -> + RnnCell w ss (T '[inputSize+attSize,bs] (Flt w)) (T '[cellSize,bs] (Flt w)) -> + RnnCell w ('[attSize,bs] ': ss) (T '[inputSize ,bs] (Flt w)) (T '[attSize,bs] (Flt w)) +attentiveWithFeedback attn cell ((F prevAttnVector :* s),x) = do + (s',y) <- cell (s,concat0 x prevAttnVector) + focus <- attn y + return ((F focus :* s'),focus) + +-- -- | LSTM for an attention model. The result of attention is fed to the next step. +-- attentiveLstm :: forall attSize n x bs t. KnownNat bs => +-- AttentionFunction t bs n attSize -> +-- LSTMP t n (x+attSize) -> +-- RnnCell t '[ '[attSize,bs], '[n,bs], '[n,bs] ] (Tensor '[x,bs] (Flt t)) (Tensor '[attSize,bs] (Flt t)) +-- attentiveLstm att w = attentiveWithFeedback att (lstm w) + + +-- | Luong attention function (following +-- https://github.com/tensorflow/nmt#background-on-the-attention-mechanism +-- commit 75aa22dfb159f10a1a5b4557777d9ff547c1975a). +-- Essentially a dense layer with tanh activation, on top of uniform attention. +luongAttention :: ∀ attnSize d m e batchSize w. KnownNat m => KnownBits w + => Tensor '[d+e,attnSize] (Flt w) -- ^ weights for the dense layer + -> AttentionScoring w batchSize e d m -- ^ scoring function + -> Tensor '[batchSize] Int32 -- ^ length of the input + -> T '[m,d,batchSize] (Flt w) -- ^ inputs + -> AttentionFunction w batchSize e attnSize +luongAttention w scoring lens hs_ ht = do + ct <- uniformAttn scoring lens hs_ ht + return (tanh (w ∙ (concat0 ct ht))) + +-- | Multiplicative scoring function +multiplicativeScoring :: forall valueSize keySize batchSize nValues t. + KnownNat batchSize => T [keySize,valueSize] ('Typ 'Float t) -- ^ weights + -> AttentionScoring t batchSize keySize valueSize nValues +multiplicativeScoring w dt hs = squeeze1 (matmul (expandDim1 ir) hs) + where ir :: T '[valueSize,batchSize] ('Typ 'Float t) + ir = w ∙ dt + + +data AdditiveScoringP sz keySize valueSize t = AdditiveScoringP + (Tensor '[sz, 1] ('Typ 'Float t)) + (Tensor '[keySize, sz] ('Typ 'Float t)) + (Tensor '[valueSize, sz] ('Typ 'Float t)) + +instance (KnownNat n, KnownNat k, KnownNat v, KnownBits t) => KnownTensors (AdditiveScoringP k v n t) where + travTensor f s (AdditiveScoringP x y z) = AdditiveScoringP <$> travTensor f (s<>"_v") x <*> travTensor f (s<>"_w1") y <*> travTensor f (s<>"_w2") z +instance (KnownNat n, KnownNat k, KnownNat v, KnownBits t) => ParamWithDefault (AdditiveScoringP k v n t) where + defaultInitializer = AdditiveScoringP glorotUniform glorotUniform glorotUniform + +-- | An additive scoring function. See https://arxiv.org/pdf/1412.7449.pdf +additiveScoring :: forall sz keySize valueSize t nValues batchSize. KnownNat sz => KnownNat keySize => (KnownNat nValues, KnownNat batchSize) => + AdditiveScoringP sz keySize valueSize t -> AttentionScoring t batchSize valueSize keySize nValues +additiveScoring (AdditiveScoringP v w1 w2) dt h = transpose r'' + where w1h :: Tensor '[sz,batchSize, nValues] ('Typ 'Float t) + w1h = transposeN01 @'[sz] (reshape @'[sz,nValues, batchSize] w1h') + w1h' = matmul (reshape @'[keySize, nValues*batchSize] (transpose01 h)) (transpose01 w1) + w2dt = w2 ∙ dt + z' = reshape @'[sz,batchSize*nValues] (tanh (w1h + w2dt)) + r'' = reshape @[batchSize,nValues] (matmul z' (transpose v)) + +--------------------------------------------------------- +-- RNN unfolding + + +-- | Build a RNN by repeating a cell @n@ times. +rnn :: ∀ n state input output t u b. + (KnownNat n, KnownShape input, KnownShape output) => + RnnCell b state (T input t) (T output u) -> RnnLayer b n state input t output u +rnn cell s0 t = do + xs <- unstack0 t + (sFin,us) <- chainForward cell (s0,xs) + return (sFin,stack0 us) +-- There will be lots of stacking and unstacking at each layer for no +-- reason; we should change the in/out from tensors to vectors of +-- tensors. + +-- | Build a RNN by repeating a cell @n@ times. However the state is +-- propagated in the right-to-left direction (decreasing indices in +-- the time dimension of the input and output tensors) +rnnBackward :: ∀ n state input output t u b. + (KnownNat n, KnownShape input, KnownShape output) => + RnnCell b state (T input t) (T output u) -> RnnLayer b n state input t output u + +rnnBackward cell s0 t = do + xs <- unstack0 t + (sFin,us) <- chainBackward cell (s0,xs) + return (sFin,stack0 us) + + + +-- | RNN helper +chainForward :: ∀ state a b n. ((state , a) -> Gen (state , b)) → (state , V n a) -> Gen (state , V n b) +chainForward _ (s0 , V []) = return (s0 , V []) +chainForward f (s0 , V (x:xs)) = do + (s1,x') <- f (s0 , x) + (sFin,V xs') <- chainForward f (s1 , V xs) + return (sFin,V (x':xs')) + +-- | RNN helper +chainBackward :: ∀ state a b n. ((state , a) -> Gen (state , b)) → (state , V n a) -> Gen (state , V n b) +chainBackward _ (s0 , V []) = return (s0 , V []) +chainBackward f (s0 , V (x:xs)) = do + (s1,V xs') <- chainBackward f (s0,V xs) + (sFin, x') <- f (s1,x) + return (sFin,V (x':xs')) + +-- | RNN helper +chainForwardWithState :: ∀ state a b n. ((state , a) -> Gen (state , b)) → (state , V n a) -> Gen (V n b, V n state) +chainForwardWithState _ (_s0 , V []) = return (V [], V []) +chainForwardWithState f (s0 , V (x:xs)) = do + (s1,x') <- f (s0 , x) + (V xs',V ss) <- chainForwardWithState f (s1 , V xs) + return (V (x':xs'), V (s1:ss) ) + +-- -- | RNN helper +-- chainBackwardWithState :: +-- ∀ state a b n. ((state , a) -> Gen (state , b)) → (state , V n a) -> Gen (state , V n b, V n state) +-- chainBackwardWithState _ (s0 , V []) = return (s0 , V [], V []) +-- chainBackwardWithState f (s0 , V (x:xs)) = do +-- (s1,V xs',V ss') <- chainBackwardWithState f (s0,V xs) +-- (sFin, x') <- f (s1,x) +-- return (sFin,V (x':xs'),V (sFin:ss')) + +-- | RNN helper +transposeV :: forall n xs t. All KnownLen xs => + SList xs -> V n (HTV (Flt t) xs) -> HTV (Flt t) (Ap (FMap (Cons n)) xs) +transposeV LZ _ = Unit +transposeV (LS _ n) xxs = F ys' :* yys' + where (ys,yys) = help @(Tail xs) xxs + ys' = stack0 ys + yys' = transposeV n yys + + help :: forall ys x tt. V n (HTV tt (x ': ys)) -> (V n (T x tt) , V n (HTV tt ys)) + help (V xs) = (V (map (fromF . hhead) xs),V (map htail xs)) + +-- | @(gatherFinalStates dynLen states)[i] = states[dynLen[i]]@ +gatherFinalStates :: KnownLen x => KnownNat n => LastEqual bs x => T '[bs] Int32 -> T (n ': x) t -> T x t +gatherFinalStates dynLen states = nth0 0 (reverseSequences dynLen states) + +-- a more efficient algorithm (perhaps:) +-- gatherFinalStates' :: forall x n bs t. KnownLen x => KnownNat n => LastEqual bs x => T '[bs] Int32 -> T (x ++ '[n,bs]) t -> T x (x ++ '[bs]) +-- gatherFinalStates' (T dynLen)t = gather (flattenN2 @x @n @bs t) indexInFlat +-- where indexInFlat = (dynLen - 1) + tf.range(0, bs) * n + +gathers :: forall n bs xs t. All (LastEqual bs) xs => All KnownLen xs => KnownNat n => + SList xs -> T '[bs] Int32 -> HTV (Flt t) (Ap (FMap (Cons n)) xs) -> HTV (Flt t) xs +gathers LZ _ Unit = Unit +gathers (LS _ n) ixs (F x :* xs) = F (gatherFinalStates ixs x) :* gathers @n n ixs xs + +-- | @rnnWithCull dynLen@ constructs an RNN as normal, but returns the +-- state after step @dynLen@ only. +rnnWithCull :: forall n bs x y t u ls b. + KnownLen ls => KnownNat n => KnownLen x => KnownLen y => All KnownLen ls => + All (LastEqual bs) ls => + T '[bs] Int32 -> RnnCell b ls (T x t) (T y u) -> RnnLayer b n ls x t y u +rnnWithCull dynLen cell s0 t = do + xs <- unstack0 t + (us,ss) <- chainForwardWithState cell (s0,xs) + let sss = transposeV @n (shapeSList @ls) ss + return (gathers @n (shapeSList @ls) dynLen sss,stack0 us) + +-- | Like @rnnWithCull@, but states are threaded backwards. +rnnBackwardsWithCull :: forall n bs x y t u ls b. + KnownLen ls => KnownNat n => KnownLen x => KnownLen y => All KnownLen ls => + All (LastEqual bs) ls => LastEqual bs x => LastEqual bs y => + T '[bs] Int32 -> RnnCell b ls (T x t) (T y u) -> RnnLayer b n ls x t y u +rnnBackwardsWithCull dynLen cell s0 t = do + (sFin,hs) <- rnnWithCull dynLen cell s0 (reverseSequences dynLen t) + hs' <- assign (reverseSequences dynLen hs) + return (sFin, hs') + diff --git a/TypedFlow/Learn.hs b/TypedFlow/Learn.hs new file mode 100644 --- /dev/null +++ b/TypedFlow/Learn.hs @@ -0,0 +1,156 @@ +{-| +Module : TypedFlow.Learn +Description : Loss functions and optimization strategies +Copyright : (c) Jean-Philippe Bernardy, 2017 +License : LGPL-3 +Maintainer : jean-philippe.bernardy@gu.se +Stability : experimental +-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveFoldable #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UnicodeSyntax #-} + +module TypedFlow.Learn where + +import TypedFlow.Types +import TypedFlow.TF +import qualified Prelude (Float) +import Prelude (($),return,Maybe(..),(=<<)) +import Text.PrettyPrint.Compact (text) +import Data.Monoid hiding (Last) +import GHC.TypeLits (KnownNat) +import Control.Monad.State (modify, gets) + + +-------------------------------- +-- Model maker. + + +-- | First type argument is the number of classes. +-- @categorical logits gold@ +-- return (prediction, accuraccy, loss) +-- accuracy and prediction are averaged over the batch. +categorical :: forall nCat bs. KnownNat nCat => Model '[nCat,bs] Float32 '[bs] Int32 +categorical logits' y = do + logits <- assign logits' + let y_ = argmax0 logits + modelY = y_ + correctPrediction <- assign (equal y_ y) + modelAccuracy <- assign (reduceMeanAll (cast @Float32 correctPrediction)) + modelLoss <- assign (reduceMeanAll (sparseSoftmaxCrossEntropyWithLogits y logits)) + return ModelOutput{..} + +-- | First type argument is the number of classes. +-- @categoricalDistribution logits gold@ +-- return (prediction, accuraccy, loss) +-- accuracy and prediction are averaged over the batch. +categoricalDistribution :: forall nCat bs. Model '[nCat,bs] Float32 '[nCat,bs] Float32 +categoricalDistribution logits' y = do + logits <- assign logits' + let y_ = softmax0 logits + modelY = y_ + correctPrediction <- assign (equal (argmax0 @'B32 logits) (argmax0 y)) + modelAccuracy <- assign (reduceMeanAll (cast @Float32 correctPrediction)) + modelLoss <- assign (reduceMeanAll (softmaxCrossEntropyWithLogits y logits)) + return ModelOutput{..} + +-- | @timedCategorical targetWeights logits y@ +-- +-- targetWeights: a zero-one matrix of the same size as +-- decoder_outputs. It is intended to mask padding positions outside +-- of the target sequence lengths with values 0. + +timedCategorical :: forall len nCat bs bits. KnownNat nCat => KnownNat bs => KnownNat len => KnownBits bits => + Tensor '[len,bs] (Flt bits) -> Tensor '[len,nCat,bs] (Flt bits) -> Tensor '[len,bs] Int32 -> Gen (ModelOutput '[len,nCat,bs] (Flt bits)) +timedCategorical targetWeights logits' y = do + logits <- assign logits' + let y_ = argmax1 logits + modelY = softmax1 logits + correctPrediction <- assign (equal y_ y) + modelAccuracy <- assign (cast @Float32 (reduceSumAll (flatten2 (cast @(Flt bits) correctPrediction ⊙ targetWeights)) ⊘ reduceSumAll targetWeights)) -- does not work + let crossEntropies = sparseSoftmaxCrossEntropyWithLogits y (transpose01 logits) + modelLoss <- assign (cast @Float32 (reduceMeanAll (crossEntropies ⊙ targetWeights))) + return ModelOutput{..} + +-- | Triple of values that are always output in a model: prediction, loss and accuracy. +data ModelOutput s t = ModelOutput {modelY :: T s t -- ^ prediction + ,modelLoss :: Scalar Float32 + ,modelAccuracy :: Scalar Float32 + } + +-- | A standard modelling function: (input value, gold value) ↦ (prediction, accuracy, loss) +type Model input tIn output tOut = T input tIn -> T output tOut -> Gen (ModelOutput output tOut) + +-- | Model with several binary outputs. +binary :: forall n bs. (KnownNat bs) => Model '[n,bs] Float32 '[n,bs] Int32 +binary logits y = do + sigy_ <- assign (sigmoid logits) + let y_ = cast @Int32 (round sigy_) + modelY = y_ + correctPrediction <- assign (equal y_ y) + modelAccuracy <- assign (reduceMeanAll (cast @Float32 correctPrediction)) + modelLoss <- assign (reduceMeanAll (sigmoidCrossEntropyWithLogits (cast @Float32 y) logits)) + return ModelOutput{..} + +-- | Model compiler options +data Options = Options {maxGradientNorm :: Maybe Prelude.Float -- ^ apply gradient clipping + } + +-- | default model compiler options +defaultOptions :: Options +defaultOptions = Options {maxGradientNorm = Nothing} + +-- | compile a standard model +compile :: forall sx tx sy ty sy_ ty_. + (KnownShape sx, KnownTyp tx, KnownShape sy, KnownTyp ty, KnownShape sy_) => + Options -> (Tensor sx tx -> Tensor sy ty -> Gen (ModelOutput sy_ ty_)) + -- Model input tIn output tOut + -> Gen () +compile options f = compileGen options $ do + x <- placeholder "x" + f x =<< placeholder "y" + + +-- | Generic a model with non-standard parameters ("x" and "y" must be +-- provided as placeholders manually). +compileGen :: forall sy ty. (KnownShape sy) => + Options -> Gen (ModelOutput sy ty) -> Gen () +compileGen Options{..} model = knownLast @sy $ do + gen (text "import tensorflow as tf") + genFun "mkModel" [text "optimizer=tf.train.AdamOptimizer()"] $ do + peekAt "optimizer" (T (text "optimizer")) + peekAt "batch_size" (T (showDim @ (Last sy))) + trainingPhasePlaceholder <- placeholder "training_phase" + modify $ \GState{..} -> GState{genTrainingPlaceholder = trainingPhasePlaceholder,..} + ModelOutput{..} <- model + y_ <- assign modelY + peekAt "y_" y_ + loss <- assign modelLoss + peekAt "loss" loss + accuracy <- assign modelAccuracy + peekAt "accuracy" accuracy + params <- getParameters + peekAt "params" (T params) + trainStep <- assign $ case maxGradientNorm of + Nothing -> T (funcall "optimizer.minimize" [fromTensor loss]) + Just clip -> T (funcall "optimizer.apply_gradients" [funcall "zip" [clipByGlobalNorm clip (grad loss params),params]]) + peekAt "train" trainStep + peeks <- gets genPeeks + gen (text "return " <> dict peeks) diff --git a/TypedFlow/TF.hs b/TypedFlow/TF.hs new file mode 100644 --- /dev/null +++ b/TypedFlow/TF.hs @@ -0,0 +1,737 @@ +{-| +Module : TypedFlow.TF +Description : Binding to tensorflow functions +Copyright : (c) Jean-Philippe Bernardy, 2017 +License : LGPL-3 +Maintainer : jean-philippe.bernardy@gu.se +Stability : experimental + +This module provides direct access to the most commonly used +TensorFlow functions. Higher-level functions are not defined here. +-} + +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveFoldable #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UnicodeSyntax #-} + +module TypedFlow.TF ( + -- * Variables, Parameters + -- ** Parameters + parameter', + parameter, + parameterDefault, + ParamWithDefault(..), + getParameters, + -- ** Persistent variables + persistent, + modifyPersistent, + -- ** Placeholders and outputs + placeholder, + peekAt, + peekAtMany, + -- * Operations + -- ** Constants + zeros, + ones, + constant, + -- ** indexwise unary operators + round, sigmoid, tanh, log, relu, floor, negate, + -- ** Indexwise binary operators + add, (+), (⊕), (⊝), (⊙), (⊘), equal, + -- ** Products + (∙), (·), matmul, + -- ** Reducers + reduceMeanAll, reduceSumAll, + reduceSum, reduceMean, + argmax, argmax0, argmax1, + softmax0, softmax1, + -- ** Gradients + grad, + clipByGlobalNorm, + clipByValue, + -- ** Indexing + last0, nth0, nth0', gather, + -- ** Split and concatenate + split0, slice, slice1, + stack0, unstack0, stackN, + stack1, + concatT, concat0, concat1, + -- ** Reshaping + expandDim, + expandDim0, squeeze0, + expandDim1, squeeze1, + flatten2, inflate2, flattenN2, + flatten3, inflate3, + reshape, flattenAll, inflateAll, + -- ** Transposition + transpose, transposeN, transposeN', transpose01, transposeN01, + -- ** Sequences + reverseSequences, sequenceMask, + -- ** Misc + cast, + convolution, + oneHot, oneHot0, oneHot1, + -- ** Testing conditions + if_, where_, + -- * Contrib + -- ** Mapping + mapT, mapTN, zipWithT, zipWithTN, + -- ** Losses + sigmoidCrossEntropyWithLogits, + softmaxCrossEntropyWithLogits, + sparseSoftmaxCrossEntropyWithLogits, + -- ** Initializers + truncatedNormal, randomUniform, randomOrthogonal, varianceScaling, glorotUniform, + + -- ** Heterogeneous vectors + repeatT, flattenHTV, inflateHTV, KnownTensors(..), LastEqual + ) where + +import Prelude hiding (tanh,Num(..),Floating(..),round,floor) +import qualified Prelude +import Prelude ((-)) +import Text.PrettyPrint.Compact hiding (Last, All,Product,Sum) +import GHC.TypeLits +import Data.Proxy +import TypedFlow.Types +import Control.Monad (when) + +-- | Repeat a flexible-shape constant vector to form a heterogeneous tensor vector. +repeatT :: forall (ss :: [Shape]) t. All KnownShape ss => KnownLen ss => + (forall s. KnownShape s => T s t) -> HTV t ss +repeatT f = zs (shapeSList @ss) + where zs :: forall (s :: [Shape]). All KnownShape s => SList s -> HTV t s + zs LZ = Unit + zs (LS _ n) = F f :* zs n + +-- | Zeros +zeros :: ∀ t (shape :: Shape). KnownShape shape => KnownTyp t => (T shape t) +zeros = T (funcall "tf.zeros" [showShape @shape, named "dtype" (showTyp @t)]) + +-- | Ones +ones :: ∀ t (shape :: Shape). KnownShape shape => KnownTyp t => (T shape t) +ones = T (funcall "tf.ones" [showShape @shape, named "dtype" (showTyp @t)]) + +-- | Constant +constant :: forall s w. KnownShape s => KnownBits w => Float -> T s ('Typ 'Float w) +constant c = T (funcall "tf.constant" [float c, named "shape" (showShape @s), named "dtype" (showTyp @(Flt w))]) + +-- TODO: use a different type for persistent? +-- | Declare variable which persists between calls to session.run. +persistent :: ∀ (shape :: Shape) t. (KnownTyp t,KnownShape shape) => Bool -> String -> T shape t -> Gen (T shape t) +persistent trainable name (T initial) = do + v <- newVar + when trainable (newParameter (ParamInfo name (shapeToList @shape) (typVal @t) (T v))) + v <-- funcall "tf.Variable" [initial, named "name" (string (show (name))), named "trainable" (bool trainable)] + return (T v) + + +-- | Declare a parameter to optimize. The shape of parameter should +-- not depend on dimensions which can change between runs, such as the +-- batch size. +parameter' :: ∀ (shape :: Shape) t. (KnownTyp t,KnownShape shape) => String -> T shape t -> Gen (T shape t) +parameter' = persistent True + +-- | Name a tensor so that it is made available for session.run. +peekAt :: String -> Tensor s t -> Gen () +peekAt p (T v) = peekAtAny p v + +peekAtMany :: String -> HTV t xs -> Gen () +peekAtMany p htv = peekAtAny p (list $ htoList $ hmap (\(F (T x)) -> K x) htv) + + +-- | Modify a mutable tensor. Attention: for the assignment to happen, +-- the resulting tensor must be evaluated! +modifyPersistent :: T s t -> T s t -> T s t +modifyPersistent (T ref) (T value) = T (funcall "tf.assign" [ref,value]) + +-- TODO: get the parameters from the genParams field +-- | Return a list of parameters. +getParameters :: Gen UntypedExpression +getParameters = do + v <- newVar + v <-- text "tf.trainable_variables()" + return v + +-- TODO: get the parameters from the genParams field + + +-- TODO: gradient wrt. a HTV +-- | Gradient of wrt. given parameters. +grad :: T s Float32 -> UntypedExpression -> UntypedExpression +grad (T y) vars = funcall "tf.gradients" [y, vars] + +-- -- | Gradient of wrt. given parameters. +-- grad' :: KnownLen xs => T s Float32 -> HHTV xs -> Gen (HHTV xs) +-- grad' (T y) vars = do +-- v <- newVar +-- v <-- funcall "tf.gradients" [y, list (htoList (hmap (\(Uncurry (T x)) -> K x) vars)) ] +-- return (mkArr 0 shapeSList v) +-- where mkArr :: forall xs. Int -> SList xs -> DOC -> HHTV xs +-- mkArr _ LZ _ = Unit +-- mkArr i (LS _ n) v = Uncurry (T (v <> brackets (int i))) :* mkArr (succ i) n v + + +-- | Clip a gradient +clipByGlobalNorm :: Float -> UntypedExpression -> UntypedExpression +clipByGlobalNorm maxNorm x = funcall "tf.clip_by_global_norm" [x,float maxNorm] <> brackets (int 0) + -- clip_by_global_norm returns a couple (clipped grads, global_norm) + +-- | Clip a tensor +clipByValue :: Float -> Float -> T s (Flt t) -> T s (Flt t) +clipByValue lo hi (T x) = T (funcall "tf.clip_by_value" [x, float lo, float hi]) + +-- | Placeholder (to fill) +placeholder :: ∀t s. (KnownShape s, KnownTyp t) => String -> Gen (T s t) +placeholder n = do + let name = text n + name <-- funcall "tf.placeholder" [showTyp @t, named "shape" (showShape @s), named "name" (text (show n))] + peekAt n (T name) + return (T name) + +-- | Internal. Use 'reduceMeanAll', etc. instead. +reduceAll :: String -> Tensor s t -> Tensor '[] t +reduceAll op = unOp ("tf.reduce_" ++ op) + +-- | Mean value of the input tensor. +reduceMeanAll, reduceSumAll :: ∀ (s :: Shape) t. Tensor s t -> Tensor '[] t +reduceMeanAll = reduceAll "mean" +reduceSumAll = reduceAll "sum" + +-- | Internal. Use 'reduceSum', etc. instead. +reduce :: ∀ n s t. (KnownLen s,KnownPeano n) => String -> T s t -> T (Take n s ++ Drop ('Succ n) s) t +reduce op (T x) = T (funcall ("tf.reduce_" ++ op) [x, text "axis=" <> integer (listLen @ s - peanoInt @n - 1)]) + +-- | Sum along a given dimension +reduceSum, reduceMean :: ∀n s t. (KnownLen s,KnownPeano n) => T s t -> T (Take n s ++ Drop ('Succ n) s) t +reduceSum = reduce @n "sum" +reduceMean = reduce @n "mean" + + +-- | Sum along the first dimension +reduceSum0 :: ∀ s' n t. KnownLen s' => Tensor (n ': s') t -> Tensor s' t +reduceSum0 = reduceSum @Dim0 + +-- | Add two tensors, broacasting along shape @s@ +add :: ∀ s d t. Tensor (d++s) t -> Tensor d t -> Tensor (d++s) t -- note ++s for for 'broadcasting' +add = binOp "tf.add" + +-- add_n :: ∀ s t. [Tensor s t] -> Tensor s t +-- add_n = error "add_n not implemented" + +-- | Add two tensors, broacasting along shape @s@ +(+) :: ∀ (d :: Shape) (s :: Shape) t. Tensor (d ++ s) t -> Tensor d t -> Tensor (d ++ s) t +(+) = add @s @d +infixl 6 + + +-- | Indexwise equality test. +equal :: Tensor d t -> Tensor d t -> Tensor d TFBool +equal = binOp "tf.equal" + +-- | Indexwise operator +(⊕), (⊝), (⊙), (⊘) :: ∀ (s :: Shape) t. Tensor s t -> Tensor s t -> Tensor s t +(⊝) = binOp "tf.subtract" +(⊙) = binOp "tf.multiply" +(⊘) = binOp "tf.divide" +(⊕) = binOp "tf.add" + +infixl 7 ⊙,⊘ +infixl 6 ⊕,⊝ + +-- | Matrix multiplication (note that shape @s@ is preserved) +matmul :: Tensor (o ': n ': s) t -> Tensor (m ': o ': s) t -> Tensor (m ': n ': s) t +matmul = binOp "tf.matmul" + +round, sigmoid, tanh, log, relu, floor + :: ∀ s t. Tensor s ('Typ 'Float t) -> Tensor s ('Typ 'Float t) +sigmoid = unOp "tf.sigmoid" +tanh = unOp "tf.tanh" +log = unOp "tf.log" +relu = unOp "tf.nn.relu" +round = unOp "tf.round" +floor = unOp "tf.floor" + +negate :: ∀ s t. T s t -> T s t +negate = unOp "-" + +-- | Split a tensor on the first dimension +split0 :: ∀ n m batchShape t. (KnownNat n, KnownNat m, KnownLen batchShape) => + Tensor ((n + m) ': batchShape) t -> Gen (Tensor (n ': batchShape) t, Tensor (m ': batchShape) t) +split0 (T x) = do + v1 <- newVar + v2 <- newVar + gen (v1 <> text "," <> v2 <> text " = " <> funcall "tf.split" [x, list [showDim @ n, showDim @ m], text "axis=" <> showShapeLen @batchShape]) + return (T v1, T v2) + +-- | Concatenate tensors on dimension @n@ +concatT :: ∀ n d1 d2 s t. (KnownPeano n, KnownLen s, (d1+d2) ~ At n s) => + T (Take n s ++ (d1 ': Drop ('Succ n) s)) t -> T (Take n s ++ (d2 ': Drop ('Succ n) s)) t -> T s t +concatT (T x) (T y) = T (funcall "tf.concat" [list [x,y], named "axis" (integer (listLen @s - peanoInt @n - 1))]) + +-- | Concatenate tensors on the first dimension +concat0 :: ∀ ys d1 d2 t. (KnownLen ys) => T (d1 ': ys) t -> T (d2 ': ys) t -> T ((d1 + d2) ': ys) t +concat0 = concatT @Dim0 + +-- | Concatenate tensors on the second dimension +concat1 :: ∀ n ys d1 d2 t. (KnownLen ys) => T (n ': d1 ': ys) t -> T (n ': d2 ': ys) t -> T (n ': (d1 + d2) ': ys) t +concat1 = concatT @Dim1 + +-- | Add an extra dimension at axis (@n@) of size 1. +expandDim :: forall n s t. (KnownLen s, KnownPeano n) => Tensor s t -> Tensor (Take n s ++ (1 ': Drop n s)) t +expandDim (T x) = (T (funcall "tf.expand_dims" [x, named "axis" (integer (listLen @s - peanoInt @n))])) + +-- | Add an extra dimension at axis (0) of size 1. +expandDim0 :: ∀ s t. KnownLen s => Tensor s t -> Tensor (1 ': s) t +expandDim0 = expandDim @Dim0 + +-- | Add an extra dimension at axis (1) of size 1. +expandDim1 :: ∀ n s t. KnownShape s => Tensor (n ': s) t -> Tensor (n ': 1 ': s) t +expandDim1 = expandDim @Dim1 + +-- -- | Tile a tensor along the first dimension +-- tile :: forall m n s t. (KnownNat m) => Tensor (n ': s) t -> Tensor ((m * n) ': s) t +-- tile (T x) = T (funcall "tf.tile" [x, integer (natVal (Proxy @m))]) +-- This implementation is incorrect. + +-- -- | Replicate a tensor +-- replicateT :: ∀ n s t. (KnownNat n, KnownLen s) => T s t -> T (n ': s) t +-- replicateT = tile @n . expandDim0 + +-- | Remove a dimension if its size is 1. +squeeze :: ∀ s0 s1 t. KnownLen s1 => Tensor (s0 ++ (1 ': s1)) t -> Tensor (s0 ++ s1) t +squeeze (T x) = T (funcall "tf.squeeze" [x, text "axis=" <> integer (listLen @ s1)]) + +-- | Remove the first dimension if its size is 1. +squeeze0 :: ∀ s t. KnownLen s => Tensor (1 ': s) t -> Tensor s t +squeeze0 = squeeze @ '[] + +-- | Remove the second dimension if its size is 1. +squeeze1 :: ∀ n s t. KnownLen s => Tensor (n ': 1 ': s) t -> Tensor (n ': s) t +squeeze1 = squeeze @ '[n] + +reshape :: ∀ s2 s1 t. KnownShape s2 => Product s1 ~ Product s2 => Tensor s1 t -> Tensor s2 t +reshape = unsafeReshape + +unsafeReshape :: ∀ s2 s1 t. KnownShape s2 => Tensor s1 t -> Tensor s2 t +unsafeReshape (T t) = T (funcall "tf.reshape" [t, showShapeMinus @s2]) + +-- | Reshape a tensor so that the first two dimensions are collapsed +flatten2 :: ∀ m n s t. (KnownNat m, KnownNat n, KnownShape s) => Tensor (m ': n ': s) t -> Tensor (m*n ': s) t +flatten2 = prodAssoc @m @n @(Product s) reshape + +-- | Reshape a tensor so that the last two dimensions are collapsed +flattenN2 :: ∀ s m n t. (KnownNat m, KnownNat n, KnownShape s) => Tensor (s ++ '[m,n]) t -> Tensor (s ++ '[m*n]) t +flattenN2 = prodHomo @s @'[m,n] $ + prodHomo @s @'[m*n] $ + knownAppend @s @'[m*n] $ + reshape + +-- | Reshape a tensor so that the first three dimensions are collapsed +flatten3 :: ∀ m n o s t. (KnownNat m, KnownNat n, KnownNat o, KnownShape s) => Tensor (m ': n ': o ': s) t -> Tensor (m*n*o ': s) t +flatten3 = -- (m * (n * (o * Product s))) + prodAssoc @m @n @(o * Product s) $ + -- (m * n) * (o * Product s) + prodAssoc @(m * n) @o @(Product s) $ + -- ((m * n) * o) * Product s + reshape + + +-- | Reshape a tensor so that the first dimension is expanded into two. +inflate2 :: ∀ m n s t. (KnownNat m, KnownNat n, KnownShape s) => Tensor (m*n ': s) t -> Tensor (m ': n ': s) t +inflate2 = prodAssoc @m @n @(Product s) reshape + +-- | Reshape a tensor so that the first dimension is expanded into three. +inflate3 :: ∀ m n o s t. (KnownNat m, KnownNat n, KnownNat o, KnownShape s) => Tensor (m*n*o ': s) t -> Tensor (m ': n ': o ': s) t +inflate3 = -- (m * (n * (o * Product s))) + prodAssoc @m @n @(o * Product s) $ + -- (m * n) * (o * Product s) + prodAssoc @(m * n) @o @(Product s) $ + -- ((m * n) * o) * Product s + reshape + +-- | Access the last element in a tensor (in the 0th dimension) +last0 :: ∀ n s t. KnownNat n => KnownLen s => T (n ': s) t -> Tensor s t +last0 = nth0 (natVal (Proxy @n) - 1) + +-- | Access the nth element in a tensor (in the 0th dimension) +nth0 :: ∀ n s t. KnownLen s => Integer -> T (n ': s) t -> Tensor s t +nth0 i (T x) = T (x <> list (replicate (fromIntegral (listLen @s)) (text ":") ++ [integer i])) + +-- | Access the nth element in a tensor (in the 0th dimension), with a static index +nth0' :: ∀ n m s t. KnownNat n => KnownLen s => n < m => T (m ': s) t -> Tensor s t +nth0' (T x) = T (x <> list (replicate (fromIntegral (listLen @s)) (text ":") ++ [integer (natVal (Proxy @n))])) + +-- | Take a slice at dimension n from i to j. +slice :: forall n i j s t. KnownNat j => KnownNat i => (i < j, j <= At n s, KnownPeano n, KnownLen s) => + Tensor s t -> Tensor (Take n s ++ ((j-i) ': Drop ('Succ n) s)) t +slice (T x) = T (x <> list (replicate (fromIntegral (listLen @s - peanoInt @n - 1)) (text ":") ++ [integer (natVal (Proxy @i)) <> text ".." <> integer (natVal (Proxy @j))])) + +slice1 :: forall i j m n s t. KnownNat j => KnownNat i => (i < j, j <= m, KnownLen s) => + Tensor (n ': m ': s) t -> Tensor (n ': (j-i) ': s) t +slice1 = slice @Dim1 @i @j + +-- | Split a tensors into @n@ tensors along the first dimension +unstack0 :: ∀ s (n::Nat) t. (KnownLen s, KnownNat n) => Tensor (n ': s) t -> Gen (V n (T s t)) +unstack0 (T x) = do + v <- newVar + v <-- funcall "tf.unstack" [x, text "axis=" <> integer (listLen @ s)] + return $ V $ [ T $ v <> brackets (integer i)| i <- [0..n Prelude.- 1] ] + where n = natVal (Proxy @ n) + +-- | Concatenate @n@ tensors along the first dimension +stack0 :: ∀ s (n::Nat) t. (KnownLen s) => V n (T s t) -> Tensor (n ': s) t +stack0 (V xs) = T (funcall "tf.stack" [list [x | T x <- xs], text "axis=" <> integer (listLen @ s)]) + +-- | Concatenate @n@ tensors along the first dimension +stack1 :: ∀ s (n::Nat) m t. (KnownLen s) => V n (T (m ': s) t) -> Tensor (m ': n ': s) t +stack1 (V xs) = T (funcall "tf.stack" [list [x | T x <- xs], text "axis=" <> integer (listLen @ s)]) + +-- | Concatenate @n@ tensors along the last dimension +stackN :: ∀ s (n::Nat) t. V n (T s t) -> Tensor (s ++ '[n]) t +stackN (V xs) = T (funcall "tf.stack" [list [x | T x <- xs], text "axis=0"]) + +-- | Transposition. See the type for the permutation of dimensions. +transpose :: ∀ s t. T (Reverse s) t -> T s t +transpose = unOp "tf.transpose" + +-- | Transposition. See the type for the permutation of dimensions. +transposeN :: ∀ s n t. KnownLen s => T (n ': s) t -> T (s ++ '[n]) t +transposeN (T x) = T (funcall "tf.transpose" [x, named "perm" (list (map integer (listLen @s:[0.. listLen @s-1])))]) + +-- | Transposition. See the type for the permutation of dimensions. +transposeN' :: ∀ s n t. KnownLen s => T (s ++ '[n]) t -> T (n ': s) t +transposeN' (T x) = T (funcall "tf.transpose" [x, named "perm" (list (map integer ([1.. listLen @s]++[0])))]) + +-- | Transposition. See the type for the permutation of dimensions. +transpose01 :: ∀ s m n t. KnownLen s => T (m ': n ': s) t -> T (n ': m ': s) t +transpose01 (T x) = T (funcall "tf.transpose" [x, named "perm" (list (map integer ([0..l-1] ++ [l Prelude.+ 1,l])))]) + where l = listLen @s + +-- | Transposition. See the type for the permutation of dimensions. +transposeN01 :: ∀ s m n t. T (s ++ [m,n]) t -> T (s ++ [n,m]) t +transposeN01 (T x) = T (funcall "tf.transpose" [x, named "perm" (list (map integer [1,0]))]) + +class LastEqual x xs +instance LastEqual x (x ': '[]) +instance LastEqual x (y2 ': xs) => LastEqual x (y ': (y2 ': xs)) + +-- | Reverse sequences. See https://www.tensorflow.org/api_docs/python/tf/reverse_sequence +reverseSequences :: forall bs n x t. KnownLen x => LastEqual bs x => T '[bs] Int32 -> T (n ': x) t -> T (n ': x) t +reverseSequences (T seqLengths) (T input) = + T (funcall "tf.reverse_sequence" [input, seqLengths, named "seq_axis" (showShapeLen @x),named "batch_axis" (int 0)]) + +-- | Generate a mask of given length for each sequence. +sequenceMask :: forall maxlen bs. KnownNat maxlen => Tensor '[bs] Int32 -> Tensor '[maxlen,bs] TFBool +sequenceMask (T x) = T (funcall "tf.sequence_mask" [x, named "maxlen" (showDim @maxlen)]) + + +-- | @(gather x ix)[k] = x[ix[k]]@. See https://www.tensorflow.org/api_docs/python/tf/gather +gather :: ∀s n indexShape t. T (s ++ '[n]) t -> T indexShape Int32 -> T (s ++ indexShape) t +gather = binOp "tf.gather" + +-- | Size-preserving convolution operation. +convolution :: forall outputChannels filterSpatialShape inChannels s t. + KnownLen filterSpatialShape + => Length filterSpatialShape <= 3 + => ((1 + Length filterSpatialShape) ~ Length s) -- the last dim of s is the batch size + => T ('[inChannels] ++ s) t -- ^ input tensor (batched) + -> T ('[outputChannels,inChannels] ++ filterSpatialShape) t -- ^ filters + -> T ('[outputChannels] ++ s) t +convolution (T input) (T filters) = T (funcall "tf.nn.convolution" [input,filters + ,named "padding" (text (show "SAME")) -- otherwise the shape s changes + ,named "data_format" (text (show dataFormat))]) + where dataFormat = case listLen @ filterSpatialShape of + 1 -> "NWC" + 2 -> "NHWC" + 3 -> "NDHWC" + _ -> error "convolution: more than 3 spatial dimensions are not supported!" + + +-- poolNC :: forall dim s inputSpatialShape channels batchSize t. +-- (inputSpatialShape ~ Take dim s, '[batchSize] ~ Drop dim s) => +-- T ('[channels] ++ s) t -> +-- Vec dim -> String -> String -> +-- T ('[channels] ++ s) t +-- poolNC (T input) windowShape poolingType padding = +-- T (funcall "tf.nn.pool" [input,list (map float (vecToList windowShape)),text poolingType,text padding,named "data_format" (text "NWC")]) + +-- Difficulty: relate windowSize, inputSpatialShape, outputSpatialShape + +-- | Softmax along the first dimension +softmax0 :: T (n ': s) ('Typ 'Float w) -> T (n ': s) ('Typ 'Float w) +softmax0 = unOp "tf.nn.softmax" + +-- | Softmax along the second dimension +softmax1 :: forall n m s w. KnownLen s => T (m ': n ': s) ('Typ 'Float w) -> T (m ': n ': s) ('Typ 'Float w) +softmax1 (T x) = T (funcall "tf.nn.softmax" [x, named "dim" (showShapeLen @s)]) + +-- | Argmax along dimension @n@ +argmax :: forall n u m s t. (KnownLen s, KnownPeano n,KnownBits u) => Tensor (Take n s ++ (m ': Drop n s)) t -> Tensor s ('Typ 'Int u) +argmax (T t) = T (funcall "tf.argmax" [t, named "axis" (integer ((listLen @ s) - peanoInt @n)) , named "output_type" (showTyp @('Typ 'Int u))]) + +-- | Argmax along the first dimension +argmax0 :: forall u n s t. (KnownLen s, KnownBits u) => T (n ': s) t -> T s ('Typ 'Int u) +argmax0 = argmax @Dim0 + +-- | Argmax along the second dimension +argmax1 :: forall u m n s t. (KnownLen s, KnownBits u) => T (m ': n ': s) t -> T (m ': s) ('Typ 'Int u) +argmax1 = argmax @Dim1 + +-- | Cast the element type. +cast :: forall u s t. KnownTyp u => T s t -> T s u +cast (T t) = T (funcall "tf.cast" [t, showTyp @ u]) + + +-- | (dense) softmax cross entropy with logits. +softmaxCrossEntropyWithLogits :: Tensor '[numClasses,batchSize] Float32 -- ^ labels + -> Tensor '[numClasses,batchSize] Float32 -- ^ logits + -> Tensor '[batchSize] Float32 +softmaxCrossEntropyWithLogits (T labels) (T logits) = + T (funcall "tf.nn.softmax_cross_entropy_with_logits" [named "labels" labels,named "logits" logits]) + +-- | Computes sigmoid cross entropy given logits. Measures the +-- probability error in discrete classification tasks in which each +-- class is independent and not mutually exclusive. For instance, one +-- could perform multilabel classification where a picture can contain +-- both an elephant and a dog at the same time. See +-- https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits +sigmoidCrossEntropyWithLogits :: Tensor s (Flt w) -- ^ labels + -> Tensor s (Flt w) -- ^ logits + -> Tensor s (Flt w) +sigmoidCrossEntropyWithLogits (T labels) (T logits) = + T (funcall "tf.nn.sigmoid_cross_entropy_with_logits" [named "labels" labels,named "logits" logits]) + +-- | sparse softmax cross entropy with logits. +sparseSoftmaxCrossEntropyWithLogits :: Tensor s Int32 -- ^ desired labels + -> Tensor (numClasses ': s) (Flt t) -- ^ predictions + -> Tensor s (Flt t) +sparseSoftmaxCrossEntropyWithLogits (T labels) (T logits) = + T (funcall "tf.nn.sparse_softmax_cross_entropy_with_logits" [named "labels" labels,named "logits" logits]) + +-- | One hot vector along axis @n@ +oneHot :: forall n numClasses s w t. KnownNat numClasses => KnownBits t => + (KnownLen s, KnownPeano n) => Tensor s ('Typ 'Int w) -> Tensor (Take n s ++ (numClasses ': Drop n s)) (Flt t) +oneHot (T x) = T (funcall "tf.one_hot" [x, named "depth" (showDim @numClasses), named "axis" (integer (listLen @s - peanoInt @n)), named "dtype" (showTyp @(Flt t))]) + +-- | One hot vector along axis 0 +oneHot0 :: forall numClasses w batchSize t. KnownNat numClasses => KnownBits t => Tensor '[batchSize] ('Typ 'Int w) -> Tensor '[numClasses,batchSize] (Flt t) +oneHot0 = oneHot @Dim0 + +-- | One hot vector along axis 1 +oneHot1 :: forall numClasses w batchSize m t. KnownNat numClasses => KnownBits t => Tensor '[m,batchSize] ('Typ 'Int w) -> Tensor '[m,numClasses,batchSize] (Flt t) +oneHot1 = oneHot @Dim1 + +-- | Generate a random tensor where each individual element is picked +-- in a normal distribution with given standard deviation. +truncatedNormal :: forall s w. KnownShape s => KnownBits w => Float -> T s ('Typ 'Float w) +truncatedNormal stddev = T (funcall "tf.truncated_normal" [showShape @s, named "stddev" (float stddev), named "dtype" (showTyp @(Flt w))]) + +-- | Generate a random tensor where each individual element is picked +-- in a uniform distribution with given bounds. +randomUniform :: forall s t. (KnownShape s, KnownTyp t) => Float -> Float -> T s t +randomUniform low high = T (funcall "tf.random_uniform" [showShape @s + ,named "minval" (float low) + ,named "maxval" (float high) + ,named "dtype" (showTyp @t)]) + + +-- | Generate an orthorgonal matrix. If the output has more dimensions +-- than 2 the matrix is reshaped. +randomOrthogonal :: forall n s t. (KnownBits t, KnownNat n, KnownShape s) => T (n ':s) ('Typ 'Float t) +randomOrthogonal = T (funcall' (funcall "tf.orthogonal_initializer" [named "dtype" (showTyp @('Typ 'Float t))]) + [named "shape" (showShape @(n ': s))]) + +--------------------------- +-- Contrib +data VarianceScaleMode = VSFanIn | VSFanOut | VSAvg +data Distrib = NormalDistr | UniformDistr + +-- | Random tensor with variance scaling according to deeplearning lore. +varianceScaling :: forall inDim outDim t. KnownNat inDim => (KnownNat outDim, KnownBits t) => + Float -> VarianceScaleMode -> Distrib -> Tensor '[inDim,outDim] ('Typ 'Float t) +varianceScaling factor mode distr = case distr of + UniformDistr -> randomUniform (-limit) limit + NormalDistr -> truncatedNormal limit + where + fan_in = fromIntegral (natVal (Proxy @inDim)) + fan_out = fromIntegral (natVal (Proxy @outDim)) + n = max 1 $ case mode of + VSFanIn -> fan_in + VSFanOut -> fan_out + VSAvg -> (fan_in Prelude.+ fan_out) / 2 + limit = Prelude.sqrt ((case distr of NormalDistr -> 1.3; UniformDistr -> 3) Prelude.* factor / n) + + +glorotUniform :: forall inDim outDim t. KnownNat inDim => (KnownNat outDim, KnownBits t) => Tensor '[inDim,outDim] ('Typ 'Float t) +glorotUniform = varianceScaling 1 VSAvg UniformDistr + +---------------- +-- Helpers + +-- matvecmulBatch :: ∀ s cols rows t. (KnownLen s) => Tensor (cols ': rows ': s) t -> Tensor (cols ': s) t -> Tensor (rows ': s) t +-- matvecmulBatch m v = squeeze0 (matmul m (expandDim0 v)) + +-- | Product of a matrix of weight with a (batched) vector . +(∙) :: Tensor '[cols, rows] t -> Tensor '[cols,batchSize] t -> Tensor '[rows,batchSize] t +m ∙ v = matmul v (transpose m) +infixl 7 ∙ + +-- | Dot product between two batched vectors. +(·) :: ∀ cols batchSize t. Tensor '[cols,batchSize] t -> Tensor '[cols,batchSize] t -> Tensor '[batchSize] t +x · y = reduceSum0 (x ⊙ y) +infixl 7 · + +-- mapT' :: forall s t r u n. KnownLen r => KnownLen s => KnownNat n => (T s t -> T r u) -> T (n ': s) t -> Gen (T (n ': r) u) +-- mapT' f t = do +-- xs <- unstack t +-- return (stack (fmap f xs)) + +-- | Map a function along the first dimension of a tensor +mapT :: forall s t r u n. KnownTyp u => KnownLen r => KnownLen s => (T s t -> T r u) -> T (n ': s) t -> Gen (T (n ': r) u) +mapT f x = do + x' <- mapTN @n f (transposeN @s @n x) + return (transposeN' @r x') + +-- | Map a function along the last dimension of a tensor +mapTN :: forall n s t r u. KnownTyp u => (T s t -> T r u) -> T (s ++ '[n]) t -> Gen(T (r ++ '[n]) u) +mapTN f t = do + fn <- lambda f + return (T (funcall "tf.map_fn" [fn, fromTensor t, named "dtype" (showTyp @u)])) + +-- TODO: separate harmless and harmful effects. (the big question: are assignments harmful?) + +zipWithT :: forall (s :: [Nat]) (t :: Typ) (s1 :: [Nat]) (t1 :: Typ) (s2 :: Shape) (n :: Nat) (t2 :: Typ). + KnownNat n => (KnownLen s, KnownLen s2, KnownLen s1) => KnownTyp t2 => + (T s t -> T s1 t1 -> T s2 t2) + -> Tensor (n ': s) t + -> Tensor (n ': s1) t1 + -> Gen (Tensor (n ': s2) t2) +zipWithT f x y = do + -- xs <- unstack x + -- ys <- unstack y + -- return (stack (f <$> xs <*> ys)) + x' <- zipWithTN @n f (transposeN @s @n x) (transposeN @s1 @n y) + return (transposeN' @s2 x') + +zipWithTN :: forall (n :: Nat) (s :: [Nat]) (t :: Typ) (s1 :: [Nat]) (t1 :: Typ) (s2 :: Shape) (t2 :: Typ). + KnownTyp t2 => + (T s t -> T s1 t1 -> T s2 t2) + -> Tensor (s ++ '[n]) t + -> Tensor (s1 ++ '[n]) t1 + -> Gen (Tensor (s2 ++ '[n]) t2) +zipWithTN f (T t) (T u) = do + fn <- lambda2 f + return (T (funcall "tf.map_fn" [fn, tuple [t,u], named "dtype" (showTyp @t2)])) + + +-- apparently tensorflow (python?) is not aware of 2-argument +-- functions; so we do this... thing. +lambda2 :: (T s t -> T s1 t1 -> T s' t') -> Gen UntypedExpression +lambda2 f = do + v <- newVar + let T body = f (T (v <> brackets (int 0))) (T (v <> brackets (int 1))) + return (text "lambda " <> v <> text ": " <> body) + +-- | Selection of a tensor (note: this is a strict operation) +if_ :: Scalar TFBool -> T s t -> T s t -> T s t +if_ (T c) (T x) (T y) = T (funcall "tf.cond" [-- named "pred" -- names have changed between TF 1.1 and TF 1.3 + c, + -- named "true_fn" + (lambda0 x), + -- named "false_fn" + (lambda0 y), + named "strict" (bool True)]) + where lambda0 z = text "lambda: " <> z + +-- | (where_ c x y)[i] = if c[i] then x[i] else y[i] +where_ :: T s TFBool -> T s t -> T s t -> T s t +where_ (T c) (T x) (T y) = T (funcall "tf.where" [c, x, y]) + +------------------------- +-- Generic parameters + +-- | Create a parameter and initialize it with a suitable default for its type. Control the exact initializer using 'parameter'. +parameterDefault :: forall p. ParamWithDefault p => String -> Gen p +parameterDefault name = parameter name defaultInitializer + +-- | Create a parameter. +parameter :: forall p. KnownTensors p => String -> p -> Gen p +parameter = travTensor parameter' + +class KnownTensors p where + -- | traverse all the tensors over tuples of tensors + travTensor :: (forall s t. (KnownTyp t, KnownShape s) => String -> T s t -> Gen (T s t)) -> String -> p -> Gen p + +instance (KnownTyp t, KnownShape shape) => KnownTensors (T shape t) where + travTensor f = f + +instance (KnownTyp t, All KnownShape ys) => KnownTensors (HTV t ys) where + travTensor f s = ttr 0 + where ttr :: forall xs. All KnownShape xs => Int -> HTV t xs -> Gen (HTV t xs) + ttr _ Unit = return Unit + ttr n (F x :* xs) = do + x' <- f (s <> "_" <> show n) x + xs' <- ttr (n Prelude.+ 1) xs + return (F x' :* xs') + +instance (KnownTensors p, KnownTensors q) => KnownTensors (p,q) where + travTensor f s (x,y) = (,) <$> travTensor f (s<>"_fst") x <*> travTensor f (s<>"_snd") y + +instance (KnownTensors p1, KnownTensors p2, KnownTensors p3) => KnownTensors (p1,p2,p3) where + travTensor f s (x,y,z) = (,,) <$> travTensor f (s<>"_1") x <*> travTensor f (s<>"_2") y <*> travTensor f (s<>"_3") z + +instance (KnownTensors p1, KnownTensors p2, KnownTensors p3, KnownTensors p4) => KnownTensors (p1,p2,p3,p4) where + travTensor f s (x,y,z,w) = (,,,) <$> travTensor f (s<>"_1") x <*> travTensor f (s<>"_2") y <*> travTensor f (s<>"_3") z <*> travTensor f (s<>"_4") w + +class KnownTensors p => ParamWithDefault p where + defaultInitializer :: p + +-- | Flatten all the dimensions of the tensor +flattenAll :: forall s t. KnownShape s => Tensor s t -> Tensor '[Product s] t +flattenAll = knownProduct @s reshape + + +flattenHTV :: KnownTyp t => All KnownShape xs => HTV t xs -> Tensor '[Sum (Ap (FMap CProduct) xs)] t +flattenHTV Unit = zeros +flattenHTV (F x :* xs) = concat0 (flattenAll x) (flattenHTV xs) + +inflateAll :: forall s t. KnownShape s => Tensor '[Product s] t -> Tensor s t +inflateAll = knownProduct @s reshape + +class CProduct (xs :: [Nat]) +instance Fun CProduct where type Ap CProduct xs = Product xs + +inflateHTV :: ∀ xs s t. (All KnownShape xs, KnownLen s, KnownLen xs) => + Tensor '[Sum (Ap (FMap CProduct) xs)] t -> Gen (HTV t xs) +inflateHTV (T x) = do + v <- newVar + gen (v <> text " = " <> funcall "tf.split" [x, showShape' (prodshape @xs shapeSList), text "axis=0"]) + return (mkArr @xs 0 shapeSList v) + where mkArr :: forall zs. All KnownShape zs => Int -> SList zs -> DOC -> HTV t zs + mkArr _ LZ _ = Unit + mkArr i (LS _ n) v = F (unsafeReshape (T (v <> brackets (int i)) )):* mkArr (succ i) n v + + prodshape :: forall zs. All KnownShape zs => SList zs -> [Integer] + prodshape LZ = [] + prodshape (LS xx xs) = product (shapeToList' (shapeSListProxy xx)) : prodshape xs + + diff --git a/TypedFlow/Types.hs b/TypedFlow/Types.hs new file mode 100644 --- /dev/null +++ b/TypedFlow/Types.hs @@ -0,0 +1,583 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE UndecidableSuperClasses #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveFoldable #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UnicodeSyntax #-} +{-# LANGUAGE OverloadedStrings #-} + +module TypedFlow.Types where + +import Text.PrettyPrint.Compact hiding (All,Last,Product,Sum) +import GHC.TypeLits +import Unsafe.Coerce +import Data.Proxy +import Control.Monad.State +import Data.Char (toLower) +-- import GHC.Prim (unsafeCoerce#) +import Data.Kind (Type,Constraint) +import Data.Type.Equality + +data Sat (a :: k -> Constraint) (b::k) where + Sat :: a b => Sat a b + +type DOC = Doc () + +type i < j = CmpNat i j ~ 'LT +-- type i <= j = (i <=? j) ~ 'True + +type family Product xs where + Product '[] = 1 + Product (x ': xs) = x * Product xs + +type family Sum xs where + Sum '[] = 0 + Sum (x ': xs) = x + Sum xs + + +type family (++) xs ys where + '[] ++ xs = xs + (x ': xs) ++ ys = x ': (xs ++ ys) + +type family Tail xs where + Tail (x ': xs) = xs + +type family Last xs where + Last '[x] = x + Last (x ': xs) = Last xs + +type family Init xs where + Init '[x] = '[] + Init (x ': xs) = x ': Init xs + +-- Some proofs. + +-- initLast' :: forall s k. ((Init s ++ '[Last s]) ~ s => k) -> k +-- initLast' k = unsafeCoerce# k -- why not? + +plusAssoc' :: forall x y z. (x + y) + z :~: x + (y + z) +plusAssoc' = unsafeCoerce Refl + +plusAssoc :: forall x y z k. (((x + y) + z) ~ (x + (y + z)) => k) -> k +plusAssoc k = case plusAssoc' @x @y @z of + Refl -> k + +prodAssoc' :: forall x y z. (x * y) * z :~: x * (y * z) +prodAssoc' = unsafeCoerce Refl + +prodAssoc :: forall x y z k. (((x * y) * z) ~ (x * (y * z)) => k) -> k +prodAssoc k = case prodAssoc' @x @y @z of + Refl -> k + +prodHomo' :: forall x y. Product (x ++ y) :~: Product x * Product y +prodHomo' = unsafeCoerce Refl + +prodHomo :: forall x y k. ((Product (x ++ y) ~ (Product x * Product y)) => k) -> k +prodHomo k = case prodHomo' @x @y of Refl -> k + +knownProduct' :: forall s k. All KnownNat s => SList s -> (KnownNat (Product s) => k) -> k +knownProduct' LZ k = k +knownProduct' (LS _ n) k = knownProduct' n k + +knownProduct :: forall s k. KnownShape s => (KnownNat (Product s) => k) -> k +knownProduct = knownProduct' @s shapeSList + +initLast' :: forall s k. SList s -> ((Init s ++ '[Last s]) ~ s => k) -> k +initLast' LZ _ = error "initLast': does not hold on empty lists" +initLast' (LS _ LZ) k = k +initLast' (LS _ (LS y ys)) k = initLast' (LS y ys) k + +initLast :: forall s k. KnownShape s => ((Init s ++ '[Last s]) ~ s => k) -> k +initLast = initLast' @s shapeSList + +knownLast' :: All KnownNat s => SList s -> (KnownNat (Last s) => k) -> k +knownLast' LZ _ = error "knownLast: does not hold on empty lists" +knownLast' (LS _ LZ) k = k +knownLast' (LS _ (LS y xs)) k = knownLast' (LS y xs) k + +knownLast :: forall s k. KnownShape s => (KnownNat (Last s) => k) -> k +knownLast = knownLast' @s shapeSList + +splitApp' :: forall ys xs k. SList xs -> ((Take (PeanoLength xs) (xs ++ ys) ~ xs, + Drop (PeanoLength xs) (xs ++ ys) ~ ys) => k) -> k +splitApp' LZ k = k +splitApp' (LS _ n) k = splitApp' @ys n k + +splitApp :: forall xs ys k. KnownLen xs => ((Take (PeanoLength xs) (xs ++ ys) ~ xs, + Drop (PeanoLength xs) (xs ++ ys) ~ ys) => k) -> k +splitApp = splitApp' @ys (shapeSList @xs) + +knownAppend' :: forall t s k. (All KnownNat s, KnownShape t) => SList s -> (KnownShape (s ++ t) => k) -> k +knownAppend' LZ k = k +knownAppend' (LS _ n) k = knownAppend' @t n k + +knownAppend :: forall s t k. (KnownShape s, KnownShape t) => (KnownShape (s ++ t) => k) -> k +knownAppend = knownAppend' @t (shapeSList @s) + +-- knownCons :: proxy x -> SList xs -> (KnownLen (x ': xs) => k) -> k +-- knownCons _ LZ k = k +-- knownCons _ (LS x n) k = knownCons x n k + +-- knownFmap' :: forall f xs. SList xs -> SList (Ap (FMap f) xs) +-- knownFmap' LZ = LZ +-- knownFmap' (LS x n) = LS Proxy (knownFmap' @f n) + +-- knownSList :: SList xs -> (KnownLen xs => k) -> k +-- knownSList LZ k = k +-- knownSList (LS _ n) k = knownSList n k + +type family Length xs where + Length '[] = 0 + Length (x ': xs) = 1 + Length xs + +type family Reverse' xs ys where + Reverse' '[] ys = ys + Reverse' (x ': xs) ys = Reverse' xs (x ': ys ) + +type family Reverse xs where + Reverse xs = Reverse' xs '[] + +newtype V (n::Nat) a = V [a] + deriving (Functor, Foldable, Traversable) + +instance KnownNat n => Applicative (V n) where + pure = V . replicate (fromIntegral (natVal (Proxy @n))) + V fs <*> V xs = V (zipWith ($) fs xs) + +-- From: https://www.cs.ox.ac.uk/projects/utgp/school/andres.pdf +data NP f (xs :: [k]) where + Unit :: NP f '[] + (:*) :: f x -> NP f xs -> NP f (x ': xs) + +newtype I a = I a +newtype K a x = K a +type HList = NP I + +pattern HSingle :: f a -> NP f '[a] +pattern HSingle x = x :* Unit + +pattern VecSing :: Tensor s t -> HTV t '[s] +pattern VecSing t1 = F t1 :* Unit + +pattern VecPair :: Tensor s t -> Tensor s' t -> HTV t '[s,s'] +pattern VecPair t1 t2 = F t1 :* F t2 :* Unit + +pattern VecTriple :: Tensor s t -> Tensor s' t -> Tensor s3 t -> HTV t '[s,s',s3] +pattern VecTriple t1 t2 t3 = F t1 :* F t2 :* F t3 :* Unit + +type family All (c :: k -> Constraint) (xs :: [k]) :: Constraint where + All c '[] = () + All c (x ': xs) = (c x, All c xs) + +class Fun (c :: k -> Constraint) where + type Ap c (t :: k) :: l + +class Cons (x :: k) (xs :: [k]) +instance Fun (Cons x) where type Ap (Cons x) xs = x ': xs + +class Snoc (x :: k) (xs :: [k]) +instance Fun (Snoc x) where + type Ap (Snoc x) '[] = '[x] + type Ap (Snoc x) (y ': ys) = y ': Ap (Snoc x) ys + +class FMap (c :: k -> Constraint) (xs :: [k]) where + +instance Fun c => Fun (FMap c) where + type Ap (FMap c) '[] = '[] + type Ap (FMap c) (x ': xs) = Ap c x ': Ap (FMap c) xs + +-- type family All2 (c :: k -> l -> Constraint) (xs :: [k]) (ys :: [l]) :: Constraint where +-- All2 c '[] '[] = () +-- All2 c (x ': xs) (y ': ys) = (c x y, All2 c xs ys) +-- All2 c '[] (y ': ys) = 'True ~ 'False +-- All2 c (y ': ys) '[] = 'True ~ 'False + +-- | Flip at type level +newtype F g t s = F {fromF :: g s t} + +-- | Heterogeneous tensor vector with the same kind of elements +type HTV t = NP (F T t) + +data Pair a b = a :& b + +type family Fst (x :: Pair a b) where Fst (x ':& y) = x +type family Snd (x :: Pair a b) where Snd (x ':& y) = y + +newtype Uncurry g (s :: Pair a b) = Uncurry {fromUncurry :: g (Fst s) (Snd s)} + +type HHTV = NP (Uncurry T) + +hhead :: NP f (x ': xs) -> f x +hhead (x :* _) = x + +htail :: NP f (x ': xs) -> NP f xs +htail (_ :* xs) = xs + +htmap :: forall f ss t u. (forall s. Tensor s t -> Tensor (Ap f s) u) -> HTV t ss -> HTV u (Ap (FMap f) ss) +htmap _ Unit = Unit +htmap f (F x :* xs) = F (f x) :* htmap @f f xs + +-- htmap' :: forall f ss t u. All KnownShape ss => (forall s. KnownShape s => Tensor (Ap f s) t -> Tensor s u) -> SList ss -> HTV t (Ap (FMap f) ss) -> HTV u ss +-- htmap' _ LZ Unit = Unit +-- htmap' f (LS _ n)(F x :* xs) = F (f x) :* htmap' @f f n xs + +hmap :: (forall x. f x -> g x) -> NP f xs -> NP g xs +hmap _ Unit = Unit +hmap f (x :* xs) = f x :* hmap f xs + +hendo :: NP Endo xs -> HList xs -> HList xs +hendo Unit Unit = Unit +hendo (Endo f :* fs) (I x :* xs) = (I (f x) :* hendo fs xs) + +happ :: NP f xs -> NP f ys -> NP f (xs ++ ys) +happ Unit xs = xs +happ (x :* xs) ys = x :* (happ xs ys) + +data Both f g x = Both (f x) (g x) + +hzip :: NP f xs -> NP g xs -> NP (Both f g) xs +hzip = hzipWith Both + +hzipWith :: (forall x. f x -> g x -> h x) -> NP f xs -> NP g xs -> NP h xs +hzipWith _ Unit Unit = Unit +hzipWith f (x :* xs) (y :* ys) = f x y :* hzipWith f xs ys + +hfor_ :: Monad m => NP f xs -> (forall x. f x -> m a) -> m () +hfor_ Unit _ = return () +hfor_ (x :* xs) f = f x >> hfor_ xs f + +htoList :: NP (K a) xs -> [a] +htoList Unit = [] +htoList (K x :* xs) = x : htoList xs + +hsplit' :: SPeano n -> NP f xs -> (NP f (Take n xs), NP f (Drop n xs)) +hsplit' SZero xs = (Unit,xs) +hsplit' (SSucc _n) Unit = (Unit,Unit) +hsplit' (SSucc n) (x :* xs) = case hsplit' n xs of + (l,r) -> (x :* l,r) + +hsplit :: forall xs ys f. KnownLen xs => NP f (xs++ys) -> (NP f xs, NP f ys) +hsplit xys = splitApp @xs @ys (hsplit' (shapePeano @xs) xys) + +hsnoc :: NP f xs -> f x -> NP f (xs ++ '[x]) +hsnoc xs x = happ xs (x :* Unit) + +infixr 5 :* + +data Peano = Zero | Succ Peano + +type Dim0 = 'Zero +type Dim1 = 'Succ Dim0 +type Dim2 = 'Succ Dim1 +type Dim3 = 'Succ Dim2 + +type Axis0 = 'Zero +type Axis1 = 'Succ Dim0 +type Axis2 = 'Succ Dim1 +type Axis3 = 'Succ Dim2 + +class KnownPeano n where peanoInt :: Integer +instance KnownPeano 'Zero where peanoInt = 0 +instance KnownPeano n => KnownPeano ('Succ n) where peanoInt = 1 + (peanoInt @n) + +data SPeano n where + SZero :: SPeano 'Zero + SSucc :: SPeano n -> SPeano ('Succ n) + +data Vec (n::Peano) a where + VNil :: Vec 'Zero a + VCons :: a -> Vec n a -> Vec ('Succ n) a + +vecToList :: Vec n a -> [a] +vecToList VNil = [] +vecToList (VCons x xs) = x : vecToList xs + +-- type family App n (xs :: Vec n a) ys where +-- App 'Zero 'VNil xs = xs +-- App ('Succ n) ('VCons x xs) ys = x ': App n xs ys + +type family Take n xs where + Take 'Zero xs = '[] + Take ('Succ n) '[] = '[] + Take ('Succ n) (x ': xs) = x ': Take n xs + +type family Drop n xs where + Drop 'Zero xs = xs + Drop ('Succ n) '[] = '[] + Drop ('Succ n) (x ': xs) = Drop n xs + +type family At n xs where + At 'Zero (x ': xs) = x + At ('Succ n) (x ': xs) = At n xs + +data Kind = Float | Int | Bool deriving Show +data NBits = B32 | B64 | B1 deriving Show +data Typ = Typ Kind NBits + +type Flt t = 'Typ 'Float t +type Float32 = 'Typ 'Float 'B32 +type Int32 = 'Typ 'Int 'B32 +type Int64 = 'Typ 'Int 'B64 +type TFBool = 'Typ 'Bool 'B1 +type Scalar t = T '[] t + +instance Show Typ where + show (Typ Bool _)= "tf.bool" + show (Typ k l) = "tf." ++ map toLower (show k) ++ drop 1 (show l) + +showTyp :: forall t. KnownTyp t => DOC +showTyp = text (show (typVal @t)) + +type Shape = [Nat] + +type UntypedExpression = DOC +data T (shape :: Shape) (t :: Typ) = T {fromTensor :: UntypedExpression} + +data SNat (n :: Nat) where + SNat :: KnownNat n => Proxy n -> SNat n + +class (KnownLen s, All KnownNat s) => KnownShape s where + +instance KnownShape '[] +instance (KnownNat x, KnownShape xs) => KnownShape (x ': xs) + +class KnownTyp t where + typVal :: Typ +class KnownBits t where + bitsVal :: NBits + +instance KnownBits 'B1 where bitsVal = B1 +instance KnownBits 'B32 where bitsVal = B32 +instance KnownBits 'B64 where bitsVal = B64 +instance (KnownBits l, KnownKind k) => KnownTyp ('Typ k l) where + typVal = Typ (kindVal @k) (bitsVal @l) + +class KnownKind t where + kindVal :: Kind + +instance KnownKind 'Bool where kindVal = Bool +instance KnownKind 'Float where kindVal = Float +instance KnownKind 'Int where kindVal = Int + +-- data SList s where +-- LZ :: SList '[] +-- LS :: forall x xs. Proxy x -> SList xs -> SList (x ': xs) + +type SList = SList' Proxy + +data SList' f s where + LZ :: SList' f '[] + LS :: forall x xs f. f x -> SList' f xs -> SList' f (x ': xs) + +type family PeanoLength xs :: Peano where + PeanoLength '[] = 'Zero + PeanoLength (x ': xs) = 'Succ (PeanoLength xs) + + +withKnownNat :: forall k. Int -> (forall (n::Nat). KnownNat n => Proxy n -> k) -> k +withKnownNat 0 f = f (Proxy @0) +withKnownNat 1 f = f (Proxy @1) +withKnownNat n f = withKnownNat (n `div` 2) (if n `mod` 2 == 0 then f2x else f2x1) + where f2x,f2x1 :: forall (n::Nat). KnownNat n => Proxy n -> k + f2x _ = f (Proxy @(n*2)) + f2x1 _ = f (Proxy @(n*2+1)) + +-- Probably a GHC bug: +-- withKnownNat'' :: forall k. Int -> (forall (n::Nat). KnownNat n => k) -> k +-- withKnownNat'' 0 f = f @0 +-- withKnownNat'' n f = withKnownNat'' (n-1) fsucc +-- where fsucc :: forall (n::Nat). KnownNat n => k +-- fsucc = f @(n+1) + +-- This also fails: +-- appProxy :: forall (n::Nat) k. KnownNat n => Proxy n -> (forall (m::Nat). KnownNat m => k) -> k +-- appProxy f _ = f @n + +-- withKnownNat :: forall k. Int -> (forall (n::Nat). KnownNat n => k) -> k +-- withKnownNat n f = withKnownNat' n (\proxy -> appProxy proxy f) + +class KnownLen s where + listLen :: Integer -- CLEAN: re + shapePeano :: SPeano (PeanoLength s) + shapeSList :: SList s + +instance KnownLen '[] where + listLen = 0 + shapePeano = SZero + shapeSList = LZ + +instance KnownLen xs => KnownLen (x ': xs) where + listLen = 1 Prelude.+ listLen @ xs + shapePeano = SSucc (shapePeano @xs) + shapeSList = LS Proxy (shapeSList @xs) + +shapeSListProxy :: KnownLen xs => proxy xs -> SList xs +shapeSListProxy _ = shapeSList + +shapeToList' :: All KnownNat s => SList s -> [Integer] +shapeToList' LZ = [] +shapeToList' (LS x xs) = natVal x : shapeToList' xs + +shapeToList :: ∀(s::Shape). KnownShape s => [Integer] +shapeToList = shapeToList' (shapeSList @ s) + +showShape' :: [Integer] -> DOC +showShape' s = list (map (showDim' "None") (reverse s)) + +showShape :: ∀ (s :: Shape). KnownShape s => DOC +showShape = showShape' (shapeToList @s) + +-- | Show a shape, but "None" is replaced by "-1" +showShapeMinus :: ∀ (s :: Shape). KnownShape s => DOC +showShapeMinus = list (map (showDim' "-1") (reverse (shapeToList @ s))) + +showShapeLen :: ∀ (s::Shape). KnownLen s => DOC +showShapeLen = (text . show) (listLen @ s) + +rememberNat :: SNat n -> (KnownNat n => r) -> r +rememberNat (SNat _) k = k + +type None = 514229 -- fibonnaci prime. +-- type None = 0 - 1 -- GHC does not like negative Nats. +-- Using a maybe type would be a RPITA. + +showDim' :: String -> Integer -> DOC +showDim' none n = text (if n == 514229 then none else show n) + +showDimM :: forall n. KnownNat n => DOC +showDimM = showDim' "-1" (natVal (Proxy @ n)) + +showDim :: forall n. KnownNat n => DOC +showDim = showDim' "None" (natVal (Proxy @ n)) + +str :: Show a => a -> DOC +str = text . show + +-------------------------------- +-- Generation Effects + +data ParamInfo = ParamInfo {paramName :: String + ,paramShape :: [Integer] + ,paramDType :: Typ + ,paramVar :: forall s t. (KnownShape s, KnownTyp t) => Tensor s t} +data GState = GState {nextVar :: Integer, -- ^ next free variable + genText :: DOC, + genParams :: [ParamInfo], -- ^ optimizable parameters + genTrainingPlaceholder :: Scalar TFBool, -- ^ flag which is true when training + genPeeks :: [(String,UntypedExpression)]} +newtype Gen x = Gen {fromGen :: State GState x} deriving (Monad, MonadState GState, Functor, Applicative) + +newParameter :: MonadState GState m => ParamInfo -> m () +newParameter p = modify $ \GState{..} -> GState{genParams = p:genParams,..} + + +-- | Name an expression so that it is made available for session.run. +peekAtAny :: String -> UntypedExpression -> Gen () +peekAtAny p v = modify $ \GState{..} -> GState{genPeeks = if p `elem` map fst genPeeks then error ("duplicate name: " ++ p) else (p,v):genPeeks,..} + + +newVar :: Gen DOC +newVar = do + n <- gets nextVar + modify $ \GState{..} -> GState {nextVar=nextVar+1,..} + return (text "var" <> integer n) + +gen :: DOC -> Gen () +gen s = modify $ \GState{..} -> GState {genText=genText $$ s,..} + +setGen :: DOC -> Gen () +setGen d = modify $ \GState{..} -> GState {genText=d,..} + +withDOC :: forall a. (DOC -> DOC) -> Gen a -> Gen a +withDOC f g = do + before <- gets genText + setGen mempty + x <- g + after <- gets genText + setGen (before $$ f after) + return x + +type Tensor shape = T shape + +----------------------------------------- +-- Generation helpers + + +(<--) :: DOC -> UntypedExpression -> Gen () +x <-- y = gen (x <> text "=" <> y) + +tuple :: [DOC] -> DOC +tuple = parens . sep . punctuate comma + +dict :: [(String,DOC)] -> DOC +dict xs = encloseSep "{" "}" "," [text (show k) <> ":" <> v | (k,v) <- xs] + +funcall :: String -> [DOC] -> DOC +funcall = funcall' . text + +funcall' :: DOC -> [DOC] -> DOC +funcall' f args = hangWith "" 2 (f <> "(") (as <> ")") + where as = sep (punctuate comma args) + +binOp :: ∀ s1 s2 s3 t1 t2 t3. String -> Tensor s1 t1 -> Tensor s2 t2 -> Tensor s3 t3 +binOp op (T x) (T y) = T (funcall op [ x , y]) + +unOp :: ∀ s1 s2 t1 t2. String -> Tensor s1 t1 -> Tensor s2 t2 +unOp op (T x) = T (funcall op [x]) + +assign :: ∀s t. T s t -> Gen (T s t) +assign (T x) = do + v <- newVar + v <-- x + return (T v) + +genFun :: forall b. String -> [DOC] -> Gen b -> Gen b +genFun name args body = do + gen (text "def " <> text name <> tuple args <> text ":") + withDOC (\b -> text " " <> b) body + +lambda :: (T s t -> T s' t') -> Gen UntypedExpression +lambda f = do + v <- newVar + let T body = f (T v) + return (text "lambda " <> v <> ": " <> body) + +generate :: Gen () -> (String,[ParamInfo]) +generate s = (renderWith (Options 92 (const id)) genText,genParams) + where GState{..} = execState (fromGen s) (GState {nextVar = 0 + ,genText = mempty + ,genParams=[] + ,genTrainingPlaceholder = T "NO TRAINING PLACEHOLDER!" + ,genPeeks=[]}) + +generateFile :: String -> Gen () -> IO () +generateFile fname g = do + putStrLn ("Parameters (total " ++ show (sum [product paramShape | ParamInfo{..} <- params]) ++ "):") + forM_ params printParam + writeFile fname output + where (output,params) = generate g + printParam ParamInfo{..} = putStrLn (paramName ++ ": " ++ "T " ++ render (showShape' paramShape) ++ " " ++ show paramDType) + +named :: String -> DOC -> DOC +named fname x = text (fname <> "=") <> x + + diff --git a/typedflow.cabal b/typedflow.cabal new file mode 100644 --- /dev/null +++ b/typedflow.cabal @@ -0,0 +1,39 @@ +name: typedflow +version: 0.9 +category: Deep Learning +synopsis: Typed frontend to TensorFlow and higher-order deep learning +description: TypedFlow is a typed, higher-order frontend to TensorFlow and a high-level library for deep-learning. + . + The main design principles are: + . + - To make the parameters of layers explicit. This choice makes sharing of parameters explicit and allows to implement "layers" as pure functions. + . + - To provide as precise as possible types. Functions are explicit about the shapes and elements of the tensors that they manipulate (they are often polymorphic in shapes and elements though.) + . + - To let combinators be as transparent as possible. If a NN layers is a simple tensor transformation it will be exposed as such. +license: LGPL-3 +license-file: LICENSE +author: Jean-Philippe Bernardy +maintainer: jean-philippe.bernardy@gu.se +Cabal-Version: >= 1.12 +build-type: Simple +source-repository head + type: git + location: git@github.com:GU-CLASP/TypedFlow.git + +library + default-language: Haskell2010 + build-depends: + base==4.*, + ghc-typelits-knownnat, + pretty-compact, + mtl + + exposed-modules: + TypedFlow, + TypedFlow.Layers, + TypedFlow.Layers.Core, + TypedFlow.Layers.RNN, + TypedFlow.Learn, + TypedFlow.TF, + TypedFlow.Types