mirror of
https://github.com/tensorflow/haskell.git
synced 2024-06-03 03:23:37 +02:00
d62c614695
Distinguish between "rendered" and "unrendered" Tensors. There are now three types of `Tensor`: - `Tensor Value a`: rendered value - `Tensor Ref a`: rendered reference - `Tensor Build a` : unrendered value The extra bookkeeping makes it easier to track (and enforce) which tensors are rendered or not. For examples where this has been confusing in the past, see With this change, pure ops look similar to before, returning `Tensor Build` instead of `Tensor Value`. "Stateful" (monadic) ops are unchanged. For example: add :: OneOf [..] t => Tensor v'1 t -> Tensor v'2 t -> Tensor Build t assign :: (MonadBuild m, TensorType t) => Tensor Ref t -> Tensor v'2 t -> m (Tensor Ref t) The `gradients` function now requires that the variables over which it's differentiating are pre-rendered: gradients :: (..., Rendered v2) => Tensor v1 a -> [Tensor v2 a] -> m [Tensor Value a] (`Rendered v2` means that `v2` is either a `Ref` or a `Value`.) Additionally, the implementation of `gradients` now takes care to render every intermediate value when performing the reverse accumulation. I suspect this fixes an exponential blowup for complicated expressions.
141 lines
5.3 KiB
Haskell
141 lines
5.3 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.Monoid ((<>))
|
|
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 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
|