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
        ]