tensorflow-haskell/tensorflow-ops/tests/EmbeddingOpsTest.hs

-- 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@(firstDim : _) <- vectorOf rank (choose (1, maxDim))
        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
       ]