feature_column_v2_test.py

# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ===================================================================
"""Tests for python.tpu.feature_column."""

import copy

from absl.testing import parameterized
from keras.feature_column import dense_features as df_lib
from keras.feature_column import sequence_feature_column as sfc_lib

from tensorflow.python.client import session
from tensorflow.python.feature_column import feature_column_lib as fc_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import parsing_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.platform import test
from tensorflow.python.tpu import feature_column_v2 as tpu_fc
from tensorflow.python.tpu import tpu
from tensorflow.python.tpu import tpu_function


def _initialized_session():
  sess = session.Session()
  sess.run(variables_lib.global_variables_initializer())
  sess.run(lookup_ops.tables_initializer())
  return sess


class _TestStateManager(fc_lib.StateManager):

  def __init__(self, trainable=True):
    self._all_variables = {}
    self._trainable = trainable

  def create_variable(self,
                      feature_column,
                      name,
                      shape,
                      dtype=None,
                      trainable=True,
                      use_resource=True,
                      initializer=None):
    if feature_column not in self._all_variables:
      self._all_variables[feature_column] = {}
    var_dict = self._all_variables[feature_column]
    if name in var_dict:
      return var_dict[name]
    else:
      var = variable_scope.get_variable(
          name=name,
          shape=shape,
          dtype=dtype,
          trainable=self._trainable and trainable,
          use_resource=use_resource,
          initializer=initializer)
      var_dict[name] = var
      return var

  def get_variable(self, feature_column, name):
    return self._all_variables[feature_column][name]


class EmbeddingColumnTestV2(test.TestCase, parameterized.TestCase):

  def test_defaults(self):
    categorical_column = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=3)
    embedding_dimension = 2
    embedding_column = tpu_fc.embedding_column_v2(
        categorical_column, dimension=embedding_dimension)
    # Can't test default initializer as it's a random function.
    self.assertIs(categorical_column, embedding_column.categorical_column)
    self.assertEqual(embedding_dimension, embedding_column.dimension)
    self.assertEqual('mean', embedding_column.combiner)
    self.assertEqual('aaa_embedding', embedding_column.name)
    self.assertEqual((embedding_dimension,), embedding_column.variable_shape)

  def test_all_constructor_args(self):
    categorical_column = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=3)
    embedding_dimension = 2
    embedding_column = tpu_fc.embedding_column_v2(
        categorical_column,
        dimension=embedding_dimension,
        combiner='my_combiner',
        initializer=lambda: 'my_initializer')
    self.assertIs(categorical_column, embedding_column.categorical_column)
    self.assertEqual(embedding_dimension, embedding_column.dimension)
    self.assertEqual('my_combiner', embedding_column.combiner)
    self.assertEqual('my_initializer', embedding_column.initializer())
    self.assertEqual('aaa_embedding', embedding_column.name)
    self.assertEqual((embedding_dimension,), embedding_column.variable_shape)
    self.assertEqual({
        'aaa': parsing_ops.VarLenFeature(dtypes.int64)
    }, embedding_column._parse_example_spec)

  @parameterized.named_parameters(
      {
          'testcase_name': 'use_safe_embedding_lookup',
          'use_safe_embedding_lookup': True,
      }, {
          'testcase_name': 'dont_use_safe_embedding_lookup',
          'use_safe_embedding_lookup': False,
      })
  @test_util.deprecated_graph_mode_only
  def test_feature_layer_cpu(self, use_safe_embedding_lookup):
    # Inputs.
    vocabulary_size = 3
    sparse_input = sparse_tensor.SparseTensorValue(
        # example 0, ids [2]
        # example 1, ids [0, 1]
        # example 2, ids []
        # example 3, ids [1]
        indices=((0, 0), (1, 0), (1, 1), (3, 0)),
        values=(2, 0, 1, 1),
        dense_shape=(4, 2))

    # Embedding variable.
    embedding_dimension = 2
    embedding_values = (
        (1., 2.),  # id 0
        (3., 5.),  # id 1
        (7., 11.)  # id 2
    )

    def _initializer(shape, dtype, partition_info=None):
      self.assertAllEqual((vocabulary_size, embedding_dimension), shape)
      self.assertEqual(dtypes.float32, dtype)
      self.assertIsNone(partition_info)
      return embedding_values

    # Expected lookup result, using combiner='mean'.
    expected_lookups = (
        # example 0, ids [2], embedding = [7, 11]
        (7., 11.),
        # example 1, ids [0, 1], embedding = mean([1, 2] + [3, 5]) = [2, 3.5]
        (2., 3.5),
        # example 2, ids [], embedding = [0, 0]
        (0., 0.),
        # example 3, ids [1], embedding = [3, 5]
        (3., 5.),
    )
    expected_lookups_sequence = (
        # example 0, ids [2], embedding = [[7, 11], [0, 0]]
        ((7., 11.), (0., 0.),),
        # example 1, ids [0, 1], embedding = [[1, 2], [3. 5]]
        ((1., 2.), (3., 5.),),
        # example 2, ids [], embedding = [0, 0]
        ((0., 0.), (0., 0.),),
        # example 3, ids [1], embedding = [3, 5]
        ((3., 5.), (0., 0.),),
    )

    # Build columns.
    categorical_column = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=vocabulary_size)
    sequence_categorical_column = (
        fc_lib.sequence_categorical_column_with_identity(
            key='bbb', num_buckets=vocabulary_size))
    embedding_column = tpu_fc.embedding_column_v2(
        categorical_column,
        dimension=embedding_dimension,
        initializer=_initializer,
        use_safe_embedding_lookup=use_safe_embedding_lookup)
    sequence_embedding_column = tpu_fc.embedding_column_v2(
        sequence_categorical_column,
        dimension=embedding_dimension,
        initializer=_initializer,
        max_sequence_length=2,
        use_safe_embedding_lookup=use_safe_embedding_lookup)

    # Provide sparse input and get dense result.
    features = {'aaa': sparse_input, 'bbb': sparse_input}
    dense_features = df_lib.DenseFeatures([embedding_column])
    sequence_features = sfc_lib.SequenceFeatures([sequence_embedding_column])
    embedding_lookup = dense_features(features)
    sequence_embedding_lookup = sequence_features(features)

    # Assert expected embedding variable and lookups.
    global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
    self.assertItemsEqual(
        ('dense_features/aaa_embedding/embedding_weights:0',
         'sequence_features/bbb_embedding/embedding_weights:0',),
        tuple([v.name for v in global_vars]))
    with _initialized_session():
      self.assertAllEqual(embedding_values, global_vars[0])
      self.assertAllEqual(expected_lookups, embedding_lookup)
      self.assertAllEqual(expected_lookups_sequence,
                          sequence_embedding_lookup[0].eval())
      # The graph will still have SparseFillEmptyRows due to sequence being
      # a Rank3 embedding lookup.
      if use_safe_embedding_lookup:
        self.assertEqual(2, [
            x.type for x in ops.get_default_graph().get_operations()
        ].count('SparseFillEmptyRows'))
      else:
        self.assertEqual(1, [
            x.type for x in ops.get_default_graph().get_operations()
        ].count('SparseFillEmptyRows'))

  def test_deepcopy(self):
    categorical_column = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=3)
    embedding_column = tpu_fc.embedding_column_v2(
        categorical_column, dimension=2)
    embedding_column_copy = copy.deepcopy(embedding_column)
    self.assertEqual(embedding_column.dimension,
                     embedding_column_copy.dimension)
    self.assertEqual(embedding_column._max_sequence_length,
                     embedding_column_copy._max_sequence_length)

  def test_with_scope_validation(self):
    categorical_column = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=3)
    embedding_dimension = 2
    initializer = init_ops.truncated_normal_initializer(mean=0.0, stddev=.5)
    embedding_column = tpu_fc._TPUEmbeddingColumnV2(
        categorical_column=categorical_column,
        dimension=embedding_dimension,
        combiner='mean',
        initializer=initializer,
        max_sequence_length=0,
        learning_rate_fn=None,
        use_safe_embedding_lookup=True,
        bypass_scope_validation=False)
    self.assertIs(categorical_column, embedding_column.categorical_column)
    self.assertEqual(embedding_dimension, embedding_column.dimension)
    state_manager = _TestStateManager()
    with tpu_function.tpu_shard_context(1):
      with variable_scope.variable_scope('tower1/scope1'):
        embedding_column.create_state(state_manager)
      with variable_scope.variable_scope('tower2/scope2'):
        # With default scope validation, the same column cannot be used in a new
        # variable scope.
        with self.assertRaisesRegex(ValueError,
                                    'the variable scope name is different'):
          embedding_column.create_state(state_manager)

  def test_bypass_scope_validation(self):
    categorical_column = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=3)
    embedding_dimension = 2
    initializer = init_ops.truncated_normal_initializer(mean=0.0, stddev=.5)
    embedding_column = tpu_fc._TPUEmbeddingColumnV2(
        categorical_column=categorical_column,
        dimension=embedding_dimension,
        combiner='mean',
        initializer=initializer,
        max_sequence_length=0,
        learning_rate_fn=None,
        use_safe_embedding_lookup=True,
        bypass_scope_validation=True)
    self.assertIs(categorical_column, embedding_column.categorical_column)
    self.assertEqual(embedding_dimension, embedding_column.dimension)
    state_manager = _TestStateManager()
    with tpu_function.tpu_shard_context(1):
      with variable_scope.variable_scope('tower1/scope1'):
        embedding_column.create_state(state_manager)
      with variable_scope.variable_scope('tower2/scope2'):
        embedding_column.create_state(state_manager)

  def test_deepcopy_with_bypass_scope_validation(self):
    categorical_column = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=3)
    embedding_dimension = 2
    initializer = init_ops.truncated_normal_initializer(mean=0.0, stddev=.5)
    embedding_column = tpu_fc._TPUEmbeddingColumnV2(
        categorical_column=categorical_column,
        dimension=embedding_dimension,
        combiner='mean',
        initializer=initializer,
        max_sequence_length=0,
        use_safe_embedding_lookup=False,
        bypass_scope_validation=True)
    embedding_column_copy = copy.deepcopy(embedding_column)
    self.assertEqual(embedding_dimension, embedding_column_copy.dimension)
    self.assertEqual(embedding_column._max_sequence_length,
                     embedding_column_copy._max_sequence_length)
    self.assertTrue(embedding_column_copy._bypass_scope_validation)
    self.assertFalse(embedding_column_copy.use_safe_embedding_lookup)


class SharedEmbeddingColumnTestV2(test.TestCase, parameterized.TestCase):

  @test_util.deprecated_graph_mode_only
  def test_defaults(self):
    vocabulary_size = 3
    categorical_column_a = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=vocabulary_size)
    categorical_column_b = fc_lib.categorical_column_with_identity(
        key='bbb', num_buckets=vocabulary_size)
    embedding_dimension = 2
    embedding_column_b, embedding_column_a = tpu_fc.shared_embedding_columns_v2(
        [categorical_column_b, categorical_column_a],
        dimension=embedding_dimension)
    self.assertIs(categorical_column_a, embedding_column_a.categorical_column)
    self.assertIs(categorical_column_b, embedding_column_b.categorical_column)
    self.assertEqual((vocabulary_size, embedding_dimension),
                     embedding_column_a.get_embedding_table_size())
    self.assertEqual((vocabulary_size, embedding_dimension),
                     embedding_column_a.get_embedding_table_size())
    self.assertEqual('mean', embedding_column_a.combiner)
    self.assertEqual('mean', embedding_column_b.combiner)
    self.assertIsNotNone(embedding_column_a.get_initializer())
    self.assertIsNotNone(embedding_column_b.get_initializer())
    self.assertEqual('aaa_bbb_shared_embedding',
                     embedding_column_a.get_embedding_var_name())
    self.assertEqual('aaa_bbb_shared_embedding',
                     embedding_column_b.get_embedding_var_name())
    self.assertEqual('aaa_shared_embedding', embedding_column_a.name)
    self.assertEqual('bbb_shared_embedding', embedding_column_b.name)
    self.assertEqual((embedding_dimension,), embedding_column_a.variable_shape)
    self.assertEqual((embedding_dimension,), embedding_column_b.variable_shape)

  @test_util.deprecated_graph_mode_only
  def test_all_constructor_args(self):
    vocabulary_size = 3
    categorical_column_a = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=vocabulary_size)
    categorical_column_b = fc_lib.categorical_column_with_identity(
        key='bbb', num_buckets=vocabulary_size)
    embedding_dimension = 2
    embedding_column_a, embedding_column_b = tpu_fc.shared_embedding_columns_v2(
        [categorical_column_a, categorical_column_b],
        dimension=embedding_dimension,
        combiner='my_combiner',
        initializer=lambda: 'my_initializer',
        shared_embedding_collection_name='var_scope_name')
    self.assertIs(categorical_column_a, embedding_column_a.categorical_column)
    self.assertIs(categorical_column_b, embedding_column_b.categorical_column)
    self.assertEqual((vocabulary_size, embedding_dimension),
                     embedding_column_a.get_embedding_table_size())
    self.assertEqual((vocabulary_size, embedding_dimension),
                     embedding_column_a.get_embedding_table_size())
    self.assertEqual('my_combiner', embedding_column_a.combiner)
    self.assertEqual('my_combiner', embedding_column_b.combiner)
    self.assertEqual('my_initializer', embedding_column_a.get_initializer()())
    self.assertEqual('my_initializer', embedding_column_b.get_initializer()())
    self.assertEqual('var_scope_name',
                     embedding_column_a.get_embedding_var_name())
    self.assertEqual('var_scope_name',
                     embedding_column_b.get_embedding_var_name())
    self.assertEqual('aaa_shared_embedding', embedding_column_a.name)
    self.assertEqual('bbb_shared_embedding', embedding_column_b.name)
    self.assertEqual((embedding_dimension,), embedding_column_a.variable_shape)
    self.assertEqual((embedding_dimension,), embedding_column_b.variable_shape)

  @parameterized.named_parameters(
      {
          'testcase_name': 'use_safe_embedding_lookup',
          'use_safe_embedding_lookup': True
      }, {
          'testcase_name': 'dont_use_safe_embedding_lookup',
          'use_safe_embedding_lookup': False
      })
  @test_util.deprecated_graph_mode_only
  def test_feature_layer_cpu(self, use_safe_embedding_lookup):
    # Inputs.
    vocabulary_size = 3
    input_a = sparse_tensor.SparseTensorValue(
        # example 0, ids [2]
        # example 1, ids [0, 1]
        indices=((0, 0), (1, 0), (1, 1)),
        values=(2, 0, 1),
        dense_shape=(2, 2))
    input_b = sparse_tensor.SparseTensorValue(
        # example 0, ids [2]
        # example 1, ids [0, 1]
        # example 2, ids []
        indices=((0, 0), (1, 0), (1, 1)),
        values=(2, 0, 1),
        dense_shape=(3, 2))
    input_features = {'aaa': input_a, 'bbb': input_b}

    # Embedding variable.
    embedding_dimension = 2
    embedding_values = (
        (1., 2.),  # id 0
        (3., 5.),  # id 1
        (7., 11.)  # id 2
    )

    def _initializer(shape, dtype, partition_info=None):
      self.assertAllEqual((vocabulary_size, embedding_dimension), shape)
      self.assertEqual(dtypes.float32, dtype)
      self.assertIsNone(partition_info)
      return embedding_values

    # Expected lookup result, using combiner='mean'.
    expected_lookups_a = (
        # example 0:
        (7., 11.),  # ids [2], embedding = [7, 11]
        # example 1:
        (2., 3.5),  # ids [0, 1], embedding = mean([1, 2] + [3, 5]) = [2, 3.5]
    )
    expected_lookups_b = (
        # example 0:
        ((7., 11.), (0., 0.),),  # ids [2], embedding = [[7, 11], [0, 0]]
        # example 1:
        ((1., 2.), (3., 5.),),  # ids [0, 1], embedding = [[1, 2], [3, 5]]
        # example 2:
        ((0., 0.), (0., 0.),),  # ids [], embedding = [[0, 0], [0, 0]]
    )

    # Build columns.
    categorical_column_a = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=vocabulary_size)
    categorical_column_b = fc_lib.sequence_categorical_column_with_identity(
        key='bbb', num_buckets=vocabulary_size)
    embedding_column_a, embedding_column_b = tpu_fc.shared_embedding_columns_v2(
        [categorical_column_a, categorical_column_b],
        dimension=embedding_dimension,
        initializer=_initializer,
        max_sequence_lengths=[0, 2],
        use_safe_embedding_lookup=use_safe_embedding_lookup)

    # Provide sparse input and get dense result.
    dense_features = df_lib.DenseFeatures([embedding_column_a])
    sequence_features = sfc_lib.SequenceFeatures([embedding_column_b])
    embedding_lookup_a = dense_features(input_features)
    embedding_lookup_b = sequence_features(input_features)

    # Assert expected embedding variable and lookups.
    global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
    self.assertItemsEqual(
        ('aaa_bbb_shared_embedding:0',),
        tuple([v.name for v in global_vars]))
    embedding_var = global_vars[0]
    with _initialized_session():
      self.assertAllEqual(embedding_values, embedding_var)
      self.assertAllEqual(expected_lookups_a, embedding_lookup_a)
      self.assertAllEqual(expected_lookups_b,
                          embedding_lookup_b[0].eval())
      # The graph will still have SparseFillEmptyRows due to sequence being
      # a Rank3 embedding lookup.
      if use_safe_embedding_lookup:
        self.assertEqual(2, [
            x.type for x in ops.get_default_graph().get_operations()
        ].count('SparseFillEmptyRows'))
      else:
        self.assertEqual(1, [
            x.type for x in ops.get_default_graph().get_operations()
        ].count('SparseFillEmptyRows'))

  def test_deepcopy(self):
    vocabulary_size = 3
    categorical_column_a = fc_lib.categorical_column_with_identity(
        key='aaa', num_buckets=vocabulary_size)
    categorical_column_b = fc_lib.categorical_column_with_identity(
        key='bbb', num_buckets=vocabulary_size)
    embedding_dimension = 2
    columns = tpu_fc.shared_embedding_columns_v2(
        [categorical_column_b, categorical_column_a],
        dimension=embedding_dimension)
    columns_copy = copy.deepcopy(columns)
    self.assertEqual(
        [column._shared_embedding_collection_name for column in columns],
        [column._shared_embedding_collection_name for column in columns_copy])


class DeviceSpecificEmbeddingColumnTestV2(test.TestCase,
                                          parameterized.TestCase):

  @parameterized.named_parameters(
      {
          'testcase_name': 'invalid_shared',
          'shared': True,
      }, {
          'testcase_name': 'invalid_not_shared',
          'shared': False,
      })
  @test_util.deprecated_graph_mode_only
  def test_invalid_cases(self, shared):

    # Inputs.
    input_sparse_tensor = sparse_tensor.SparseTensorValue(
        indices=((0, 0), (1, 0), (1, 1), (1, 4)),
        values=(2, 0, 1, 3),
        dense_shape=(2, 5))
    input_features = {'inp': input_sparse_tensor}

    # Build columns.
    categorical_column_input = fc_lib.categorical_column_with_identity(
        key='inp', num_buckets=3)

    # Training on TPU with cpu embedding lookups is not supported.
    if shared:
      embedding_column = tpu_fc.shared_embedding_columns_v2(
          [categorical_column_input],
          dimension=2,
          embedding_lookup_device='cpu',
          tensor_core_shape=[None, 3])
    else:
      embedding_column = tpu_fc.embedding_column_v2(
          categorical_column_input,
          dimension=2,
          embedding_lookup_device='cpu',
          tensor_core_shape=[None, 3])
    dense_features = df_lib.DenseFeatures(embedding_column)
    with self.assertRaisesRegex(
        ValueError,
        r'.*embedding_lookup_device=\"cpu\" during training is not'):
      dense_features(input_features)

    # Inference on with TPU Embedding Hardware is not supported.
    if shared:
      embedding_column = tpu_fc.shared_embedding_columns_v2(
          [categorical_column_input],
          dimension=2,
          embedding_lookup_device='tpu_embedding_core',
          tensor_core_shape=[None, 3])
    else:
      embedding_column = tpu_fc.embedding_column_v2(
          categorical_column_input,
          dimension=2,
          embedding_lookup_device='tpu_embedding_core',
          tensor_core_shape=[None, 3])
    context = tpu._TPUInferenceContext('tpu_inference')
    context.Enter()
    dense_features = df_lib.DenseFeatures(embedding_column)
    with self.assertRaisesRegex(
        ValueError,
        r'Using embedding_lookup_device=tpu_embedding_core during inference is '
    ):
      dense_features(input_features)
    context.Exit()

  @parameterized.named_parameters(
      {
          'testcase_name': 'combiner_mean_shared',
          'shared': True,
          'combiner': 'mean'
      }, {
          'testcase_name': 'combiner_sum_shared',
          'shared': True,
          'combiner': 'sum'
      }, {
          'testcase_name': 'combiner_sqrtn_shared',
          'shared': True,
          'combiner': 'sqrtn'
      }, {
          'testcase_name': 'combiner_mean_not_shared',
          'shared': False,
          'combiner': 'mean'
      }, {
          'testcase_name': 'combiner_sum_not_shared',
          'shared': False,
          'combiner': 'sum'
      }, {
          'testcase_name': 'combiner_sqrtn_not_shared',
          'shared': False,
          'combiner': 'sqrtn'
      })
  @test_util.deprecated_graph_mode_only
  def test_dense_embedding_lookup(self, shared, combiner):
    # Inputs.
    vocabulary_size = 3
    input_sparse_tensor = sparse_tensor.SparseTensorValue(
        # example 0, ids [2]
        # example 1, ids [0, 1, 3]
        indices=((0, 0), (1, 0), (1, 1), (1, 4)),
        values=(2, 0, 1, 3),
        dense_shape=(2, 5))
    input_features = {'inp': input_sparse_tensor}

    # Embedding variable.
    embedding_dimension = 2
    embedding_values = (
        (1., 2.),  # id 0
        (3., 5.),  # id 1
        (7., 11.),  # id 2
        (13., 17.)  # id 3
    )

    def _initializer(shape, dtype, partition_info=None):
      self.assertAllEqual((vocabulary_size, embedding_dimension), shape)
      self.assertEqual(dtypes.float32, dtype)
      self.assertIsNone(partition_info)
      return embedding_values

    # Build columns.
    categorical_column_input = fc_lib.categorical_column_with_identity(
        key='inp', num_buckets=vocabulary_size)

    # Set tensor_core_shape to be [None, 20] to ensure some padding and
    # dynamic batch size.
    if shared:
      embedding_column = tpu_fc.shared_embedding_columns_v2(
          [categorical_column_input],
          dimension=embedding_dimension,
          initializer=_initializer,
          combiner=combiner,
          embedding_lookup_device='tpu_tensor_core',
          tensor_core_shape=[None, 3])
    else:
      embedding_column = tpu_fc.embedding_column_v2(
          categorical_column_input,
          dimension=embedding_dimension,
          initializer=_initializer,
          combiner=combiner,
          embedding_lookup_device='tpu_tensor_core',
          tensor_core_shape=[None, 3])

    # Run in TPUContexts so that we hit the intended densification case.
    context = tpu._TPUInferenceContext('tpu_inference')
    context.Enter()
    with tpu_function.tpu_shard_context(1):
      dense_features = df_lib.DenseFeatures(embedding_column)
      # Sqrtn combiner not supported for now.
      if combiner == 'sqrtn':
        with self.assertRaisesRegex(
            ValueError, 'Dense TPU Embedding does not support combiner'):
          embedding_lookup = dense_features(input_features)
        return
      if combiner == 'mean':
        expected_lookups = (
            # example 0:
            (7., 11.),  # ids [2], embedding = [7, 11]
            # example 1:
            (2., 3.5),  # ids [0, 1], embedding = mean([1, 2] + [3, 5]) =
            # [2, 3.5]
        )
      elif combiner == 'sum':
        expected_lookups = (
            # example 0:
            (7., 11.),  # ids [2], embedding = [7, 11]
            # example 1:
            (4., 7),  # ids [0, 1], embedding = sum([1, 2] + [3, 5]) = [4, 7]
        )

      embedding_lookup = dense_features(input_features)

      # Assert expected embedding variable and lookups.
      global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
      if shared:
        self.assertCountEqual(('inp_shared_embedding:0',),
                              tuple([v.name for v in global_vars]))
      else:
        self.assertCountEqual(
            ('dense_features/inp_embedding/embedding_weights:0',),
            tuple([v.name for v in global_vars]))

      embedding_var = global_vars[0]
      with _initialized_session():
        self.assertAllEqual(embedding_values, embedding_var)
        eval_res = embedding_lookup.eval()
        self.assertAllEqual(expected_lookups, eval_res)
      context.Exit()

  @test_util.deprecated_graph_mode_only
  def test_empty_row(self):
    # Inputs.
    vocabulary_size = 3
    input_sparse_tensor = sparse_tensor.SparseTensorValue(
        # example 0, ids []
        # example 1, ids [0, 1, 3]
        indices=((1, 0), (1, 1), (1, 4)),
        values=(0, 1, 3),
        dense_shape=(2, 5))
    input_features = {'inp': input_sparse_tensor}

    # Embedding variable.
    embedding_dimension = 2
    embedding_values = (
        (1., 2.),  # id 0
        (3., 5.),  # id 1
        (7., 11.),  # id 2
        (13., 17.)  # id 3
    )

    def _initializer(shape, dtype, partition_info=None):
      self.assertAllEqual((vocabulary_size, embedding_dimension), shape)
      self.assertEqual(dtypes.float32, dtype)
      self.assertIsNone(partition_info)
      return embedding_values

    # Build columns.
    categorical_column_input = fc_lib.categorical_column_with_identity(
        key='inp', num_buckets=vocabulary_size)

    # Set tensor_core_shape to be [None, 20] to ensure some padding and
    # dynamic batch size.
    embedding_column = tpu_fc.embedding_column_v2(
        categorical_column_input,
        dimension=embedding_dimension,
        initializer=_initializer,
        combiner='mean',
        embedding_lookup_device='tpu_tensor_core',
        tensor_core_shape=[None, 3])

    # Run in TPUContexts so that we hit the intended densification case.
    context = tpu._TPUInferenceContext('tpu_inference')
    context.Enter()
    with tpu_function.tpu_shard_context(1):
      dense_features = df_lib.DenseFeatures(embedding_column)
      expected_lookups = (
          # example 0:
          (0., 0.),  # ids [], embedding = [0, 0]
          # example 1:
          (2., 3.5),  # ids [0, 1], embedding = mean([1, 2] + [3, 5]) = [2, 3.5]
      )

      embedding_lookup = dense_features(input_features)

      # Assert expected embedding variable and lookups.
      global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
      self.assertCountEqual(
          ('dense_features/inp_embedding/embedding_weights:0',),
          tuple([v.name for v in global_vars]))

      embedding_var = global_vars[0]
      with _initialized_session():
        self.assertAllEqual(embedding_values, embedding_var)
        eval_res = embedding_lookup.eval()
        self.assertAllEqual(expected_lookups, eval_res)
      context.Exit()

  @test_util.deprecated_graph_mode_only
  def test_error_dense_shape_invalid(self):
    categorical_column_input = fc_lib.categorical_column_with_identity(
        key='inp', num_buckets=5)
    with self.assertRaisesRegex(ValueError, 'tensor_core_shape must be size 2'):
      tpu_fc.shared_embedding_columns_v2([categorical_column_input],
                                         dimension=20,
                                         tensor_core_shape=[None, 20, 15])


if __name__ == '__main__':
  test.main()