1
0
Fork 0
mirror of https://github.com/tensorflow/haskell.git synced 2024-11-19 01:19:43 +01:00
tensorflow-haskell/tensorflow-queue/tests/QueueTest.hs
Judah Jacobson 2c5c879037 Introduce a MonadBuild class, and remove buildAnd. (#83)
This change adds a class that both `Build` and `Session` are instances of:

    class MonadBuild m where
        build :: Build a -> m a

All stateful ops (generated and manually written) now have a signature that returns
an instance of `MonadBuild` (rather than just `Build`).  For example:

    assign_ :: (MonadBuild m, TensorType t)
            => Tensor Ref t -> Tensor v t -> m (Tensor Ref t)

This lets us remove a bunch of spurious calls to `build` in user code.  It also
lets us replace the pattern `buildAnd run foo` with the simpler pattern `foo >>= run`
(or `run =<< foo`, which is sometimes nicer when foo is a complicated expression).

I went ahead and deleted `buildAnd` altogether since it seems to lead to
confusion; in particular a few tests had `buildAnd run . pure` which is
actually equivalent to just `run`.
2017-03-18 12:08:53 -07:00

91 lines
3.1 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 OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Main where
import Control.Monad.IO.Class (liftIO)
import Data.Int (Int64)
import Google.Test (googleTest)
import TensorFlow.Types (ListOf(..), Scalar(..), (/:/))
import TensorFlow.Ops (scalar)
import TensorFlow.Queue
import TensorFlow.Session
( asyncProdNodes
, build
, run
, runSession
, run_
)
import Test.Framework (Test)
import Test.Framework.Providers.HUnit (testCase)
import Test.HUnit ((@=?))
import qualified Data.ByteString as BS
-- | Test basic queue behaviors.
testBasic :: Test
testBasic = testCase "testBasic" $ runSession $ do
q :: Queue [Int64, BS.ByteString] <- build $ makeQueue 1 ""
run_ =<< enqueue q (42 :/ scalar "Hi" :/ Nil)
x <- run =<< dequeue q
liftIO $ (Scalar 42 /:/ Scalar "Hi" /:/ Nil) @=? x
run_ =<< enqueue q (56 :/ scalar "Bar" :/ Nil)
y <- run =<< dequeue q
-- Note: we use explicit "Scalar" here to specify the type that was
-- fetched. Equivalently we could write
-- 56 /:/ "Bar" /:/ Nil :: List [Scalar Int64, Scalar BS.ByteString]
-- or else allow the types to be determined by future use of the fetched
-- value.
let expected = Scalar 56 /:/ Scalar "Bar" /:/ Nil
liftIO $ expected @=? y
-- | Test queue pumping.
testPump :: Test
testPump = testCase "testPump" $ runSession $ do
(deq, pump) <- build $ do
q :: Queue [Int64, BS.ByteString] <- makeQueue 2 "ThePumpQueue"
(,) <$> dequeue q
<*> enqueue q (31 :/ scalar "Baz" :/ Nil)
-- This is a realistic use. The pump inputs are pre-bound to some
-- nodes that produce values when pumped (e.g. read from a
-- file).
run_ (pump, pump)
(x, y) <- run (deq, deq)
let expected = Scalar 31 /:/ Scalar "Baz" /:/ Nil
liftIO $ expected @=? x
liftIO $ expected @=? y
testAsync :: Test
testAsync = testCase "testAsync" $ runSession $ do
(deq, pump) <- do
q :: Queue [Int64, BS.ByteString] <- makeQueue 2 ""
(,) <$> dequeue q
<*> enqueue q (10 :/ scalar "Async" :/ Nil)
-- Pumps the queue until canceled by runSession exiting.
asyncProdNodes pump
-- Picks up a couple values and verifies they are as expected.
let expected = Scalar 10 /:/ Scalar "Async" /:/ Nil
run deq >>= liftIO . (expected @=?)
run deq >>= liftIO . (expected @=?)
main :: IO ()
main = googleTest [ testBasic
, testPump
, testAsync
]