[Embedding] Fix SharedEmbeddingColumn with PartitionedEmbedingVariable shape validation error. (#948)

Signed-off-by: candy.dc <[email protected]>

Author: candyzone

tensorflow/python/feature_column/feature_column.py

@@ -2675,6 +2675,9 @@ def create_embedding(self,
         embedding_weights = shared_embedding_collection[0]
         if isinstance(embedding_weights, kv_variable_ops.EmbeddingVariable):
           embedding_shape = (self.dimension)
+        elif isinstance(embedding_weights, variables.PartitionedVariable):
+          if isinstance(embedding_weights._get_variable_list()[0], kv_variable_ops.EmbeddingVariable):
+            embedding_shape = (self.dimension)
         if embedding_weights.get_shape() != embedding_shape:
           raise ValueError(
               'Shared embedding collection {} contains variable {} of '

tensorflow/python/feature_column/feature_column_v2_test.py

@@ -7705,6 +7705,41 @@ def testEmbeddingVariableForSharedEmbeddingColumnsMultiCol(self):
         for j in range(3):
           self.assertAlmostEqual(emb_r[i][j], emb_right[i][j])
 
+  def testEmbeddingVariableForSharedPartitionedEmbeddingColumnsMultiCol(self):
+    columns_list=[]
+    columns_list.append(fc.categorical_column_with_embedding("col_emb", dtype=dtypes.string))
+    columns_list.append(fc.categorical_column_with_embedding("col_emb2", dtype=dtypes.string))
+    W = fc.shared_embedding_columns(columns_list,
+            dimension=3,
+            initializer=init_ops.ones_initializer(dtypes.float32),
+            shared_embedding_collection_name="xxxxx_shared")
+
+    ids={}
+    ids["col_emb"] = sparse_tensor.SparseTensor(indices=[[0,0],[1,0],[2,0],[3,0],[4,0]], values=["aaaa","bbbbb","ccc","4nn","5b"], dense_shape=[5, 5])
+    ids["col_emb2"] = sparse_tensor.SparseTensor(indices=[[0,0],[1,0],[2,0],[3,0],[4,0]], values=["aaaa","bbbbb","ccc","4nn","5b"], dense_shape=[5, 5])
+    with variable_scope.variable_scope("scope",partitioner=partitioned_variables.fixed_size_partitioner(4)):
+      emb = fc_old.input_layer(ids, W)
+    fun = math_ops.multiply(emb, 2.0, name='multiply')
+    loss = math_ops.reduce_sum(fun, name='reduce_sum')
+    opt = ftrl.FtrlOptimizer(0.1, l1_regularization_strength=2.0, l2_regularization_strength=0.00001)
+    g_v = opt.compute_gradients(loss)
+    train_op = opt.apply_gradients(g_v)
+    init = variables_lib.global_variables_initializer()
+
+    with self.test_session() as sess:
+      sess.run(init)
+      sess.run([emb, train_op,loss])
+      sess.run([emb, train_op,loss])
+      emb_r, _, _ = sess.run([emb, train_op,loss])
+      emb_right = [[0.7221214, 0.7221214, 0.7221214],
+       [0.7221214, 0.7221214, 0.7221214],
+       [0.7221214, 0.7221214, 0.7221214],
+       [0.7221214, 0.7221214, 0.7221214],
+       [0.7221214, 0.7221214, 0.7221214]]
+      for i in range(5):
+        for j in range(3):
+          self.assertAlmostEqual(emb_r[i][j], emb_right[i][j])
+
   @test_util.run_deprecated_v1
   def testEmbeddingVariableForSharedEmbeddingColumnsWithPartitionNum(self):
     columns_list=[]

tensorflow/python/ops/variables.py

@@ -3100,6 +3100,9 @@ def __init__(self, name, shape, dtype, variable_list, partitions):
 
     self._name = name
     self._shape = shape
+    from tensorflow.python.ops import kv_variable_ops
+    if isinstance(self._variable_list[0], kv_variable_ops.EmbeddingVariable):
+      self._shape = shape[1:]
     self._dtype = dtype
     self._partitions = partitions
     self._as_tensor = None