content: -- 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)
         ]
  