146 lines
5.5 KiB
Haskell
146 lines
5.5 KiB
Haskell
-- Copyright 2016 TensorFlow authors.
|
|
--
|
|
-- Licensed under the Apache License, Version 2.0 (the "License");
|
|
-- you may not use this file except in compliance with the License.
|
|
-- You may obtain a copy of the License at
|
|
--
|
|
-- http://www.apache.org/licenses/LICENSE-2.0
|
|
--
|
|
-- Unless required by applicable law or agreed to in writing, software
|
|
-- distributed under the License is distributed on an "AS IS" BASIS,
|
|
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
-- See the License for the specific language governing permissions and
|
|
-- limitations under the License.
|
|
|
|
{-# LANGUAGE DataKinds #-}
|
|
{-# LANGUAGE FlexibleContexts #-}
|
|
{-# LANGUAGE FlexibleInstances #-}
|
|
{-# LANGUAGE MultiParamTypeClasses #-}
|
|
{-# LANGUAGE RankNTypes #-}
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
|
{-# LANGUAGE TypeFamilies #-}
|
|
{-# LANGUAGE TypeOperators #-}
|
|
{-# LANGUAGE UndecidableInstances #-} -- For Fetchable (TensorExpr a)
|
|
module TensorFlow.Nodes where
|
|
|
|
import Control.Applicative (liftA2, liftA3)
|
|
import Data.Functor.Identity (Identity)
|
|
import Data.Map.Strict (Map)
|
|
import Data.Set (Set)
|
|
import Data.Text (Text)
|
|
import qualified Data.Map.Strict as Map
|
|
import qualified Data.Set as Set
|
|
|
|
import TensorFlow.Build
|
|
import TensorFlow.Output
|
|
import TensorFlow.Tensor
|
|
import TensorFlow.Types
|
|
import qualified TensorFlow.Internal.FFI as FFI
|
|
|
|
-- | Types that contain ops which can be run.
|
|
class Nodes t where
|
|
getNodes :: t -> Build (Set NodeName)
|
|
|
|
-- | Types that tensor representations (e.g. 'Tensor', 'ControlNode') can be
|
|
-- fetched into.
|
|
--
|
|
-- Includes collections of tensors (e.g. tuples).
|
|
class Nodes t => Fetchable t a where
|
|
getFetch :: t -> Build (Fetch a)
|
|
|
|
-- | Fetch action. Keeps track of what needs to be fetched and how to decode
|
|
-- the fetched data.
|
|
data Fetch a = Fetch
|
|
{ -- | Nodes to fetch
|
|
fetches :: Set Text
|
|
-- | Function to create an 'a' from the fetched data.
|
|
, fetchRestore :: Map Text FFI.TensorData -> a
|
|
}
|
|
|
|
instance Functor Fetch where
|
|
fmap f (Fetch fetch restore) = Fetch fetch (f . restore)
|
|
|
|
instance Applicative Fetch where
|
|
pure x = Fetch Set.empty (const x)
|
|
Fetch fetch restore <*> Fetch fetch' restore' =
|
|
Fetch (fetch <> fetch') (restore <*> restore')
|
|
|
|
nodesUnion :: (Monoid b, Traversable t, Applicative f) => t (f b) -> f b
|
|
nodesUnion = fmap (foldMap id) . sequenceA
|
|
|
|
instance (Nodes t1, Nodes t2) => Nodes (t1, t2) where
|
|
getNodes (x, y) = nodesUnion [getNodes x, getNodes y]
|
|
|
|
instance (Nodes t1, Nodes t2, Nodes t3) => Nodes (t1, t2, t3) where
|
|
getNodes (x, y, z) = nodesUnion [getNodes x, getNodes y, getNodes z]
|
|
|
|
instance (Fetchable t1 a1, Fetchable t2 a2) => Fetchable (t1, t2) (a1, a2) where
|
|
getFetch (x, y) = liftA2 (,) <$> getFetch x <*> getFetch y
|
|
|
|
instance (Fetchable t1 a1, Fetchable t2 a2, Fetchable t3 a3)
|
|
=> Fetchable (t1, t2, t3) (a1, a2, a3) where
|
|
getFetch (x, y, z) =
|
|
liftA3 (,,) <$> getFetch x <*> getFetch y <*> getFetch z
|
|
|
|
instance Nodes t => Nodes [t] where
|
|
getNodes = nodesUnion . map getNodes
|
|
|
|
instance Fetchable t a => Fetchable [t] [a] where
|
|
getFetch ts = sequenceA <$> mapM getFetch ts
|
|
|
|
instance Nodes t => Nodes (Maybe t) where
|
|
getNodes = nodesUnion . fmap getNodes
|
|
|
|
instance Fetchable t a => Fetchable (Maybe t) (Maybe a) where
|
|
getFetch = fmap sequenceA . mapM getFetch
|
|
|
|
instance Nodes ControlNode where
|
|
getNodes (ControlNode o) = pure $ Set.singleton o
|
|
|
|
-- We use the constraint @(a ~ ())@ to help with type inference. For example,
|
|
-- if @t :: ControlNode@, then this constraint ensures that @run t :: Session
|
|
-- ()@. If we used @instance Fetchable ControlNode ()@ instead, then that
|
|
-- expression would be ambiguous without explicitly specifying the return type.
|
|
instance a ~ () => Fetchable ControlNode a where
|
|
getFetch _ = return $ pure ()
|
|
|
|
instance Nodes (ListOf f '[]) where
|
|
getNodes _ = return Set.empty
|
|
|
|
instance (Nodes (f a), Nodes (ListOf f as)) => Nodes (ListOf f (a ': as)) where
|
|
getNodes (x :/ xs) = liftA2 Set.union (getNodes x) (getNodes xs)
|
|
|
|
instance l ~ List '[] => Fetchable (ListOf f '[]) l where
|
|
getFetch _ = return $ pure Nil
|
|
|
|
instance (Fetchable (f t) a, Fetchable (ListOf f ts) (List as), i ~ Identity)
|
|
=> Fetchable (ListOf f (t ': ts)) (ListOf i (a ': as)) where
|
|
getFetch (x :/ xs) = liftA2 (\y ys -> y /:/ ys) <$> getFetch x <*> getFetch xs
|
|
|
|
instance Nodes (Tensor v a) where
|
|
getNodes (Tensor o) = Set.singleton . outputNodeName <$> toBuild o
|
|
|
|
fetchTensorVector :: forall a v . (TensorType a)
|
|
=> Tensor v a -> Build (Fetch (TensorData a))
|
|
fetchTensorVector (Tensor o) = do
|
|
outputName <- encodeOutput <$> toBuild o
|
|
pure $ Fetch (Set.singleton outputName) $ \tensors ->
|
|
let tensorData = tensors Map.! outputName
|
|
expectedType = tensorType (undefined :: a)
|
|
actualType = FFI.tensorDataType tensorData
|
|
badTypeError = error $ "Bad tensor type: expected "
|
|
++ show expectedType
|
|
++ ", got "
|
|
++ show actualType
|
|
in if expectedType /= actualType
|
|
then badTypeError
|
|
else TensorData tensorData
|
|
|
|
-- The constraint "a ~ a'" means that the input/output of fetch can constrain
|
|
-- the TensorType of each other.
|
|
instance (TensorType a, a ~ a') => Fetchable (Tensor v a) (TensorData a') where
|
|
getFetch = fetchTensorVector
|
|
|
|
instance (TensorType a, TensorDataType s a, a ~ a') => Fetchable (Tensor v a) (s a') where
|
|
getFetch t = fmap decodeTensorData <$> fetchTensorVector t
|