From 9c81241439bd7eb8fe754abaf81b6633a61f3432 Mon Sep 17 00:00:00 2001 From: silky Date: Wed, 9 Nov 2016 10:30:05 +1100 Subject: [PATCH] Tests for "embedding_lookup" and minor fix - added a test that fails for a partitioned embedding - added a test that passes for a single embedding --- tensorflow-ops/src/TensorFlow/EmbeddingOps.hs | 9 ++-- tensorflow-ops/tests/EmbeddingOpsTest.hs | 53 +++++++++++++++++++ 2 files changed, 59 insertions(+), 3 deletions(-) diff --git a/tensorflow-ops/src/TensorFlow/EmbeddingOps.hs b/tensorflow-ops/src/TensorFlow/EmbeddingOps.hs index 9eb396b..96ee040 100644 --- a/tensorflow-ops/src/TensorFlow/EmbeddingOps.hs +++ b/tensorflow-ops/src/TensorFlow/EmbeddingOps.hs @@ -60,9 +60,12 @@ embeddingLookup :: forall a b v . -- fewer than 2^31 entries. -> Build (Tensor Value a) -- ^ A dense tensor with shape `shape(ids) + shape(params)[1:]`. -embeddingLookup params ids = - CoreOps.dynamicStitch pindices <$> partitionedResult - where np = genericLength params +embeddingLookup params@(p1 : _) ids = + go (np :: Int32) + where + go 1 = colocateWith p1 (render $ CoreOps.gather p1 ids) + go _ = CoreOps.dynamicStitch pindices <$> partitionedResult + np = genericLength params pAssignments = CoreOps.cast (ids `CoreOps.mod` np) newIds = ids `CoreOps.div` np originalIndices = CoreOps.range 0 (CoreOps.size ids) 1 diff --git a/tensorflow-ops/tests/EmbeddingOpsTest.hs b/tensorflow-ops/tests/EmbeddingOpsTest.hs index 0a6b97d..722492b 100644 --- a/tensorflow-ops/tests/EmbeddingOpsTest.hs +++ b/tensorflow-ops/tests/EmbeddingOpsTest.hs @@ -24,6 +24,7 @@ import Google.Test (googleTest) import TensorFlow.EmbeddingOps (embeddingLookup) 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) @@ -34,6 +35,56 @@ import qualified TensorFlow.Session as TF import qualified TensorFlow.Tensor as TF import qualified TensorFlow.Types as TF + +buildAndRun = TF.runSession . TF.buildAnd TF.run + +-- | Tries to perform a simple embedding lookup, with two partitions. +testEmbeddingLookupHasRightShapeWithPartition = testCase "testEmbeddingLookupHasRightShapeWithPartition" $ do + let shape = 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 embedding = [ TF.constant shape embedding1 + , TF.constant shape embedding2 + ] + + let idValues = [0, 1] :: [Int32] + let ids = TF.constant (TF.Shape [1, 2]) idValues + let op = embeddingLookup embedding ids + + (values, shape) <- buildAndRun $ do + vs <- op + return (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 = testCase "testEmbeddingLookupHasRightShape" $ do + let shape = TF.Shape [2, 3] -- Consider a 3-dim embedding of two items. + let embeddingInit = [ 1, 1, 1 + , 0, 0, 0 ] :: [Int32] + + let embedding = TF.constant shape embeddingInit + let idValues = [0, 1] :: [Int32] + let ids = TF.constant (TF.Shape [1, 2]) idValues + let op = embeddingLookup [embedding] ids + + (values, shape) <- buildAndRun $ do + vs <- op + return (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 ] + + -- Verifies that direct gather is the same as dynamic split into -- partitions, followed by embedding lookup. testEmbeddingLookupUndoesSplit :: forall a. (TF.TensorType a, Show a, Eq a) @@ -85,4 +136,6 @@ main :: IO () main = googleTest [ testProperty "EmbeddingLookupUndoesSplit" (testEmbeddingLookupUndoesSplit :: LookupExample Double -> Property) + , testEmbeddingLookupHasRightShape + , testEmbeddingLookupHasRightShapeWithPartition ]