tensorflow-haskell/tensorflow-ops/tests/VariableTest.hs

80 lines
2.4 KiB
Haskell

{-# LANGUAGE OverloadedLists #-}
module Main (main) where
import Control.Monad.IO.Class (liftIO)
import qualified Data.Vector.Storable as V
import TensorFlow.Core
( unScalar
, render
, run_
, runSession
, run
, withControlDependencies)
import qualified TensorFlow.Ops as Ops
import TensorFlow.Variable
( readValue
, initializedVariable
, assign
, assignAdd
, variable
)
import Test.Framework (defaultMain, Test)
import Test.Framework.Providers.HUnit (testCase)
import Test.HUnit ((@=?))
main :: IO ()
main = defaultMain
[ testInitializedVariable
, testInitializedVariableShape
, testDependency
, testRereadRef
, testAssignAdd
]
testInitializedVariable :: Test
testInitializedVariable =
testCase "testInitializedVariable" $ runSession $ do
(formula, reset) <- do
v <- initializedVariable 42
r <- assign v 24
return (1 + readValue v, r)
result <- run formula
liftIO $ 43 @=? (unScalar result :: Float)
run_ reset -- Updates v to a different value
rerunResult <- run formula
liftIO $ 25 @=? (unScalar rerunResult :: Float)
testInitializedVariableShape :: Test
testInitializedVariableShape =
testCase "testInitializedVariableShape" $ runSession $ do
vector <- initializedVariable (Ops.constant [1] [42 :: Float])
result <- run (readValue vector)
liftIO $ [42] @=? (result :: V.Vector Float)
testDependency :: Test
testDependency =
testCase "testDependency" $ runSession $ do
v <- variable []
a <- assign v 24
r <- withControlDependencies a $ render (readValue v + 18)
result <- run r
liftIO $ (42 :: Float) @=? unScalar result
-- | See https://github.com/tensorflow/haskell/issues/92.
-- Even though we're not explicitly evaluating `f0` until the end,
-- it should hold the earlier value of the variable.
testRereadRef :: Test
testRereadRef = testCase "testReRunAssign" $ runSession $ do
w <- initializedVariable 0
f0 <- run (readValue w)
run_ =<< assign w (Ops.scalar (0.1 :: Float))
f1 <- run (readValue w)
liftIO $ (0.0, 0.1) @=? (unScalar f0, unScalar f1)
testAssignAdd :: Test
testAssignAdd = testCase "testAssignAdd" $ runSession $ do
w <- initializedVariable 42
run_ =<< assignAdd w 17
f1 <- run (readValue w)
liftIO $ (42 + 17 :: Float) @=? unScalar f1