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 RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  -- | Tests for EmbeddingOps.
  module Main where
  
  import Control.Monad.IO.Class (liftIO)
  import Data.Int (Int32, Int64)
  import Data.List (genericLength)
  import TensorFlow.EmbeddingOps (embeddingLookup)
  import Test.Framework (defaultMain, Test)
  import Test.Framework.Providers.QuickCheck2 (testProperty)
  import Test.HUnit ((@=?))
  import Test.Framework.Providers.HUnit (testCase)
  import Test.QuickCheck (Arbitrary(..), Property, choose, vectorOf)
  import Test.QuickCheck.Monadic (monadicIO, run)
  import TensorFlow.Test (assertAllClose)
  
  import qualified Data.Vector as V
  import qualified TensorFlow.GenOps.Core as CoreOps
  import qualified TensorFlow.Ops as TF
  import qualified TensorFlow.Session as TF
  import qualified TensorFlow.Tensor as TF
  import qualified TensorFlow.Types as TF
  import qualified TensorFlow.Gradient as TF
  import qualified TensorFlow.Build as TF
  
  
  -- | Tries to perform a simple embedding lookup, with two partitions.
  testEmbeddingLookupHasRightShapeWithPartition :: Test
  testEmbeddingLookupHasRightShapeWithPartition =
          testCase "testEmbeddingLookupHasRightShapeWithPartition" $ do
      let embShape     = TF.Shape [1, 3] -- Consider a 3-dim embedding of two items.
      let embedding1  = [1, 1, 1 :: Int32]
      let embedding2  = [0, 0, 0 :: Int32]
  
      let idValues  = [0, 1 :: Int32]
  
      (values, shape) <- TF.runSession $ do
          embedding   <- mapM TF.render [ TF.constant embShape embedding1
                          , TF.constant embShape embedding2
                          ]
          let ids     = TF.constant (TF.Shape [1, 2]) idValues
          vs          <- embeddingLookup embedding ids
          TF.run (vs, TF.shape vs)
  
      -- This is the shape that is returned in the equiv. Python.
      shape  @=? V.fromList [1, 2, 3]
  
      -- "[0, 1]" should pull out the resulting vector.
      values @=? V.fromList [1, 1, 1, 0, 0, 0]
  
  
  -- | Tries to perform a simple embedding lookup, with only a single partition.
  testEmbeddingLookupHasRightShape :: Test
  testEmbeddingLookupHasRightShape =
          testCase "testEmbeddingLookupHasRightShape" $ do
      -- Consider a 3-dim embedding of two items
      let embShape      = TF.Shape [2, 3]
      let embeddingInit = [ 1, 1, 1
                          , 0, 0, 0 :: Int32
                          ]
  
      let idValues  = [0, 1 :: Int32]
  
      (values, shape) <- TF.runSession $ do
          embedding <- TF.render $ TF.constant embShape embeddingInit
          let ids = TF.constant (TF.Shape [1, 2]) idValues
          vs <- embeddingLookup [embedding] ids
          TF.run (vs, TF.shape vs)
  
      -- This is the shape that is returned in the equiv. Python.
      shape  @=? V.fromList [1, 2, 3]
  
      -- "[0, 1]" should pull out the resulting vector.
      values @=? V.fromList [1, 1, 1, 0, 0, 0]
  
  -- | Check that we can calculate gradients w.r.t embeddings.
  testEmbeddingLookupGradients :: Test
  testEmbeddingLookupGradients = testCase "testEmbeddingLookupGradients" $ do
      -- Agrees with "embedding", so gradient should be zero.
      let xVals = V.fromList ([20, 20 :: Float])
      let shape = TF.Shape [2]
  
      gs <- TF.runSession $ do
              let embShape      = TF.Shape [2, 1]
              let embeddingInit = [1, 20 ::Float]
              let idValues      = [1, 1 :: Int32]
              let ids           = TF.constant (TF.Shape [1, 2]) idValues
  
              x <- TF.placeholder (TF.Shape [2])
              embedding <- TF.initializedVariable
                              (TF.constant embShape embeddingInit)
  
              op <- embeddingLookup [embedding] ids
              let twoNorm = CoreOps.square $ TF.abs (op `TF.sub` x)
                  loss    = TF.mean twoNorm (TF.scalar (0 :: Int32))
  
              grad <- fmap head (TF.gradients loss [embedding])
              TF.runWithFeeds
                  [TF.feed x $ TF.encodeTensorData shape xVals]
                  grad
      -- Gradients should be zero (or close)
      assertAllClose gs (V.fromList ([0, 0 :: Float]))
  
  
  -- Verifies that direct gather is the same as dynamic split into
  -- partitions, followed by embedding lookup.
  testEmbeddingLookupUndoesSplit ::
      forall a. (TF.TensorDataType V.Vector a, Show a, Eq a)
      => LookupExample a -> Property
  testEmbeddingLookupUndoesSplit
      (LookupExample numParts
                     shape@(TF.Shape (firstDim : restDims))
                     values
                     indices) = monadicIO $ run $ TF.runSession $ do
      let shapedValues = TF.constant shape values
      indicesVector <- TF.render $ TF.vector indices
      let directs = CoreOps.gather shapedValues indicesVector
      let cyclicCounter :: TF.Tensor TF.Build Int32 =
              TF.vector [0..fromIntegral firstDim-1]
              `CoreOps.mod` fromIntegral numParts
      modShardedValues :: [TF.Tensor TF.Value a] <-
              mapM TF.render $ CoreOps.dynamicPartition numParts shapedValues cyclicCounter
      embeddings <- embeddingLookup modShardedValues indicesVector
      (shapeOut, got, want :: V.Vector a) <-
              TF.run (TF.cast (TF.shape embeddings), embeddings, directs)
      -- Checks the explicitly documented invariant of embeddingLookup.
      liftIO $ shapeOut @=? V.fromList (genericLength indices : restDims)
      liftIO $ got @=? want
  testEmbeddingLookupUndoesSplit _ = error "Bug in Arbitrary (LookupExample)"
  
  -- | Consistent set of parameters for EmbeddingLookupUndoesSplit.
  data LookupExample a = LookupExample
                         Int64  -- ^ number of ways to split.
                         TF.Shape  -- ^ shape of the generated tensor
                         [a]       -- ^ data for the tensor
                         [Int32]   -- ^ indices to split the tensor by
      deriving Show
  
  instance Arbitrary a => Arbitrary (LookupExample a) where
      arbitrary = do
          rank <- choose (1, 4)
          -- Takes rank-th root of 100 to cap the tensor size.
          let maxDim = fromIntegral (ceiling doubleMaxDim :: Int64)
              doubleMaxDim :: Double
              doubleMaxDim = 100 ** (1 / fromIntegral rank)
          shape <- vectorOf rank (choose (1, maxDim))
          let firstDim = head shape
          values <- vectorOf (fromIntegral $ product shape) arbitrary
          numParts <- choose (2, 15)
          indSize <- choose (0, fromIntegral $ firstDim - 1)
          indices <- vectorOf indSize (choose (0, fromIntegral firstDim - 1))
          return $ LookupExample numParts (TF.Shape shape) values indices
  
  main :: IO ()
  main = defaultMain
         [ testProperty "EmbeddingLookupUndoesSplit"
           (testEmbeddingLookupUndoesSplit :: LookupExample Double -> Property)
         , testEmbeddingLookupHasRightShape
         , testEmbeddingLookupHasRightShapeWithPartition
         , testEmbeddingLookupGradients
         ]
  