1
0
Fork 0
mirror of https://github.com/tensorflow/haskell.git synced 2024-11-05 18:49:41 +01:00
tensorflow-haskell/tensorflow-ops/tests/DataFlowOpsTest.hs
Judah Jacobson 64971c876a Consolidate some packages. (#111)
- Merge tensorflow-nn and tensorflow-queue into tensorflow-ops.
  They don't add extra dependencies and each contain a single module, so I
  don't think it's worth separating them at the package level.
- Remove google-shim in favor of direct use of test-framework.
2017-05-10 15:26:03 -07:00

65 lines
2.6 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 FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
import Data.Int (Int32, Int64)
import Data.List (genericLength)
import Test.Framework (defaultMain)
import Test.Framework.Providers.QuickCheck2 (testProperty)
import Test.HUnit ((@=?))
import Test.QuickCheck (Arbitrary(..), Property, choose, vectorOf)
import Test.QuickCheck.Monadic (monadicIO, run)
import qualified Data.Vector as V
import qualified TensorFlow.GenOps.Core as CoreOps
import qualified TensorFlow.Ops as TF
import qualified TensorFlow.Core as TF
-- DynamicSplit is undone with DynamicStitch to get the original input
-- back.
testDynamicPartitionStitchInverse :: forall a.
(TF.TensorDataType V.Vector a, Show a, Eq a) => StitchExample a -> Property
testDynamicPartitionStitchInverse (StitchExample numParts values partitions) =
let splitParts :: [TF.Tensor TF.Build a] =
CoreOps.dynamicPartition numParts (TF.vector values) partTensor
partTensor = TF.vector partitions
restitchIndices = CoreOps.dynamicPartition numParts
(TF.vector [0..genericLength values-1])
partTensor
-- drop (numParts - 2) from both args to expose b/27343984
restitch = CoreOps.dynamicStitch restitchIndices splitParts
in monadicIO $ run $ do
fromIntegral numParts @=? length splitParts
valuesOut <- TF.runSession $ TF.run restitch
V.fromList values @=? valuesOut
data StitchExample a = StitchExample Int64 [a] [Int32]
deriving Show
instance Arbitrary a => Arbitrary (StitchExample a) where
arbitrary = do
-- Limits the size of the vector.
size <- choose (1, 100)
values <- vectorOf size arbitrary
numParts <- choose (2, 15)
partitions <- vectorOf size (choose (0, fromIntegral numParts - 1))
return $ StitchExample numParts values partitions
main :: IO ()
main = defaultMain
[ testProperty "DynamicPartitionStitchInverse"
(testDynamicPartitionStitchInverse :: StitchExample Int64 -> Property)
]