mirror of
https://github.com/tensorflow/haskell.git
synced 2024-12-12 12:49:48 +01:00
47 lines
1.5 KiB
Haskell
47 lines
1.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 OverloadedLists #-}
|
||
|
|
{-# LANGUAGE RankNTypes #-}
|
||
|
|
|
||
|
|
module Main where
|
||
|
|
|
||
|
|
import Control.Monad.IO.Class (liftIO)
|
||
|
|
import Data.Int (Int32)
|
||
|
|
import Test.Framework.Providers.HUnit (testCase)
|
||
|
|
import Test.HUnit ((@=?))
|
||
|
|
import Google.Test
|
||
|
|
import qualified Data.Vector as V
|
||
|
|
|
||
|
|
import TensorFlow.Ops
|
||
|
|
import TensorFlow.Session
|
||
|
|
|
||
|
|
-- | Test fetching multiple outputs from an op.
|
||
|
|
testMultipleOutputs = testCase "testMultipleOutputs" $
|
||
|
|
runSession $ do
|
||
|
|
(values, indices) <- run $ topK 2 $ constant [1, 4] [10, 40, 20, 30]
|
||
|
|
liftIO $ [40, 30] @=? V.toList (values :: V.Vector Float)
|
||
|
|
liftIO $ [1, 3] @=? V.toList (indices :: V.Vector Int32)
|
||
|
|
|
||
|
|
-- | Test op with variable number of inputs.
|
||
|
|
testVarargs = testCase "testVarargs" $
|
||
|
|
runSession $ do
|
||
|
|
xs <- run $ pack $ map scalar [1..8]
|
||
|
|
liftIO $ [1..8] @=? V.toList (xs :: V.Vector Float)
|
||
|
|
|
||
|
|
main :: IO ()
|
||
|
|
main = googleTest [ testMultipleOutputs
|
||
|
|
, testVarargs
|
||
|
|
]
|