Fixing sparse_ops_test string comparison failure. Fixing embedding_lookup_unique to support >2d params as per description.

Illia Polosukhin · tensorflower-gardener · commit fe954af593fa · 2016-10-31T12:17:52.000-07:00
Change: 137730923
diff --git a/tensorflow/contrib/layers/python/layers/embedding_ops.py b/tensorflow/contrib/layers/python/layers/embedding_ops.py
@@ -346,31 +346,34 @@ def embedding_lookup_unique(params, ids, name=None):
   """Version of embedding_lookup that avoids duplicate lookups.
 
   This can save communication in the case of repeated ids.
-  Same interface as embedding_lookup.
+  Same interface as embedding_lookup. Except it supports multi-dimensional `ids`
+  which allows to not reshape input/output to fit gather.
 
   Args:
     params: A list of tensors with the same shape and type, or a
-      `PartitionedVariable`.
+      `PartitionedVariable`. Shape `[index, d1, d2, ...]`.
     ids: A one-dimensional `Tensor` with type `int32` or `int64` containing
-      the ids to be looked up in `params`.
+      the ids to be looked up in `params`. Shape `[ids1, ids2, ...]`.
     name: A name for this operation (optional).
 
   Returns:
-    A `Tensor` with the same type as the tensors in `params`.
+    A `Tensor` with the same type as the tensors in `params` and dimension of
+    `[ids1, ids2, d1, d2, ...]`.
 
   Raises:
     ValueError: If `params` is empty.
   """
   with ops.name_scope(name, "EmbeddingLookupUnique", [params, ids]):
-    params = ops.convert_to_tensor(params)
     ids = ops.convert_to_tensor(ids)
     shape = array_ops.shape(ids)
     ids_flat = array_ops.reshape(
         ids, math_ops.reduce_prod(shape, keep_dims=True))
     unique_ids, idx = array_ops.unique(ids_flat)
     unique_embeddings = embedding_ops.embedding_lookup(params, unique_ids)
     embeds_flat = array_ops.gather(unique_embeddings, idx)
-    embed_shape = array_ops.concat(0, [shape, [-1]])
+    embed_shape = array_ops.concat(
+        0, [shape, array_ops.shape(unique_embeddings)[1:]])
     embeds = array_ops.reshape(embeds_flat, embed_shape)
-    embeds.set_shape(ids.get_shape().concatenate(params.get_shape()[1:]))
+    embeds.set_shape(ids.get_shape().concatenate(
+        unique_embeddings.get_shape()[1:]))
     return embeds
diff --git a/tensorflow/contrib/layers/python/layers/embedding_ops_test.py b/tensorflow/contrib/layers/python/layers/embedding_ops_test.py
@@ -371,6 +371,28 @@ def test_embedding_lookup_unique(self):
     self.assertEqual(embedded_np.shape, embedded_tf.shape)
     np.testing.assert_almost_equal(embedded_np, embedded_tf)
 
+  def test_embedding_lookup_unique_param3d(self):
+    embeds = np.random.randn(5, 3, 3)
+    idx = np.random.randint(0, 5, 10)
+    idx2d = np.random.randint(0, 5, (10, 2))
+
+    with self.test_session():
+      embedded_np = embeds[idx]
+      embedded_np2d = embeds[idx2d]
+      embedded_tf = tf.contrib.layers.embedding_lookup_unique(
+          embeds, idx).eval()
+      embedded_tf_lst = tf.contrib.layers.embedding_lookup_unique(
+          [embeds], idx).eval()
+      embedded_tf2d = tf.contrib.layers.embedding_lookup_unique(
+          embeds, idx2d).eval()
+
+    self.assertEqual(embedded_np.shape, embedded_tf.shape)
+    np.testing.assert_almost_equal(embedded_np, embedded_tf)
+    self.assertEqual(embedded_np.shape, embedded_tf_lst.shape)
+    np.testing.assert_almost_equal(embedded_np, embedded_tf_lst)
+    self.assertEqual(embedded_np2d.shape, embedded_tf2d.shape)
+    np.testing.assert_almost_equal(embedded_np2d, embedded_tf2d)
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/tensorflow/contrib/layers/python/ops/sparse_ops_test.py b/tensorflow/contrib/layers/python/ops/sparse_ops_test.py
@@ -61,25 +61,25 @@ def test_dense_to_sparse_tensor_1d_bool(self):
 
   def test_dense_to_sparse_tensor_1d_str(self):
     with self.test_session() as sess:
-      st = sparse_ops.dense_to_sparse_tensor(['qwe', '', 'ewq', ''])
+      st = sparse_ops.dense_to_sparse_tensor([b'qwe', b'', b'ewq', b''])
       result = sess.run(st)
     self.assertEqual(result.indices.dtype, np.int64)
     self.assertEqual(result.values.dtype, np.object)
     self.assertEqual(result.shape.dtype, np.int64)
     self.assertAllEqual([[0], [2]], result.indices)
-    self.assertAllEqual(['qwe', 'ewq'], result.values)
+    self.assertAllEqual([b'qwe', b'ewq'], result.values)
     self.assertAllEqual([4], result.shape)
 
   def test_dense_to_sparse_tensor_1d_str_special_ignore(self):
     with self.test_session() as sess:
       st = sparse_ops.dense_to_sparse_tensor(
-          ['qwe', '', 'ewq', ''], ignore_value='qwe')
+          [b'qwe', b'', b'ewq', b''], ignore_value=b'qwe')
       result = sess.run(st)
     self.assertEqual(result.indices.dtype, np.int64)
     self.assertEqual(result.values.dtype, np.object)
     self.assertEqual(result.shape.dtype, np.int64)
     self.assertAllEqual([[1], [2], [3]], result.indices)
-    self.assertAllEqual(['', 'ewq', ''], result.values)
+    self.assertAllEqual([b'', b'ewq', b''], result.values)
     self.assertAllEqual([4], result.shape)
 
   def test_dense_to_sparse_tensor_2d(self):
