code: # Copyright 2018 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 third_party.tensorflow.python.ops.ragged_tensor."""
  
  import re
  
  from absl.testing import parameterized
  
  from tensorflow.python.eager import context
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops.ragged import ragged_factory_ops
  from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
  
  from tensorflow.python.platform import googletest
  
  
  class _SliceBuilder:
    """Helper to construct arguments for __getitem__.
  
    Usage: _SliceBuilder()[<expr>] slice_spec Python generates for <expr>.
    """
  
    def __getitem__(self, slice_spec):
      return slice_spec
  
  
  SLICE_BUILDER = _SliceBuilder()
  
  
  def _make_tensor_slice_spec(slice_spec, use_constant=True):
    """Wraps all integers in an extended slice spec w/ a tensor.
  
    This function is used to help test slicing when the slice spec contains
    tensors, rather than integers.
  
    Args:
      slice_spec: The extended slice spec.
      use_constant: If true, then wrap each integer with a tf.constant.  If false,
        then wrap each integer with a tf.placeholder.
  
    Returns:
      A copy of slice_spec, but with each integer i replaced with tf.constant(i).
    """
  
    def make_piece_scalar(piece):
      if isinstance(piece, int):
        scalar = constant_op.constant(piece)
        if use_constant:
          return scalar
        else:
          return array_ops.placeholder_with_default(scalar, [])
      elif isinstance(piece, slice):
        return slice(
            make_piece_scalar(piece.start), make_piece_scalar(piece.stop),
            make_piece_scalar(piece.step))
      else:
        return piece
  
    if isinstance(slice_spec, tuple):
      return tuple(make_piece_scalar(piece) for piece in slice_spec)
    else:
      return make_piece_scalar(slice_spec)
  
  
  # Example 2D ragged tensor value with one ragged dimension and with scalar
  # values, expressed as nested python lists and as splits+values.
  EXAMPLE_RAGGED_TENSOR_2D = [[b'a', b'b'], [b'c', b'd', b'e'], [b'f'], [],
                              [b'g']]
  EXAMPLE_RAGGED_TENSOR_2D_SPLITS = [0, 2, 5, 6, 6, 7]
  EXAMPLE_RAGGED_TENSOR_2D_VALUES = ['a', 'b', 'c', 'd', 'e', 'f', 'g']
  
  # Example 4D ragged tensor value, with two ragged dimensions and with values
  # whose shape is [2], expressed as nested python lists and as splits+values.
  EXAMPLE_RAGGED_TENSOR_4D = [
      [                                       # rt[0]
          [[1, 2], [3, 4], [5, 6]],           # rt[0][0]
          [[7, 8], [9, 10], [11, 12]]],       # rt[0][1]
      [],                                     # rt[1]
      [                                       # rt[2]
          [[13, 14], [15, 16], [17, 18]]],    # rt[2][0]
      [                                       # rt[3]
          [[19, 20]]]                         # rt[3][0]
  ]  # pyformat: disable
  EXAMPLE_RAGGED_TENSOR_4D_SPLITS1 = [0, 2, 2, 3, 4]
  EXAMPLE_RAGGED_TENSOR_4D_SPLITS2 = [0, 3, 6, 9, 10]
  EXAMPLE_RAGGED_TENSOR_4D_VALUES = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10],
                                     [11, 12], [13, 14], [15, 16], [17, 18],
                                     [19, 20]]
  
  # Example 3D ragged tensor with uniform_row_lengths.
  EXAMPLE_RAGGED_TENSOR_3D = [[[1, 2, 3], [4], [5, 6]], [[], [7, 8, 9], []]]
  EXAMPLE_RAGGED_TENSOR_3D_ROWLEN = 3
  EXAMPLE_RAGGED_TENSOR_3D_SPLITS = [0, 3, 4, 6, 6, 9, 9]
  EXAMPLE_RAGGED_TENSOR_3D_VALUES = [1, 2, 3, 4, 5, 6, 7, 8, 9]
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class RaggedGetItemTest(test_util.TensorFlowTestCase, parameterized.TestCase):
    longMessage = True  # Property in unittest.Testcase. pylint: disable=invalid-name
  
    #=============================================================================
    # RaggedTensor.__getitem__
    #=============================================================================
  
    def _TestGetItem(self, rt, slice_spec, expected, expected_shape=None):
      """Helper function for testing RaggedTensor.__getitem__.
  
      Checks that calling `rt.__getitem__(slice_spec) returns the expected value.
      Checks three different configurations for each slice spec:
  
        * Call __getitem__ with the slice spec as-is (with int values)
        * Call __getitem__ with int values in the slice spec wrapped in
          `tf.constant()`.
        * Call __getitem__ with int values in the slice spec wrapped in
          `tf.compat.v1.placeholder()` (so value is not known at graph
          construction time).
  
      Args:
        rt: The RaggedTensor to test.
        slice_spec: The slice spec.
        expected: The expected value of rt.__getitem__(slice_spec), as a python
          list; or an exception class.
        expected_shape: The expected shape for `rt.__getitem__(slice_spec)`.
      """
      tensor_slice_spec1 = _make_tensor_slice_spec(slice_spec, True)
      tensor_slice_spec2 = _make_tensor_slice_spec(slice_spec, False)
      value1 = rt.__getitem__(slice_spec)
      value2 = rt.__getitem__(tensor_slice_spec1)
      value3 = rt.__getitem__(tensor_slice_spec2)
      self.assertAllEqual(value1, expected, 'slice_spec=%s' % (slice_spec,))
      self.assertAllEqual(value2, expected, 'slice_spec=%s' % (slice_spec,))
      self.assertAllEqual(value3, expected, 'slice_spec=%s' % (slice_spec,))
      if expected_shape is not None:
        value1.shape.assert_is_compatible_with(expected_shape)
        value2.shape.assert_is_compatible_with(expected_shape)
        value3.shape.assert_is_compatible_with(expected_shape)
  
    def _TestGetItemException(self, rt, slice_spec, expected, message):
      """Helper function for testing RaggedTensor.__getitem__ exceptions."""
      tensor_slice_spec = _make_tensor_slice_spec(slice_spec, True)
      with self.assertRaisesRegex(expected, message):
        self.evaluate(rt.__getitem__(slice_spec))
      with self.assertRaisesRegex(expected, message):
        self.evaluate(rt.__getitem__(tensor_slice_spec))
  
    @parameterized.parameters(
        # Tests for rt[i]
        (SLICE_BUILDER[-5], EXAMPLE_RAGGED_TENSOR_2D[-5]),
        (SLICE_BUILDER[-4], EXAMPLE_RAGGED_TENSOR_2D[-4]),
        (SLICE_BUILDER[-1], EXAMPLE_RAGGED_TENSOR_2D[-1]),
        (SLICE_BUILDER[0], EXAMPLE_RAGGED_TENSOR_2D[0]),
        (SLICE_BUILDER[1], EXAMPLE_RAGGED_TENSOR_2D[1]),
        (SLICE_BUILDER[4], EXAMPLE_RAGGED_TENSOR_2D[4]),
  
        # Tests for rt[i:]
        (SLICE_BUILDER[-6:], EXAMPLE_RAGGED_TENSOR_2D[-6:]),
        (SLICE_BUILDER[-3:], EXAMPLE_RAGGED_TENSOR_2D[-3:]),
        (SLICE_BUILDER[-1:], EXAMPLE_RAGGED_TENSOR_2D[-1:]),
        (SLICE_BUILDER[0:], EXAMPLE_RAGGED_TENSOR_2D[0:]),
        (SLICE_BUILDER[3:], EXAMPLE_RAGGED_TENSOR_2D[3:]),
        (SLICE_BUILDER[5:], EXAMPLE_RAGGED_TENSOR_2D[5:]),
  
        # Tests for rt[:j]
        (SLICE_BUILDER[:-6], EXAMPLE_RAGGED_TENSOR_2D[:-6]),
        (SLICE_BUILDER[:-3], EXAMPLE_RAGGED_TENSOR_2D[:-3]),
        (SLICE_BUILDER[:-1], EXAMPLE_RAGGED_TENSOR_2D[:-1]),
        (SLICE_BUILDER[:0], EXAMPLE_RAGGED_TENSOR_2D[:0]),
        (SLICE_BUILDER[:3], EXAMPLE_RAGGED_TENSOR_2D[:3]),
        (SLICE_BUILDER[:5], EXAMPLE_RAGGED_TENSOR_2D[:5]),
  
        # Tests for rt[i:j]
        (SLICE_BUILDER[0:3], EXAMPLE_RAGGED_TENSOR_2D[0:3]),
        (SLICE_BUILDER[3:5], EXAMPLE_RAGGED_TENSOR_2D[3:5]),
        (SLICE_BUILDER[-5:3], EXAMPLE_RAGGED_TENSOR_2D[-5:3]),
        (SLICE_BUILDER[3:1], EXAMPLE_RAGGED_TENSOR_2D[3:1]),
        (SLICE_BUILDER[-1:1], EXAMPLE_RAGGED_TENSOR_2D[-1:1]),
        (SLICE_BUILDER[1:-1], EXAMPLE_RAGGED_TENSOR_2D[1:-1]),
  
        # Tests for rt[i, j]
        (SLICE_BUILDER[0, 1], EXAMPLE_RAGGED_TENSOR_2D[0][1]),
        (SLICE_BUILDER[1, 2], EXAMPLE_RAGGED_TENSOR_2D[1][2]),
        (SLICE_BUILDER[-1, 0], EXAMPLE_RAGGED_TENSOR_2D[-1][0]),
        (SLICE_BUILDER[-3, 0], EXAMPLE_RAGGED_TENSOR_2D[-3][0]),
        (SLICE_BUILDER[:], EXAMPLE_RAGGED_TENSOR_2D),
        (SLICE_BUILDER[:, :], EXAMPLE_RAGGED_TENSOR_2D),
  
        # Empty slice spec.
        ([], EXAMPLE_RAGGED_TENSOR_2D),
  
        # Test for ellipsis
        (SLICE_BUILDER[...], EXAMPLE_RAGGED_TENSOR_2D),
        (SLICE_BUILDER[2, ...], EXAMPLE_RAGGED_TENSOR_2D[2]),
        (SLICE_BUILDER[..., :], EXAMPLE_RAGGED_TENSOR_2D),
        (SLICE_BUILDER[..., 2, 0], EXAMPLE_RAGGED_TENSOR_2D[2][0]),
        (SLICE_BUILDER[2, ..., 0], EXAMPLE_RAGGED_TENSOR_2D[2][0]),
        (SLICE_BUILDER[2, 0, ...], EXAMPLE_RAGGED_TENSOR_2D[2][0]),
  
        # Test for array_ops.newaxis
        (SLICE_BUILDER[array_ops.newaxis, :], [EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[:, array_ops.newaxis],
         [[row] for row in EXAMPLE_RAGGED_TENSOR_2D]),
  
        # Slicing inner ragged dimensions.
        (SLICE_BUILDER[-1:,
                       1:4], [row[1:4] for row in EXAMPLE_RAGGED_TENSOR_2D[-1:]]),
        (SLICE_BUILDER[:, 1:4], [row[1:4] for row in EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[:, -2:], [row[-2:] for row in EXAMPLE_RAGGED_TENSOR_2D]),
  
        # Strided slices
        (SLICE_BUILDER[::2], EXAMPLE_RAGGED_TENSOR_2D[::2]),
        (SLICE_BUILDER[::-1], EXAMPLE_RAGGED_TENSOR_2D[::-1]),
        (SLICE_BUILDER[::-2], EXAMPLE_RAGGED_TENSOR_2D[::-2]),
        (SLICE_BUILDER[::-3], EXAMPLE_RAGGED_TENSOR_2D[::-3]),
        (SLICE_BUILDER[:, ::2], [row[::2] for row in EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[:, ::-1], [row[::-1] for row in EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[:, ::-2], [row[::-2] for row in EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[:, ::-3], [row[::-3] for row in EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[:, 2::-1],
         [row[2::-1] for row in EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[:, -1::-1],
         [row[-1::-1] for row in EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[..., -1::-1],
         [row[-1::-1] for row in EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[:, 2::-2],
         [row[2::-2] for row in EXAMPLE_RAGGED_TENSOR_2D]),
        (SLICE_BUILDER[::-1, ::-1],
         [row[::-1] for row in EXAMPLE_RAGGED_TENSOR_2D[::-1]]),
    )  # pyformat: disable
    def testWithRaggedRank1(self, slice_spec, expected):
      """Test that rt.__getitem__(slice_spec) == expected."""
      # Ragged tensor
      rt = RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES,
                                        EXAMPLE_RAGGED_TENSOR_2D_SPLITS)
  
      self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_2D)
      self._TestGetItem(rt, slice_spec, expected)
  
    # pylint: disable=g-complex-comprehension
    @parameterized.parameters([(start, stop)
                               for start in [-2, -1, None, 0, 1, 2]
                               for stop in [-2, -1, None, 0, 1, 2]])
    def testWithStridedSlices(self, start, stop):
      test_value = [[1, 2, 3, 4, 5], [6, 7], [8, 9, 10], [], [9],
                    [1, 2, 3, 4, 5, 6, 7, 8]]
      rt = ragged_factory_ops.constant(test_value)
      for step in [-3, -2, -1, 1, 2, 3]:
        # Slice outer dimension
        self.assertAllEqual(rt[start:stop:step], test_value[start:stop:step],
                            'slice=%s:%s:%s' % (start, stop, step))
        # Slice inner dimension
        self.assertAllEqual(rt[:, start:stop:step],
                            [row[start:stop:step] for row in test_value],
                            'slice=%s:%s:%s' % (start, stop, step))
  
    # pylint: disable=invalid-slice-index
    @parameterized.parameters(
        # Tests for out-of-bound errors
        (SLICE_BUILDER[5], (IndexError, ValueError, errors.InvalidArgumentError),
         '.*out of bounds.*'),
        (SLICE_BUILDER[-6], (IndexError, ValueError, errors.InvalidArgumentError),
         '.*out of bounds.*'),
        (SLICE_BUILDER[0, 2], (IndexError, ValueError,
                               errors.InvalidArgumentError), '.*out of bounds.*'),
        (SLICE_BUILDER[3, 0], (IndexError, ValueError,
                               errors.InvalidArgumentError), '.*out of bounds.*'),
  
        # Indexing into an inner ragged dimension
        (SLICE_BUILDER[:, 3], ValueError,
         'Cannot index into an inner ragged dimension'),
        (SLICE_BUILDER[:1, 3], ValueError,
         'Cannot index into an inner ragged dimension'),
        (SLICE_BUILDER[..., 3], ValueError,
         'Cannot index into an inner ragged dimension'),
  
        # Tests for type errors
        (SLICE_BUILDER[0.5], TypeError, re.escape(array_ops._SLICE_TYPE_ERROR)),
        (SLICE_BUILDER[1:3:0.5], TypeError, re.escape(
            array_ops._SLICE_TYPE_ERROR)),
        (SLICE_BUILDER[:, 1:3:0.5], TypeError,
         'slice strides must be integers or None'),
        (SLICE_BUILDER[:, 0.5:1.5], TypeError,
         'slice offsets must be integers or None'),
        (SLICE_BUILDER['foo'], TypeError, re.escape(array_ops._SLICE_TYPE_ERROR)),
        (SLICE_BUILDER[:, 'foo':'foo'], TypeError,
         'slice offsets must be integers or None'),
  
        # Tests for other errors
        (SLICE_BUILDER[..., 0, 0,
                       0], IndexError, 'Too many indices for RaggedTensor'),
    )
    def testErrorsWithRaggedRank1(self, slice_spec, expected, message):
      """Test that rt.__getitem__(slice_spec) == expected."""
      # Ragged tensor
      rt = RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES,
                                        EXAMPLE_RAGGED_TENSOR_2D_SPLITS)
  
      self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_2D)
      self._TestGetItemException(rt, slice_spec, expected, message)
  
    @parameterized.parameters(
        # Tests for rt[index, index, ...]
        (SLICE_BUILDER[2, 0], EXAMPLE_RAGGED_TENSOR_4D[2][0]),
        (SLICE_BUILDER[2, 0, 1], EXAMPLE_RAGGED_TENSOR_4D[2][0][1]),
        (SLICE_BUILDER[2, 0, 1, 1], EXAMPLE_RAGGED_TENSOR_4D[2][0][1][1]),
        (SLICE_BUILDER[2, 0, 1:], EXAMPLE_RAGGED_TENSOR_4D[2][0][1:]),
        (SLICE_BUILDER[2, 0, 1:, 1:], [[16], [18]]),
        (SLICE_BUILDER[2, 0, :, 1], [14, 16, 18]),
        (SLICE_BUILDER[2, 0, 1, :], EXAMPLE_RAGGED_TENSOR_4D[2][0][1]),
  
        # Tests for rt[index, slice, ...]
        (SLICE_BUILDER[0, :], EXAMPLE_RAGGED_TENSOR_4D[0]),
        (SLICE_BUILDER[1, :], EXAMPLE_RAGGED_TENSOR_4D[1]),
        (SLICE_BUILDER[0, :, :, 1], [[2, 4, 6], [8, 10, 12]]),
        (SLICE_BUILDER[1, :, :, 1], []),
        (SLICE_BUILDER[2, :, :, 1], [[14, 16, 18]]),
        (SLICE_BUILDER[3, :, :, 1], [[20]]),
  
        # Tests for rt[slice, slice, ...]
        (SLICE_BUILDER[:, :], EXAMPLE_RAGGED_TENSOR_4D),
        (SLICE_BUILDER[:, :, :, 1], [[[2, 4, 6], [8, 10, 12]], [], [[14, 16, 18]],
                                     [[20]]]),
        (SLICE_BUILDER[1:, :, :, 1], [[], [[14, 16, 18]], [[20]]]),
        (SLICE_BUILDER[-3:, :, :, 1], [[], [[14, 16, 18]], [[20]]]),
  
        # Test for ellipsis
        (SLICE_BUILDER[...], EXAMPLE_RAGGED_TENSOR_4D),
        (SLICE_BUILDER[2, ...], EXAMPLE_RAGGED_TENSOR_4D[2]),
        (SLICE_BUILDER[2, 0, ...], EXAMPLE_RAGGED_TENSOR_4D[2][0]),
        (SLICE_BUILDER[..., 0], [[[1, 3, 5], [7, 9, 11]], [], [[13, 15, 17]],
                                 [[19]]]),
        (SLICE_BUILDER[2, ..., 0], [[13, 15, 17]]),
        (SLICE_BUILDER[2, 0, ..., 0], [13, 15, 17]),
  
        # Test for array_ops.newaxis
        (SLICE_BUILDER[array_ops.newaxis, :], [EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[:, array_ops.newaxis],
         [[row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
  
        # Empty slice spec.
        ([], EXAMPLE_RAGGED_TENSOR_4D),
  
        # Slicing inner ragged dimensions.
        (SLICE_BUILDER[:, 1:4], [row[1:4] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[:, -2:], [row[-2:] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[:, :, :-1],
         [[v[:-1] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[:, :, 1:2],
         [[v[1:2] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[1:, 1:3, 1:2],
         [[v[1:2] for v in row[1:3]] for row in EXAMPLE_RAGGED_TENSOR_4D[1:]]),
  
        # Strided slices
        (SLICE_BUILDER[::2], EXAMPLE_RAGGED_TENSOR_4D[::2]),
        (SLICE_BUILDER[::-1], EXAMPLE_RAGGED_TENSOR_4D[::-1]),
        (SLICE_BUILDER[::-2], EXAMPLE_RAGGED_TENSOR_4D[::-2]),
        (SLICE_BUILDER[1::2], EXAMPLE_RAGGED_TENSOR_4D[1::2]),
        (SLICE_BUILDER[:, ::2], [row[::2] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[:, 1::2], [row[1::2] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[:, :, ::2],
         [[v[::2] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[:, :, 1::2],
         [[v[1::2] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[:, :, ::-1],
         [[v[::-1] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[:, :, ::-2],
         [[v[::-2] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[..., ::-1, :],
         [[v[::-1] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
        (SLICE_BUILDER[..., ::-1], [[[v[::-1] for v in col] for col in row]
                                    for row in EXAMPLE_RAGGED_TENSOR_4D]),
    )  # pyformat: disable
    def testWithRaggedRank2(self, slice_spec, expected):
      """Test that rt.__getitem__(slice_spec) == expected."""
      rt = RaggedTensor.from_nested_row_splits(
          EXAMPLE_RAGGED_TENSOR_4D_VALUES,
          [EXAMPLE_RAGGED_TENSOR_4D_SPLITS1, EXAMPLE_RAGGED_TENSOR_4D_SPLITS2])
      self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_4D)
      self._TestGetItem(rt, slice_spec, expected)
  
    @parameterized.parameters(
        # Test for errors in unsupported cases
        (SLICE_BUILDER[:, 0], ValueError,
         'Cannot index into an inner ragged dimension.'),
        (SLICE_BUILDER[:, :, 0], ValueError,
         'Cannot index into an inner ragged dimension.'),
  
        # Test for out-of-bounds errors.
        (SLICE_BUILDER[1, 0], (IndexError, ValueError,
                               errors.InvalidArgumentError), '.*out of bounds.*'),
        (SLICE_BUILDER[0, 0, 3],
         (IndexError, ValueError,
          errors.InvalidArgumentError), '.*out of bounds.*'),
        (SLICE_BUILDER[5], (IndexError, ValueError, errors.InvalidArgumentError),
         '.*out of bounds.*'),
        (SLICE_BUILDER[0, 5], (IndexError, ValueError,
                               errors.InvalidArgumentError), '.*out of bounds.*'),
    )
    def testErrorsWithRaggedRank2(self, slice_spec, expected, message):
      """Test that rt.__getitem__(slice_spec) == expected."""
      rt = RaggedTensor.from_nested_row_splits(
          EXAMPLE_RAGGED_TENSOR_4D_VALUES,
          [EXAMPLE_RAGGED_TENSOR_4D_SPLITS1, EXAMPLE_RAGGED_TENSOR_4D_SPLITS2])
      self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_4D)
      self._TestGetItemException(rt, slice_spec, expected, message)
  
    @parameterized.parameters(
        (SLICE_BUILDER[:], []),
        (SLICE_BUILDER[2:], []),
        (SLICE_BUILDER[:-3], []),
    )
    def testWithEmptyTensor(self, slice_spec, expected):
      """Test that rt.__getitem__(slice_spec) == expected."""
      rt = RaggedTensor.from_row_splits([], [0])
      self._TestGetItem(rt, slice_spec, expected)
  
    @parameterized.parameters(
        (SLICE_BUILDER[0], (IndexError, ValueError, errors.InvalidArgumentError),
         '.*out of bounds.*'),
        (SLICE_BUILDER[-1], (IndexError, ValueError, errors.InvalidArgumentError),
         '.*out of bounds.*'),
    )
    def testErrorsWithEmptyTensor(self, slice_spec, expected, message):
      """Test that rt.__getitem__(slice_spec) == expected."""
      rt = RaggedTensor.from_row_splits([], [0])
      self._TestGetItemException(rt, slice_spec, expected, message)
  
    @parameterized.parameters(
        (SLICE_BUILDER[-4], EXAMPLE_RAGGED_TENSOR_2D[-4]),
        (SLICE_BUILDER[0], EXAMPLE_RAGGED_TENSOR_2D[0]),
        (SLICE_BUILDER[-3:], EXAMPLE_RAGGED_TENSOR_2D[-3:]),
        (SLICE_BUILDER[:3], EXAMPLE_RAGGED_TENSOR_2D[:3]),
        (SLICE_BUILDER[3:5], EXAMPLE_RAGGED_TENSOR_2D[3:5]),
        (SLICE_BUILDER[0, 1], EXAMPLE_RAGGED_TENSOR_2D[0][1]),
        (SLICE_BUILDER[-3, 0], EXAMPLE_RAGGED_TENSOR_2D[-3][0]),
    )
    def testWithPlaceholderShapes(self, slice_spec, expected):
      """Test that rt.__getitem__(slice_spec) == expected."""
      # Intentionally use an unknown shape for `splits`, to force the code path
      # that deals with having nrows unknown at graph construction time.
      splits = constant_op.constant(
          EXAMPLE_RAGGED_TENSOR_2D_SPLITS, dtype=dtypes.int64)
      splits = array_ops.placeholder_with_default(splits, None)
      rt = RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES, splits)
      self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_2D)
      self._TestGetItem(rt, slice_spec, expected)
  
    @parameterized.parameters(
        (SLICE_BUILDER[..., 2], ValueError,
         'Ellipsis not supported for unknown shape RaggedTensors'),)
    def testErrorsWithPlaceholderShapes(self, slice_spec, expected, message):
      """Test that rt.__getitem__(slice_spec) == expected."""
      if not context.executing_eagerly():
        # Intentionally use an unknown shape for `values`.
        values = array_ops.placeholder_with_default([0], None)
        rt = RaggedTensor.from_row_splits(values, [0, 1])
        self._TestGetItemException(rt, slice_spec, expected, message)
  
    def testNewAxis(self):
      # rt: [[[['a', 'b'], ['c', 'd']], [], [['e', 'f']]], []]
      splits1 = [0, 3, 3]
      splits2 = [0, 2, 2, 3]
      values = constant_op.constant([['a', 'b'], ['c', 'd'], ['e', 'f']])
      rt = RaggedTensor.from_nested_row_splits(values, [splits1, splits2])
      rt_newaxis0 = rt[array_ops.newaxis]
      rt_newaxis1 = rt[:, array_ops.newaxis]
      rt_newaxis2 = rt[:, :, array_ops.newaxis]
      rt_newaxis3 = rt[:, :, :, array_ops.newaxis]
      rt_newaxis4 = rt[:, :, :, :, array_ops.newaxis]
  
      self.assertAllEqual(
          rt, [[[[b'a', b'b'], [b'c', b'd']], [], [[b'e', b'f']]], []])
      self.assertAllEqual(
          rt_newaxis0, [[[[[b'a', b'b'], [b'c', b'd']], [], [[b'e', b'f']]], []]])
      self.assertAllEqual(
          rt_newaxis1,
          [[[[[b'a', b'b'], [b'c', b'd']], [], [[b'e', b'f']]]], [[]]])
      self.assertAllEqual(
          rt_newaxis2,
          [[[[[b'a', b'b'], [b'c', b'd']]], [[]], [[[b'e', b'f']]]], []])
      self.assertAllEqual(
          rt_newaxis3,
          [[[[[b'a', b'b']], [[b'c', b'd']]], [], [[[b'e', b'f']]]], []])
      self.assertAllEqual(
          rt_newaxis4,
          [[[[[b'a'], [b'b']], [[b'c'], [b'd']]], [], [[[b'e'], [b'f']]]], []])
  
      self.assertEqual(rt.ragged_rank, 2)
      self.assertEqual(rt_newaxis0.ragged_rank, 3)
      self.assertEqual(rt_newaxis1.ragged_rank, 3)
      self.assertEqual(rt_newaxis2.ragged_rank, 3)
      self.assertEqual(rt_newaxis3.ragged_rank, 2)
      self.assertEqual(rt_newaxis4.ragged_rank, 2)
  
      self.assertEqual(rt_newaxis0.shape.as_list(), [1, 2, None, None, 2])
      self.assertEqual(rt_newaxis1.shape.as_list(), [2, 1, None, None, 2])
      self.assertEqual(rt_newaxis2.shape.as_list(), [2, None, 1, None, 2])
      self.assertEqual(rt_newaxis3.shape.as_list(), [2, None, None, 1, 2])
      self.assertEqual(rt_newaxis4.shape.as_list(), [2, None, None, 2, 1])
  
    @parameterized.parameters(
        # EXAMPLE_RAGGED_TENSOR_3D.shape = [2, 3, None]
  
        # Indexing into uniform_row_splits dimension:
        (SLICE_BUILDER[:, 1], [r[1] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, None]),
        (SLICE_BUILDER[:, 2], [r[2] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, None]),
        (SLICE_BUILDER[:, -2], [r[-2] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, None]),
        (SLICE_BUILDER[:, -3], [r[-3] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, None]),
        (SLICE_BUILDER[1:, 2], [r[2] for r in EXAMPLE_RAGGED_TENSOR_3D[1:]],
         [1, None]),
        (SLICE_BUILDER[:, 1, 1:], [r[1][1:] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, None]),
        (SLICE_BUILDER[1:, 1, 1:],
         [r[1][1:] for r in EXAMPLE_RAGGED_TENSOR_3D[1:]],
         [1, None]),
  
        # Slicing uniform_row_splits dimension:
        (SLICE_BUILDER[:, 2:], [r[2:] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, 1, None]),
        (SLICE_BUILDER[:, -2:], [r[-2:] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, 2, None]),
        (SLICE_BUILDER[:, :, 1:],
         [[c[1:] for c in r] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, 3, None]),
        (SLICE_BUILDER[:, 5:], [r[5:] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, 0, None]),
  
        # Slicing uniform_row_splits dimension with a non-default step size:
        (SLICE_BUILDER[:, ::2], [r[::2] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, 2, None]),
        (SLICE_BUILDER[:, ::-1], [r[::-1] for r in EXAMPLE_RAGGED_TENSOR_3D],
         [2, 3, None]),
    )  # pyformat: disable
    def testWithUniformRowLength(self, slice_spec, expected, expected_shape):
      """Test that rt.__getitem__(slice_spec) == expected."""
      rt = RaggedTensor.from_uniform_row_length(
          RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_3D_VALUES,
                                       EXAMPLE_RAGGED_TENSOR_3D_SPLITS),
          EXAMPLE_RAGGED_TENSOR_3D_ROWLEN)
      self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_3D)
      self.assertIsNot(rt.uniform_row_length, None)
      self._TestGetItem(rt, slice_spec, expected, expected_shape)
  
      # If the result is 3D, then check that it still has a uniform row length:
      actual = rt.__getitem__(slice_spec)  # pylint: disable=assignment-from-no-return
      if actual.shape.rank == 3:
        self.assertIsNot(actual.uniform_row_length, None)
        self.assertAllEqual(actual.uniform_row_length, expected_shape[1])
  
    @parameterized.parameters(
        (SLICE_BUILDER[:, 3], errors.InvalidArgumentError, 'out of bounds'),
        (SLICE_BUILDER[:, -4], errors.InvalidArgumentError, 'out of bounds'),
        (SLICE_BUILDER[:, 10], errors.InvalidArgumentError, 'out of bounds'),
        (SLICE_BUILDER[:, -10], errors.InvalidArgumentError, 'out of bounds'),
    )
    def testErrorsWithUniformRowLength(self, slice_spec, expected, message):
      """Test that rt.__getitem__(slice_spec) == expected."""
      rt = RaggedTensor.from_uniform_row_length(
          RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_3D_VALUES,
                                       EXAMPLE_RAGGED_TENSOR_3D_SPLITS),
          EXAMPLE_RAGGED_TENSOR_3D_ROWLEN)
      self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_3D)
      self._TestGetItemException(rt, slice_spec, expected, message)
  
  
  if __name__ == '__main__':
    googletest.main()
   # Copyright 2015 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 tensorflow.python.framework.sparse_tensor."""
  from absl.testing import parameterized
  import numpy as np
  
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import sparse_ops
  from tensorflow.python.platform import googletest
  
  
  class SparseTensorTest(test_util.TensorFlowTestCase):
  
    def testPythonConstruction(self):
      indices = [[1, 2], [2, 0], [3, 4]]
      values = [b"a", b"b", b"c"]
      shape = [4, 5]
      sp_value = sparse_tensor.SparseTensorValue(indices, values, shape)
      for sp in [
          sparse_tensor.SparseTensor(indices, values, shape),
          sparse_tensor.SparseTensor.from_value(sp_value),
          sparse_tensor.SparseTensor.from_value(
              sparse_tensor.SparseTensor(indices, values, shape))]:
        self.assertEqual(sp.indices.dtype, dtypes.int64)
        self.assertEqual(sp.values.dtype, dtypes.string)
        self.assertEqual(sp.dense_shape.dtype, dtypes.int64)
        self.assertEqual(sp.get_shape(), (4, 5))
  
        value = self.evaluate(sp)
        self.assertAllEqual(indices, value.indices)
        self.assertAllEqual(values, value.values)
        self.assertAllEqual(shape, value.dense_shape)
        sp_value = self.evaluate(sp)
        self.assertAllEqual(sp_value.indices, value.indices)
        self.assertAllEqual(sp_value.values, value.values)
        self.assertAllEqual(sp_value.dense_shape, value.dense_shape)
  
    def testShape(self):
  
      @def_function.function
      def test_fn(tensor):
        tensor = sparse_ops.sparse_transpose(tensor)
        self.assertEqual(tensor.shape.rank, 2)
        return tensor
  
      tensor = sparse_tensor.SparseTensor(
          indices=[[0, 0], [1, 2]], values=[1., 2], dense_shape=[3, 4])
      test_fn(tensor)
  
    def testIsSparse(self):
      self.assertFalse(sparse_tensor.is_sparse(3))
      self.assertFalse(sparse_tensor.is_sparse("foo"))
      self.assertFalse(sparse_tensor.is_sparse(np.array(3)))
      self.assertTrue(
          sparse_tensor.is_sparse(sparse_tensor.SparseTensor([[0]], [0], [1])))
      self.assertTrue(
          sparse_tensor.is_sparse(
              sparse_tensor.SparseTensorValue([[0]], [0], [1])))
  
    def testConsumers(self):
      with context.graph_mode():
        sp = sparse_tensor.SparseTensor([[0, 0], [1, 2]], [1.0, 3.0], [3, 4])
        w = ops.convert_to_tensor(np.ones([4, 1], np.float32))
        out = sparse_ops.sparse_tensor_dense_matmul(sp, w)
        self.assertEqual(len(sp.consumers()), 1)
        self.assertEqual(sp.consumers()[0], out.op)
  
        dense = sparse_ops.sparse_tensor_to_dense(sp)
        self.assertEqual(len(sp.consumers()), 2)
        self.assertIn(dense.op, sp.consumers())
        self.assertIn(out.op, sp.consumers())
  
    def testWithValues(self):
      source = sparse_tensor.SparseTensor(
          indices=[[0, 0], [1, 2]], values=[1., 2], dense_shape=[3, 4])
      new_tensor = source.with_values([5.0, 1.0])
      self.assertAllEqual(new_tensor.indices, source.indices)
      self.assertAllEqual(new_tensor.values, [5.0, 1.0])
      self.assertAllEqual(new_tensor.dense_shape, source.dense_shape)
  
      # ensure new value's shape is checked
      with self.assertRaises((errors.InvalidArgumentError, ValueError)):
        source.with_values([[5.0, 1.0]])
  
  
  class ConvertToTensorOrSparseTensorTest(test_util.TensorFlowTestCase):
  
    def test_convert_dense(self):
      value = [42, 43]
      from_value = sparse_tensor.convert_to_tensor_or_sparse_tensor(
          value)
      self.assertAllEqual(value, self.evaluate(from_value))
  
    def test_convert_sparse(self):
      indices = [[0, 1], [1, 0]]
      values = [42, 43]
      shape = [2, 2]
      sparse_tensor_value = sparse_tensor.SparseTensorValue(
          indices, values, shape)
      st = sparse_tensor.SparseTensor.from_value(sparse_tensor_value)
      from_value = self.evaluate(
          sparse_tensor.convert_to_tensor_or_sparse_tensor(sparse_tensor_value))
      from_tensor = self.evaluate(
          sparse_tensor.convert_to_tensor_or_sparse_tensor(st))
      for convertee in [from_value, from_tensor]:
        self.assertAllEqual(sparse_tensor_value.indices, convertee.indices)
        self.assertAllEqual(sparse_tensor_value.values, convertee.values)
        self.assertAllEqual(
            sparse_tensor_value.dense_shape, convertee.dense_shape)
  
  
  class SparseTensorShapeTest(test_util.TensorFlowTestCase):
  
    def test_simple(self):
      indices = [[0, 2]]
      values = [1]
      dense_shape = [5, 5]
      sp = sparse_tensor.SparseTensor(indices, values, dense_shape)
  
      self.assertIsInstance(sp.shape, tensor_shape.TensorShape)
      self.assertIsInstance(sp.dense_shape, ops.Tensor)
      self.assertEqual(sp.shape.as_list(), [5, 5])
  
    def test_unknown_shape(self):
  
      @def_function.function
      def my_func(dense_shape):
        indices = [[0, 2]]
        values = [1]
        sp = sparse_tensor.SparseTensor(indices, values, dense_shape)
        self.assertEqual(sp.shape.as_list(), [None, None])
        return sp
  
      my_func.get_concrete_function(
          dense_shape=tensor_spec.TensorSpec(
              dtype=dtypes.int64, shape=[2,]))
  
    def test_partial_shape(self):
  
      @def_function.function
      def my_func(x):
        indices = [[0, 2]]
        values = [1]
        y = ops.convert_to_tensor(3, dtype=dtypes.int64)
        dense_shape = [x, y]
        sp = sparse_tensor.SparseTensor(indices, values, dense_shape)
        self.assertEqual(sp.shape.as_list(), [None, 3])
        return sp
  
      my_func.get_concrete_function(
          x=tensor_spec.TensorSpec(dtype=dtypes.int64, shape=[]))
  
    def test_neg_shape(self):
      indices = [[0, 2]]
      values = [1]
      dense_shape = [-1, 5]
      sp = sparse_tensor.SparseTensor(indices, values, dense_shape)
      self.assertEqual(sp.shape.as_list(), [None, 5])
  
    def test_unknown_tensor_shape(self):
  
      @def_function.function
      def my_func(x):
        indices = [[0, 0]]
        values = [1]
        dense_shape = array_ops.shape(x)
        dense_shape = math_ops.cast(dense_shape, dtypes.int64)
  
        sp = sparse_tensor.SparseTensor(indices, values, dense_shape)
        self.assertEqual(sp.shape.as_list(), [None, None])
        return sp
  
      my_func.get_concrete_function(
          x=tensor_spec.TensorSpec(dtype=dtypes.int64, shape=[None, None]))
  
    def test_unknown_rank(self):
  
      @def_function.function
      def my_func(dense_shape):
        indices = [[0, 0]]
        values = [1]
        sp = sparse_tensor.SparseTensor(indices, values, dense_shape)
        self.assertEqual(sp.shape.rank, None)
        return sp
  
      my_func.get_concrete_function(
          dense_shape=tensor_spec.TensorSpec(dtype=dtypes.int64, shape=[None]))
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class SparseTensorSpecTest(test_util.TensorFlowTestCase,
                             parameterized.TestCase):
  
    def assertAllTensorsEqual(self, list1, list2):
      self.assertLen(list1, len(list2))
      for (t1, t2) in zip(list1, list2):
        self.assertAllEqual(t1, t2)
  
    def testConstruction(self):
      spec1 = sparse_tensor.SparseTensorSpec()
      self.assertEqual(spec1.shape.rank, None)
      self.assertEqual(spec1.dtype, dtypes.float32)
  
      spec2 = sparse_tensor.SparseTensorSpec([None, None], dtypes.string)
      self.assertEqual(spec2.shape.as_list(), [None, None])
      self.assertEqual(spec2.dtype, dtypes.string)
  
    def testValueType(self):
      spec1 = sparse_tensor.SparseTensorSpec()
      self.assertEqual(spec1.value_type, sparse_tensor.SparseTensor)
  
    @parameterized.parameters([
        (sparse_tensor.SparseTensorSpec(),
         (tensor_shape.TensorShape(None), dtypes.float32)),
        (sparse_tensor.SparseTensorSpec(shape=[5, None, None]),
         (tensor_shape.TensorShape([5, None, None]), dtypes.float32)),
        (sparse_tensor.SparseTensorSpec(dtype=dtypes.int32),
         (tensor_shape.TensorShape(None), dtypes.int32)),
    ])  # pyformat: disable
    def testSerialize(self, st_spec, expected):
      serialization = st_spec._serialize()
      # TensorShape has an unconventional definition of equality, so we can't use
      # assertEqual directly here.  But repr() is deterministic and lossless for
      # the expected values, so we can use that instead.
      self.assertEqual(repr(serialization), repr(expected))
  
    @parameterized.parameters([
        (sparse_tensor.SparseTensorSpec(dtype=dtypes.string), [
            tensor_spec.TensorSpec([None, None], dtypes.int64),
            tensor_spec.TensorSpec([None], dtypes.string),
            tensor_spec.TensorSpec([None], dtypes.int64)
        ]),
        (sparse_tensor.SparseTensorSpec(shape=[5, None, None]), [
            tensor_spec.TensorSpec([None, 3], dtypes.int64),
            tensor_spec.TensorSpec([None], dtypes.float32),
            tensor_spec.TensorSpec([3], dtypes.int64)
        ]),
    ])
    def testComponentSpecs(self, st_spec, expected):
      self.assertEqual(st_spec._component_specs, expected)
  
    @parameterized.parameters([
        {
            "st_spec": sparse_tensor.SparseTensorSpec(),
            "indices": [[0, 1], [10, 8]],
            "values": [3.0, 5.0],
            "dense_shape": [100, 100]
        },
        {
            "st_spec": sparse_tensor.SparseTensorSpec([100, None, None]),
            "indices": [[0, 1, 3], [10, 8, 2]],
            "values": [3.0, 5.0],
            "dense_shape": [100, 20, 20]
        },
    ])
    def testToFromComponents(self, st_spec, indices, values, dense_shape):
      st = sparse_tensor.SparseTensor(indices, values, dense_shape)
      actual_components = st_spec._to_components(st)
      self.assertAllTensorsEqual(actual_components,
                                 [indices, values, dense_shape])
      st_reconstructed = st_spec._from_components(actual_components)
      self.assertAllEqual(st.indices, st_reconstructed.indices)
      self.assertAllEqual(st.values, st_reconstructed.values)
      self.assertAllEqual(st.dense_shape, st_reconstructed.dense_shape)
  
    @test_util.run_v1_only("SparseTensorValue is deprecated in v2")
    def testFromNumpyComponents(self):
      indices = np.array([[0], [8]])
      values = np.array([1.0, 9.0])
      dense_shape = np.array([100])
      spec = sparse_tensor.SparseTensorSpec()
      st = spec._from_components([indices, values, dense_shape])
      self.assertIsInstance(st, sparse_tensor.SparseTensorValue)
      self.assertAllEqual(st.indices, indices)
      self.assertAllEqual(st.values, values)
      self.assertAllEqual(st.dense_shape, dense_shape)
  
    @parameterized.parameters([
        sparse_tensor.SparseTensorSpec(dtype=dtypes.string),
        sparse_tensor.SparseTensorSpec(shape=[5, None, None]),
    ])
    def testFlatTensorSpecs(self, st_spec):
      self.assertEqual(st_spec._flat_tensor_specs,
                       [tensor_spec.TensorSpec(None, dtypes.variant)])
  
    @parameterized.parameters([
        {
            "st_spec": sparse_tensor.SparseTensorSpec(),
            "indices": [[0, 1], [10, 8]],
            "values": [3.0, 5.0],
            "dense_shape": [100, 100]
        },
        {
            "st_spec": sparse_tensor.SparseTensorSpec([100, None, None]),
            "indices": [[0, 1, 3], [10, 8, 2]],
            "values": [3.0, 5.0],
            "dense_shape": [100, 20, 20]
        },
    ])
    def testToFromTensorList(self, st_spec, indices, values, dense_shape):
      st = sparse_tensor.SparseTensor(indices, values, dense_shape)
      tensor_list = st_spec._to_tensor_list(st)
      st_reconstructed = st_spec._from_tensor_list(tensor_list)
      self.assertAllEqual(st.indices, st_reconstructed.indices)
      self.assertAllEqual(st.values, st_reconstructed.values)
      self.assertAllEqual(st.dense_shape, st_reconstructed.dense_shape)
  
    @parameterized.parameters([
        (sparse_tensor.SparseTensorSpec([2, None], dtypes.float32), 32,
         sparse_tensor.SparseTensorSpec([32, 2, None], dtypes.float32)),
        (sparse_tensor.SparseTensorSpec([4, None], dtypes.float32), None,
         sparse_tensor.SparseTensorSpec([None, 4, None], dtypes.float32)),
        (sparse_tensor.SparseTensorSpec([2], dtypes.float32), 32,
         sparse_tensor.SparseTensorSpec([32, 2], dtypes.float32)),
    ])
    def testBatch(self, spec, batch_size, expected):
      self.assertEqual(spec._batch(batch_size), expected)
  
    @parameterized.parameters([
        (sparse_tensor.SparseTensorSpec([32, None, None], dtypes.float32),
         sparse_tensor.SparseTensorSpec([None, None], dtypes.float32)),
        (sparse_tensor.SparseTensorSpec([None, None, None], dtypes.float32),
         sparse_tensor.SparseTensorSpec([None, None], dtypes.float32)),
        (sparse_tensor.SparseTensorSpec([32, 2], dtypes.float32),
         sparse_tensor.SparseTensorSpec([2], dtypes.float32)),
    ])
    def testUnbatch(self, spec, expected):
      self.assertEqual(spec._unbatch(), expected)
  
  
  if __name__ == "__main__":
    googletest.main()
   # 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.
  # ==============================================================================
  """Definitions for resource-type trackable object classes."""
  
  import contextlib
  import copy
  import weakref
  
  import six
  
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import ops
  from tensorflow.python.training.tracking import base
  from tensorflow.python.util import tf_contextlib
  from tensorflow.python.util.tf_export import tf_export
  
  # global _RESOURCE_TRACKER_STACK
  _RESOURCE_TRACKER_STACK = []
  
  
  class ResourceTracker(object):
    """An object that tracks a list of resources."""
  
    __slots__ = ["_resources"]
  
    def __init__(self):
      self._resources = []
  
    @property
    def resources(self):
      return self._resources
  
    def add_resource(self, resource):
      self._resources.append(resource)
  
  
  @tf_contextlib.contextmanager
  def resource_tracker_scope(resource_tracker):
    """A context to manage resource trackers.
  
    Use this in order to collect up all resources created within a block of code.
    Example usage:
  
    ```python
    resource_tracker = ResourceTracker()
    with resource_tracker_scope(resource_tracker):
      resource = TrackableResource()
  
    assert resource_tracker.resources == [resource]
  
    Args:
      resource_tracker: The passed in ResourceTracker object
  
    Yields:
      A scope in which the resource_tracker is active.
    """
    global _RESOURCE_TRACKER_STACK
    old = list(_RESOURCE_TRACKER_STACK)
    _RESOURCE_TRACKER_STACK.append(resource_tracker)
    try:
      yield
    finally:
      _RESOURCE_TRACKER_STACK = old
  
  
  def _make_getter(captured_getter, captured_previous):
    """To avoid capturing loop variables."""
  
    def getter(*args, **kwargs):
      return captured_getter(captured_previous, *args, **kwargs)
  
    return getter
  
  
  class _ResourceMetaclass(type):
    """Metaclass for CapturableResource."""
  
    def __call__(cls, *args, **kwargs):
  
      def default_resource_creator(next_creator, *a, **kw):
        assert next_creator is None
        obj = cls.__new__(cls, *a, **kw)
        obj.__init__(*a, **kw)
        return obj
  
      previous_getter = lambda *a, **kw: default_resource_creator(None, *a, **kw)
      resource_creator_stack = ops.get_default_graph()._resource_creator_stack
      for getter in resource_creator_stack[cls._resource_type()]:
        previous_getter = _make_getter(getter, previous_getter)
  
      return previous_getter(*args, **kwargs)
  
  
  class CapturableResource(six.with_metaclass(_ResourceMetaclass,
                                              base.Trackable)):
    """Holds a Tensor which a tf.function can capture.
  
    `CapturableResource`s are discovered by traversing the graph of object
    attributes, e.g. during `tf.saved_model.save`. They are excluded from the
    scope-based tracking of `TrackableResource`; generally things that require
    initialization should inherit from `TrackableResource` instead of
    `CapturableResource` directly.
    """
  
    def __init__(self, device=""):
      """Initialize the `CapturableResource`.
  
      Args:
        device: A string indicating a required placement for this resource,
          e.g. "CPU" if this resource must be created on a CPU device. A blank
          device allows the user to place resource creation, so generally this
          should be blank unless the resource only makes sense on one device.
      """
      self._resource_handle_value = None
      self._resource_device = device
      self._self_destruction_context = (
          context.eager_mode if context.executing_eagerly()
          else ops.get_default_graph().as_default)
  
    @classmethod
    def _resource_type(cls):
      return cls.__name__
  
    @property
    def _destruction_context(self):
      return getattr(self, "_self_destruction_context",
                     # no-op context
                     contextlib.suppress)
  
    @_destruction_context.setter
    def _destruction_context(self, destruction_context):
      self._self_destruction_context = destruction_context
  
    def _create_resource(self):
      """A function that creates a resource handle."""
      raise NotImplementedError("TrackableResource._create_resource not "
                                "implemented.")
  
    @property
    def _resource_handle(self):
      return self._resource_handle_value
  
    @_resource_handle.setter
    def _resource_handle(self, value):
      if isinstance(value, (ops.Tensor, ops.EagerTensor)):
        value._parent_trackable = weakref.ref(self)  # pylint: disable=protected-access
      self._resource_handle_value = value
  
    def _initialize(self):
      """A function that initializes the resource. Optional."""
      pass
  
    def _destroy_resource(self):
      """A function that destroys the resource. Optional."""
      pass
  
    @property
    def resource_handle(self):
      """Returns the resource handle associated with this Resource."""
      if self._resource_handle is None:
        with ops.device(self._resource_device):
          self._resource_handle = self._create_resource()
      return self._resource_handle
  
    def _map_resources(self, _):
      """For implementing `Trackable`."""
      new_obj = copy.copy(self)
      # pylint: disable=protected-access
      with ops.device(self._resource_device):
        new_resource = new_obj._create_resource()
      new_obj._resource_handle = new_resource
      # pylint: enable=protected-access
      obj_map = {self: new_obj}
      resource_map = {self.resource_handle: new_resource}
      return obj_map, resource_map
  
    def _trackable_children(self, save_type, **kwargs):
      children = super()._trackable_children(save_type, **kwargs)
      if save_type == "savedmodel":
        @def_function.function(input_signature=[], autograph=False)
        def _creator():
          resource = self._create_resource()
          return resource
  
        @def_function.function(input_signature=[], autograph=False)
        def _initializer():
          self._initialize()
          return 1  # Dummy return
  
        @def_function.function(input_signature=[], autograph=False)
        def _destroyer():
          self._destroy_resource()
          return 1  # Dummy return
  
        children.update({
            "_create_resource": _creator,
            "_initialize": _initializer,
            "_destroy_resource": _destroyer,
        })
      return children
  
    def __del__(self):
      try:
        # Outer race condition: on program exit, the destruction context may be
        # deleted before this __del__ is called. At this point we can safely
        # exit without calling _destroy_resource() and let Python handle things.
        with self._destruction_context():
          # Inner race condition: possible between this and `ScopedTFFunction`
          # whereby if an entire garbage collection chain containing both
          # objects is moved to unreachable during the same garbage collection
          # cycle, the __del__ for `ScopedTFFunction` can be collected before
          # this method is called. In that case, we can't do much but
          # continue.
          self._destroy_resource()
      except Exception:  # pylint: disable=broad-except
        # Silence all error logs that occur when attempting to destroy this
        # resource.
        pass
  
  
  @tf_export("saved_model.experimental.TrackableResource")
  class TrackableResource(CapturableResource):
    """Holds a Tensor which a tf.function can capture.
  
    A TrackableResource is most useful for stateful Tensors that require
    initialization, such as `tf.lookup.StaticHashTable`. `TrackableResource`s
    are discovered by traversing the graph of object attributes, e.g. during
    `tf.saved_model.save`.
  
    A TrackableResource has three methods to override:
  
    * `_create_resource` should create the resource tensor handle.
    * `_initialize` should initialize the resource held at `self.resource_handle`.
    * `_destroy_resource` is called upon a `TrackableResource`'s destruction
      and should decrement the resource's ref count. For most resources, this
      should be done with a call to `tf.raw_ops.DestroyResourceOp`.
  
    Example usage:
  
    >>> class DemoResource(tf.saved_model.experimental.TrackableResource):
    ...   def __init__(self):
    ...     super().__init__()
    ...     self._initialize()
    ...   def _create_resource(self):
    ...     return tf.raw_ops.VarHandleOp(dtype=tf.float32, shape=[2])
    ...   def _initialize(self):
    ...     tf.raw_ops.AssignVariableOp(
    ...         resource=self.resource_handle, value=tf.ones([2]))
    ...   def _destroy_resource(self):
    ...     tf.raw_ops.DestroyResourceOp(resource=self.resource_handle)
    >>> class DemoModule(tf.Module):
    ...   def __init__(self):
    ...     self.resource = DemoResource()
    ...   def increment(self, tensor):
    ...     return tensor + tf.raw_ops.ReadVariableOp(
    ...         resource=self.resource.resource_handle, dtype=tf.float32)
    >>> demo = DemoModule()
    >>> demo.increment([5, 1])
    <tf.Tensor: shape=(2,), dtype=float32, numpy=array([6., 2.], dtype=float32)>
    """
  
    def __init__(self, device=""):
      """Initialize the `TrackableResource`.
  
      Args:
        device: A string indicating a required placement for this resource,
          e.g. "CPU" if this resource must be created on a CPU device. A blank
          device allows the user to place resource creation, so generally this
          should be blank unless the resource only makes sense on one device.
      """
      global _RESOURCE_TRACKER_STACK
      for resource_tracker in _RESOURCE_TRACKER_STACK:
        resource_tracker.add_resource(self)
      super(TrackableResource, self).__init__(device=device)
  
  
  # TODO(b/124205571,b/124092991): Solve destruction of resources.
  class RestoredResource(TrackableResource):
    """Restored SavedResource."""
  
    def __init__(self, device=""):
      super(RestoredResource, self).__init__(device=device)
  
    @classmethod
    def _deserialize_from_proto(cls, object_proto, dependencies, **unused_kwargs):
      obj = cls(device=object_proto.resource.device)
      resource_creator = dependencies.get("_create_resource")
      if resource_creator is not None:
        obj._create_resource = resource_creator  # pylint: disable=protected-access
      return obj
  
    def _add_trackable_child(self, name, value):
      setattr(self, name, value)
      if (isinstance(value, base.Trackable) and
          not isinstance(value, def_function.Function)):
        self._track_trackable(value, name)
   # Copyright 2018 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.
  # ==============================================================================
  """Implementation of Cluster Resolvers for Kubernetes."""
  
  from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
  from tensorflow.python.distribute.cluster_resolver.cluster_resolver import format_master_url
  from tensorflow.python.training import server_lib
  from tensorflow.python.util.tf_export import tf_export
  
  
  @tf_export('distribute.cluster_resolver.KubernetesClusterResolver')
  class KubernetesClusterResolver(ClusterResolver):
    """ClusterResolver for Kubernetes.
  
    This is an implementation of cluster resolvers for Kubernetes. When given the
    the Kubernetes namespace and label selector for pods, we will retrieve the
    pod IP addresses of all running pods matching the selector, and return a
    ClusterSpec based on that information.
  
    Note: it cannot retrieve `task_type`, `task_id` or `rpc_layer`. To use it
    with some distribution strategies like
    `tf.distribute.experimental.MultiWorkerMirroredStrategy`, you will need to
    specify `task_type` and `task_id` by setting these attributes.
  
    Usage example with tf.distribute.Strategy:
  
      ```Python
      # On worker 0
      cluster_resolver = KubernetesClusterResolver(
          {"worker": ["job-name=worker-cluster-a", "job-name=worker-cluster-b"]})
      cluster_resolver.task_type = "worker"
      cluster_resolver.task_id = 0
      strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
          cluster_resolver=cluster_resolver)
  
      # On worker 1
      cluster_resolver = KubernetesClusterResolver(
          {"worker": ["job-name=worker-cluster-a", "job-name=worker-cluster-b"]})
      cluster_resolver.task_type = "worker"
      cluster_resolver.task_id = 1
      strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
          cluster_resolver=cluster_resolver)
      ```
    """
  
    def __init__(self,
                 job_to_label_mapping=None,
                 tf_server_port=8470,
                 rpc_layer='grpc',
                 override_client=None):
      """Initializes a new KubernetesClusterResolver.
  
      This initializes a new Kubernetes ClusterResolver. The ClusterResolver
      will attempt to talk to the Kubernetes master to retrieve all the instances
      of pods matching a label selector.
  
      Args:
        job_to_label_mapping: A mapping of TensorFlow jobs to label selectors.
          This allows users to specify many TensorFlow jobs in one Cluster
          Resolver, and each job can have pods belong with different label
          selectors. For example, a sample mapping might be
          ```
          {'worker': ['job-name=worker-cluster-a', 'job-name=worker-cluster-b'],
           'ps': ['job-name=ps-1', 'job-name=ps-2']}
          ```
        tf_server_port: The port the TensorFlow server is listening on.
        rpc_layer: (Optional) The RPC layer TensorFlow should use to communicate
          between tasks in Kubernetes. Defaults to 'grpc'.
        override_client: The Kubernetes client (usually automatically retrieved
          using `from kubernetes import client as k8sclient`). If you pass this
          in, you are responsible for setting Kubernetes credentials manually.
  
      Raises:
        ImportError: If the Kubernetes Python client is not installed and no
          `override_client` is passed in.
        RuntimeError: If autoresolve_task is not a boolean or a callable.
      """
      try:
        from kubernetes import config as k8sconfig  # pylint: disable=g-import-not-at-top
  
        k8sconfig.load_kube_config()
      except ImportError:
        if not override_client:
          raise ImportError('The Kubernetes Python client must be installed '
                            'before using the Kubernetes Cluster Resolver. '
                            'To install the Kubernetes Python client, run '
                            '`pip install kubernetes` on your command line.')
  
      if not job_to_label_mapping:
        job_to_label_mapping = {'worker': ['job-name=tensorflow']}
  
      self._job_to_label_mapping = job_to_label_mapping
      self._tf_server_port = tf_server_port
      self._override_client = override_client
  
      self.task_type = None
      self.task_id = None
      self.rpc_layer = rpc_layer
  
    def master(self, task_type=None, task_id=None, rpc_layer=None):
      """Returns the master address to use when creating a session.
  
      You must have set the task_type and task_id object properties before
      calling this function, or pass in the `task_type` and `task_id`
      parameters when using this function. If you do both, the function parameters
      will override the object properties.
  
      Note: this is only useful for TensorFlow 1.x.
  
      Args:
        task_type: (Optional) The type of the TensorFlow task of the master.
        task_id: (Optional) The index of the TensorFlow task of the master.
        rpc_layer: (Optional) The RPC protocol for the given cluster.
  
      Returns:
        The name or URL of the session master.
      """
      task_type = task_type if task_type is not None else self.task_type
      task_id = task_id if task_id is not None else self.task_id
  
      if task_type is not None and task_id is not None:
        return format_master_url(
            self.cluster_spec().task_address(task_type, task_id),
            rpc_layer or self.rpc_layer)
  
      return ''
  
    def cluster_spec(self):
      """Returns a ClusterSpec object based on the latest info from Kubernetes.
  
      We retrieve the information from the Kubernetes master every time this
      method is called.
  
      Returns:
        A ClusterSpec containing host information returned from Kubernetes.
  
      Raises:
        RuntimeError: If any of the pods returned by the master is not in the
          `Running` phase.
      """
      if self._override_client:
        client = self._override_client
      else:
        from kubernetes import config as k8sconfig  # pylint: disable=g-import-not-at-top
        from kubernetes import client as k8sclient  # pylint: disable=g-import-not-at-top
  
        k8sconfig.load_kube_config()
        client = k8sclient.CoreV1Api()
  
      cluster_map = {}
  
      for tf_job in self._job_to_label_mapping:
        all_pods = []
        for selector in self._job_to_label_mapping[tf_job]:
          ret = client.list_pod_for_all_namespaces(label_selector=selector)
          selected_pods = []
  
          # Sort the list by the name to make sure it doesn't change call to call.
          for pod in sorted(ret.items, key=lambda x: x.metadata.name):
            if pod.status.phase == 'Running':
              selected_pods.append(
                  '%s:%s' % (pod.status.host_ip, self._tf_server_port))
            else:
              raise RuntimeError('Pod "%s" is not running; phase: "%s"' %
                                 (pod.metadata.name, pod.status.phase))
          all_pods.extend(selected_pods)
        cluster_map[tf_job] = all_pods
  
      return server_lib.ClusterSpec(cluster_map)
   # Copyright 2020 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.
  # ==============================================================================
  """Mid level API for TPU Embeddings."""
  
  import functools
  from typing import Any, Callable, Dict, Iterable, List, Optional, Text, Tuple, Union
  
  from absl import logging
  
  from tensorflow.core.framework import attr_value_pb2
  from tensorflow.core.protobuf.tpu import tpu_embedding_configuration_pb2
  from tensorflow.python.distribute import device_util
  from tensorflow.python.distribute import distribute_utils
  from tensorflow.python.distribute import distribution_strategy_context
  from tensorflow.python.distribute import sharded_variable
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import device as tf_device
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.framework.tensor_shape import TensorShape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import embedding_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import sparse_ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables as tf_variables
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow.python.saved_model import save_context
  from tensorflow.python.tpu import tpu
  from tensorflow.python.tpu import tpu_embedding_v2_utils
  from tensorflow.python.tpu.ops import tpu_ops
  from tensorflow.python.training.saving import saveable_hook
  from tensorflow.python.training.tracking import base
  from tensorflow.python.training.tracking import tracking
  from tensorflow.python.types import core
  from tensorflow.python.types import internal as internal_types
  from tensorflow.python.util import compat
  from tensorflow.python.util import nest
  from tensorflow.python.util import tf_inspect
  from tensorflow.python.util.tf_export import tf_export
  
  
  _HOOK_KEY = "TPUEmbedding_saveable"
  _NAME_KEY = "_tpu_embedding_layer"
  
  
  class TPUShardedVariable(sharded_variable.ShardedVariableMixin):
    """A ShardedVariable class for TPU."""
  
    @property
    def _in_graph_mode(self):
      return self.variables[0]._in_graph_mode  # pylint: disable=protected-access
  
  
  def _add_key_attr(op, name):
    op._set_attr(_NAME_KEY, attr_value_pb2.AttrValue(s=compat.as_bytes(name)))  # pylint: disable=protected-access
  
  
  @tf_export("tpu.experimental.embedding.TPUEmbedding")
  class TPUEmbedding(tracking.AutoTrackable):
    """The TPUEmbedding mid level API.
  
    NOTE: When instantiated under a TPUStrategy, this class can only be created
    once per call to `tf.tpu.experimental.initialize_tpu_system`. If you wish to
    re-initialize the embedding engine you must re-initialize the tpu as well.
    Doing this will clear any variables from TPU, so ensure you have checkpointed
    before you do this. If a further instances of the class are needed,
    set the `initialize_tpu_embedding` argument to `False`.
  
    This class can be used to support training large embeddings on TPU. When
    creating an instance of this class, you must specify the complete set of
    tables and features you expect to lookup in those tables. See the
    documentation of `tf.tpu.experimental.embedding.TableConfig` and
    `tf.tpu.experimental.embedding.FeatureConfig` for more details on the complete
    set of options. We will cover the basic usage here.
  
    NOTE: multiple `FeatureConfig` objects can use the same `TableConfig` object,
    allowing different features to share the same table:
  
    ```python
    table_config_one = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
    table_config_two = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
    feature_config = {
        'feature_one': tf.tpu.experimental.embedding.FeatureConfig(
            table=table_config_one),
        'feature_two': tf.tpu.experimental.embedding.FeatureConfig(
            table=table_config_one),
        'feature_three': tf.tpu.experimental.embedding.FeatureConfig(
            table=table_config_two)}
    ```
  
    There are two modes under which the `TPUEmbedding` class can used. This
    depends on if the class was created under a `TPUStrategy` scope or not.
  
    Under `TPUStrategy`, we allow access to the method `enqueue`, `dequeue` and
    `apply_gradients`. We will show examples below of how to use these to train
    and evaluate your model. Under CPU, we only access to the `embedding_tables`
    property which allow access to the embedding tables so that you can use them
    to run model evaluation/prediction on CPU.
  
    First lets look at the `TPUStrategy` mode. Initial setup looks like:
  
    ```python
    strategy = tf.distribute.TPUStrategy(...)
    with strategy.scope():
      embedding = tf.tpu.experimental.embedding.TPUEmbedding(
          feature_config=feature_config,
          optimizer=tf.tpu.experimental.embedding.SGD(0.1))
    ```
  
    When creating a distributed dataset that is to be passed to the enqueue
    operation a special input option must be specified:
  
    ```python
    distributed_dataset = (
        strategy.distribute_datasets_from_function(
            dataset_fn=...,
            options=tf.distribute.InputOptions(
                experimental_fetch_to_device=False))
    dataset_iterator = iter(distributed_dataset)
    ```
  
    Different feature inputs can have different shapes. For dense and sparse
    tensor, rank 2 and above is supported. For ragged tensor, although only rank 2
    is supported, you can specify the output shape to be rank 2 and above. The
    output shape specified in the FeatureConfig has the first priority. The input
    shape passed in build method has second priority and the input shapes
    auto detected from input feature has the lowest priority. The latter two will
    be converted to output shapes by omitting the last dimension. If the lower
    priority one has output shapes which don't match the former one. A ValueError
    will be raised. Only when the former one has undefined output shapes, the
    latter one can override.
  
    NOTE: All batches passed to the layer can have different input shapes. But
    these input shapes need to match with the output shapes set by either
    `FeatureConfig` or build method except for ragged tensor. Only 2D
    ragged tensor with output shape set to higher dimensions is allowed as
    long as the total number of elements matches. All subsequent calls must have
    the same input shapes. In the event that the input shapes cannot be
    automatically determined by the enqueue method, you must call
    the build method with the input shapes or provide output shapes in the
    `FeatureConfig` to initialize the layer.
  
    To use this API on TPU you should use a custom training loop. Below is an
    example of a training and evaluation step:
  
    ```python
    @tf.function
    def training_step(dataset_iterator, num_steps):
      def tpu_step(tpu_features):
        with tf.GradientTape() as tape:
          activations = embedding.dequeue()
          tape.watch(activations)
          model_output = model(activations)
          loss = ...  # some function of labels and model_output
  
        embedding_gradients = tape.gradient(loss, activations)
        embedding.apply_gradients(embedding_gradients)
        # Insert your model gradient and optimizer application here
  
      for _ in tf.range(num_steps):
        embedding_features, tpu_features = next(dataset_iterator)
        embedding.enqueue(embedding_features, training=True)
        strategy.run(tpu_step, args=(tpu_features, ))
  
    @tf.function
    def evalution_step(dataset_iterator, num_steps):
      def tpu_step(tpu_features):
        activations = embedding.dequeue()
        model_output = model(activations)
        # Insert your evaluation code here.
  
      for _ in tf.range(num_steps):
        embedding_features, tpu_features = next(dataset_iterator)
        embedding.enqueue(embedding_features, training=False)
        strategy.run(tpu_step, args=(tpu_features, ))
    ```
  
    NOTE: The calls to `enqueue` have `training` set to `True` when
    `embedding.apply_gradients` is used and set to `False` when
    `embedding.apply_gradients` is not present in the function. If you don't
    follow this pattern you may cause an error to be raised or the tpu may
    deadlock.
  
    In the above examples, we assume that the user has a dataset which returns
    a tuple where the first element of the tuple matches the structure of what
    was passed as the `feature_config` argument to the object initializer. Also we
    utilize `tf.range` to get a `tf.while_loop` in order to increase performance.
  
    When checkpointing your model, you should include your
    `tf.tpu.experimental.embedding.TPUEmbedding` object in the checkpoint. It is a
    trackable object and saving it will save the embedding tables and their
    optimizer slot variables:
  
    ```python
    checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
    checkpoint.save(...)
    ```
  
    On CPU, only the `embedding_table` property is usable. This will allow you to
    restore a checkpoint to the object and have access to the table variables:
  
    ```python
    model = model_fn(...)
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=feature_config,
        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
    checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
    checkpoint.restore(...)
  
    tables = embedding.embedding_tables
    ```
  
    You can now use table in functions like `tf.nn.embedding_lookup` to perform
    your embedding lookup and pass to your model.
  
    """
  
    def __init__(
        self,
        feature_config: Union[tpu_embedding_v2_utils.FeatureConfig, Iterable],  # pylint:disable=g-bare-generic
        optimizer: Optional[tpu_embedding_v2_utils._Optimizer],  # pylint:disable=protected-access
        pipeline_execution_with_tensor_core: bool = False):
      """Creates the TPUEmbedding mid level API object.
  
      ```python
      strategy = tf.distribute.TPUStrategy(...)
      with strategy.scope():
        embedding = tf.tpu.experimental.embedding.TPUEmbedding(
            feature_config=tf.tpu.experimental.embedding.FeatureConfig(
                table=tf.tpu.experimental.embedding.TableConfig(
                    dim=...,
                    vocabulary_size=...)))
      ```
  
      Args:
        feature_config: A nested structure of
          `tf.tpu.experimental.embedding.FeatureConfig` configs.
        optimizer: An instance of one of `tf.tpu.experimental.embedding.SGD`,
          `tf.tpu.experimental.embedding.Adagrad` or
          `tf.tpu.experimental.embedding.Adam`. When not created under
          TPUStrategy may be set to None to avoid the creation of the optimizer
          slot variables, useful for optimizing memory consumption when exporting
          the model for serving where slot variables aren't needed.
        pipeline_execution_with_tensor_core: If True, the TPU embedding
          computations will overlap with the TensorCore computations (and hence
          will be one step old). Set to True for improved performance.
  
      Raises:
        ValueError: If optimizer is not one of tf.tpu.experimental.embedding.(SGD,
        Adam or Adagrad) or None when created under a TPUStrategy.
      """
      self._strategy = distribution_strategy_context.get_strategy()
      self._using_tpu = isinstance(self._strategy, (tpu_strategy.TPUStrategy,
                                                    tpu_strategy.TPUStrategyV2))
      self._pipeline_execution_with_tensor_core = (
          pipeline_execution_with_tensor_core)
  
      self._feature_config = feature_config
      self._output_shapes = []
      for feature in nest.flatten(feature_config):
        self._output_shapes.append(feature.output_shape)
  
      # The TPU embedding ops are slightly inconsistent with how they refer to
      # tables:
      # * The enqueue op takes a parallel list of tensors for input, one of those
      #   is the table id for the feature which matches the integer index of the
      #   table in the proto created by _create_config_proto().
      # * The recv_tpu_embedding_activations op emits lookups per table in the
      #   order from the config proto.
      # * The send_tpu_embedding_gradients expects input tensors to be per table
      #   in the same order as the config proto.
      # * Per optimizer load and retrieve ops are specified per table and take the
      #   table name rather than the table id.
      # Thus we must fix a common order to tables and ensure they have unique
      # names.
  
      # Set table order here to the order of the first occurence of the table in a
      # feature provided by the user. The order of this struct must be fixed
      # to provide the user with deterministic behavior over multiple
      # instantiations.
      self._table_config = []
      for feature in nest.flatten(feature_config):
        if feature.table not in self._table_config:
          self._table_config.append(feature.table)
  
      # Ensure tables have unique names. Also error check the optimizer as we
      # specifically don't do that in the TableConfig class to allow high level
      # APIs that are built on this to use strings/other classes to represent
      # optimizers (before they are passed to this class).
      table_names = []
      for i, table in enumerate(self._table_config):
        if table.optimizer is None:
          # TODO(bfontain) Should we allow some sort of optimizer merging here?
          table.optimizer = optimizer
        if ((table.optimizer is not None or self._using_tpu) and
            not isinstance(table.optimizer, tpu_embedding_v2_utils._Optimizer)):  # pylint: disable=protected-access
          raise ValueError("{} is an unsupported optimizer class. Please pass an "
                           "instance of one of the optimizer classes under "
                           "tf.tpu.experimental.embedding.".format(
                               type(table.optimizer)))
        if table.name is None:
          table.name = "table_{}".format(i)
        if table.name in table_names:
          raise ValueError("Tables must have a unique name. "
                           f"Multiple tables with name {table.name} found.")
        table_names.append(table.name)
  
      if self._using_tpu:
        # Extract a list of callable learning rates also in fixed order. Each
        # table in the confix proto will get a index into this list and we will
        # pass this list in the same order after evaluation to the
        # send_tpu_embedding_gradients op.
        self._dynamic_learning_rates = list({
            table.optimizer.learning_rate for table in self._table_config if
            callable(table.optimizer.learning_rate)})
  
        # We need to list of host devices for the load/retrieve operations.
        self._hosts = get_list_of_hosts(self._strategy)
  
      self._built = False
      self._verify_output_shapes_on_enqueue = True
  
    def build(self, per_replica_input_shapes=None, per_replica_batch_size=None):  # pylint:disable=g-bare-generic
      """Create the underlying variables and initializes the TPU for embeddings.
  
      This method creates the underlying variables (including slot variables). If
      created under a TPUStrategy, this will also initialize the TPU for
      embeddings.
  
      This function will automatically get called by enqueue, which will try to
      determine your output shapes. If this fails, you must manually
      call this method before you call enqueue.
  
      Args:
        per_replica_input_shapes: A nested structure of The per replica input
          shapes that matches the structure of the feature config. The input
          shapes should be the same as the input shape of the feature (except for
          ragged tensor) Note that it is fixed and the same per replica input
          shapes must be used for both training and evaluation. If you want to
          calculate this from the global input shapes, you can use
          `num_replicas_in_sync` property of your strategy object. May be set to
          None if not created under a TPUStrategy.
        per_replica_batch_size: (Deprecated) The per replica batch size that you
          intend to use. Note that is fixed and the same batch size must be used
          for both training and evaluation. If you want to calculate this from the
          global batch size, you can use `num_replicas_in_sync` property of your
          strategy object. May be set to None if not created under a TPUStrategy.
  
      Raises:
        ValueError: If per_replica_input_shapes is inconsistent with the output
        shapes stored in the feature config or the output shapes get from the
        input shapes are not fully defined.
        RuntimeError: If tpu embedding is already initialized on TPU.
      """
      if self._built:
        return
  
      if self._using_tpu:
        # If the tpu embedding is already initialized on TPU, raise runtime error.
        # Below logic is not added in `initialize_system_for_tpu_embedding`
        # because doing exception control flow in graph mode is difficult.
        if tpu_ops.is_tpu_embedding_initialized():
          raise RuntimeError(
              "TPU is already initialized for embeddings. This may be caused by "
              "using multiple TPUEmbedding instances in a TPU scope which is "
              "unsupported")
        self._get_and_update_output_shapes_from_input(per_replica_input_shapes,
                                                      per_replica_batch_size)
  
        self._config_proto = self._create_config_proto()
  
        logging.info("Initializing TPU Embedding engine.")
        tpu_embedding_v2_utils.log_tpu_embedding_configuration(self._config_proto)
  
        @def_function.function
        def load_config():
          tpu.initialize_system_for_tpu_embedding(self._config_proto)
  
        load_config()
        logging.info("Done initializing TPU Embedding engine.")
  
      # Create and load variables and slot variables into the TPU.
      # Note that this is a dict of dicts. Keys to the first dict are table names.
      # We would prefer to use TableConfigs, but then these variables won't be
      # properly tracked by the tracking API.
      self._variables = self._create_variables_and_slots()
  
      self._built = True
  
      # This is internally conditioned self._built and self._using_tpu
      self._load_variables()
  
    def _maybe_build(self,
                     output_shapes: Optional[Union[List[int], Iterable]] = None):  # pylint:disable=g-bare-generic
      if not self._built:
        # This can be called while tracing a function, so we wrap the
        # initialization code with init_scope so it runs eagerly, this means that
        # it will not be included the function graph generated by tracing so that
        # we can be sure that we only initialize the TPU for embeddings exactly
        # once.
        with ops.init_scope():
          self.build(output_shapes)
  
    def _get_and_update_output_shapes_from_input(
        self,
        per_replica_input_shapes: Optional[List[TensorShape]] = None,
        per_replica_batch_size: Optional[int] = None):
      """Get and update the per replica output shapes from the input."""
      per_replica_output_shapes = None
      if per_replica_batch_size and per_replica_input_shapes is None:
        logging.warning(
            "per_replica_batch_size argument will be deprecated, please specify "
            "all the input shapes using per_replica_input_shapes argument.")
        per_replica_output_shapes = self._get_output_shapes_from_batch_size(
            per_replica_batch_size)
  
      # Update the input shapes if provided.
      if per_replica_input_shapes is not None:
        if isinstance(per_replica_input_shapes, int):
          logging.warning(
              "Passing batch size to per_replica_input_shapes argument will be"
              " deprecated, please specify all the input shapes using"
              " per_replica_input_shapes argument.")
          per_replica_output_shapes = self._get_output_shapes_from_batch_size(
              per_replica_input_shapes)
        else:
          nest.assert_same_structure(
              nest.flatten(per_replica_input_shapes),
              nest.flatten(self._feature_config))
  
          # Convert the nested structure to list.
          per_replica_input_shapes = nest.flatten(per_replica_input_shapes)
  
          per_replica_output_shapes = self._get_output_shapes_from_input_shapes(
              per_replica_input_shapes)
  
      if per_replica_output_shapes is not None:
  
        # Check the output shapes with existing output shapes setting.
        self._check_output_shapes(per_replica_output_shapes)
  
        # Update the output shapes with existing output shapes setting.
        # This is necessary Because the output shapes might be missing from
        # the feature config, the usr can set it:
        #  1. calling the build method
        #  2. output shapes auto detected when calling the dequeue method for
        #     for the first time. The dequeue method will call build method
        #     with the output shapes.
        # Either these two situations will lead to an update to the existing
        # output shapes.
        self._update_output_shapes(per_replica_output_shapes)
  
      # Check if the output shapes are fully defined. This is required in order
      # to set them in the feature descriptor field of the tpu embedding config
      # proto.
      self._check_output_shapes_fully_defined()
  
    def _get_output_shapes_from_input_shapes(
        self, input_shapes: List[TensorShape]) -> List[TensorShape]:
      """Get output shapes from the flattened input shapes list."""
      output_shapes = []
      for input_shape, feature in zip(input_shapes,
                                      nest.flatten(self._feature_config)):
        if input_shape.rank is None or input_shape.rank < 1:
          raise ValueError(
              "Received input tensor of shape {}. Rank must be 1 and above"
              .format(input_shape))
        # Update the input shape with the max sequence length. Only update when
        # 1. Input feature is 2D ragged or sparse tensor.
        # 2. Output shape is not set in the feature config and the max sequence
        #    length is set.
        if (len(input_shape) == 2 and input_shape[-1] != 1 and
            not feature.output_shape and feature.max_sequence_length > 0):
          input_shape_list = input_shape.as_list()
          input_shape_list.insert(
              len(input_shape_list) - 1, feature.max_sequence_length)
          input_shape = TensorShape(input_shape_list)
        if input_shape.rank == 1:
          output_shapes.append(input_shape)
        else:
          output_shapes.append(input_shape[:-1])
      return output_shapes
  
    @property
    def embedding_tables(
        self
    ) -> Dict[tpu_embedding_v2_utils.TableConfig, tf_variables.Variable]:
      """Returns a dict of embedding tables, keyed by `TableConfig`.
  
      This property only works when the `TPUEmbedding` object is created under a
      non-TPU strategy. This is intended to be used to for CPU based lookup when
      creating a serving checkpoint.
  
      Returns:
        A dict of embedding tables, keyed by `TableConfig`.
  
      Raises:
        RuntimeError: If object was created under a `TPUStrategy`.
      """
      # We don't support returning tables on TPU due to their sharded nature and
      # the fact that when using a TPUStrategy:
      # 1. Variables are stale and are only updated when a checkpoint is made.
      # 2. Updating the variables won't affect the actual tables on the TPU.
      if self._using_tpu:
        if save_context.in_save_context():
          return {table: self._variables[table.name]["parameters"].variables[0]
                  for table in self._table_config}
        raise RuntimeError("Unable to retrieve embedding tables when using a TPU "
                           "strategy. If you need access, save your model, "
                           "create this object under a CPU strategy and restore.")
  
      self._maybe_build(None)
  
      # Only return the tables and not the slot variables. On CPU this are honest
      # tf.Variables.
      return {table: self._variables[table.name]["parameters"]
              for table in self._table_config}
  
    def _create_config_proto(
        self
    ) -> tpu_embedding_configuration_pb2.TPUEmbeddingConfiguration:
      """Creates the TPUEmbeddingConfiguration proto.
  
      This proto is used to initialize the TPU embedding engine.
  
      Returns:
        A TPUEmbeddingConfiguration proto.
      """
  
      config_proto = tpu_embedding_configuration_pb2.TPUEmbeddingConfiguration()
  
      # Map each callable dynamic learning rate to its in index in the list.
      # The learning rate index is the index of the dynamic learning rate for this
      # table (if it exists) in the list we created at initialization. We don't
      # simply create one learning rate index per table as this has extremely bad
      # performance characteristics. The more separate optimization configurations
      # we have, the worse the performance will be.
      learning_rate_index = {r: i for i, r in enumerate(
          self._dynamic_learning_rates)}
  
      for table in self._table_config:
        table_descriptor = config_proto.table_descriptor.add()
        table_descriptor.name = table.name
  
        # For small tables, we pad to the number of hosts so that at least one
        # id will be assigned to each host.
        table_descriptor.vocabulary_size = max(table.vocabulary_size,
                                               self._strategy.extended.num_hosts)
        table_descriptor.dimension = table.dim
  
        parameters = table_descriptor.optimization_parameters
  
        # We handle the learning rate separately here and don't allow the
        # optimization class to handle this, as it doesn't know about dynamic
        # rates.
        if callable(table.optimizer.learning_rate):
          parameters.learning_rate.dynamic.tag = (
              learning_rate_index[table.optimizer.learning_rate])
        else:
          parameters.learning_rate.constant = table.optimizer.learning_rate
  
        # Use optimizer to handle the rest of the parameters.
        table.optimizer._set_optimization_parameters(parameters)  # pylint: disable=protected-access
  
      table_to_id = {table: i for i, table in enumerate(self._table_config)}
  
      # Set feature descriptor field in the config proto.
      for feature, output_shape in zip(
          nest.flatten(self._feature_config), self._output_shapes):
        feature_descriptor = config_proto.feature_descriptor.add()
  
        if feature.name:
          feature_descriptor.name = feature.name
  
        feature_descriptor.table_id = table_to_id[feature.table]
        # The input shape of the feature is the actual shape of the input tensor
        # except the last dimension because the last dimension will always be
        # reduced.
        feature_descriptor.input_shape.extend(output_shape.as_list())
  
      # Always set mode to training, we override the mode during enqueue.
      config_proto.mode = (
          tpu_embedding_configuration_pb2.TPUEmbeddingConfiguration.TRAINING)
  
      config_proto.num_hosts = self._strategy.extended.num_hosts
      config_proto.num_tensor_cores = self._strategy.num_replicas_in_sync
  
      # TODO(bfontain): Allow users to pick MOD for the host sharding.
      config_proto.sharding_strategy = (
          tpu_embedding_configuration_pb2.TPUEmbeddingConfiguration.DIV_DEFAULT)
      config_proto.pipeline_execution_with_tensor_core = (
          self._pipeline_execution_with_tensor_core)
  
      return config_proto
  
    def apply_gradients(self, gradients, name: Optional[Text] = None):
      """Applies the gradient update to the embedding tables.
  
      If a gradient of `None` is passed in any position of the nested structure,
      then an gradient update with a zero gradient is applied for that feature.
      For optimizers like SGD or Adagrad, this is the same as applying no update
      at all. For lazy Adam and other sparsely applied optimizers with decay,
      ensure you understand the effect of applying a zero gradient.
  
      ```python
      strategy = tf.distribute.TPUStrategy(...)
      with strategy.scope():
        embedding = tf.tpu.experimental.embedding.TPUEmbedding(...)
  
      distributed_dataset = (
          strategy.distribute_datasets_from_function(
              dataset_fn=...,
              options=tf.distribute.InputOptions(
                  experimental_fetch_to_device=False))
      dataset_iterator = iter(distributed_dataset)
  
      @tf.function
      def training_step():
        def tpu_step(tpu_features):
          with tf.GradientTape() as tape:
            activations = embedding.dequeue()
            tape.watch(activations)
  
            loss = ... #  some computation involving activations
  
          embedding_gradients = tape.gradient(loss, activations)
          embedding.apply_gradients(embedding_gradients)
  
        embedding_features, tpu_features = next(dataset_iterator)
        embedding.enqueue(embedding_features, training=True)
        strategy.run(tpu_step, args=(tpu_features, ))
  
      training_step()
      ```
  
      Args:
        gradients: A nested structure of gradients, with structure matching the
          `feature_config` passed to this object.
        name: A name for the underlying op.
  
      Raises:
        RuntimeError: If called when object wasn't created under a `TPUStrategy`
          or if not built (either by manually calling build or calling enqueue).
        ValueError: If a non-`tf.Tensor` non-`None` gradient is passed in, or a
          `tf.Tensor` of the incorrect shape is passed in. Also if
          the size of any sequence in `gradients` does not match corresponding
          sequence in `feature_config`.
        TypeError: If the type of any sequence in `gradients` does not match
          corresponding sequence in `feature_config`.
      """
      if not self._using_tpu:
        raise RuntimeError("apply_gradients is not valid when TPUEmbedding "
                           "object is not created under a TPUStrategy.")
  
      if not self._built:
        raise RuntimeError("apply_gradients called on unbuilt TPUEmbedding "
                           "object. Please either call enqueue first or manually "
                           "call the build method.")
  
      nest.assert_same_structure(self._feature_config, gradients)
      updated_gradients = []
      for (path, gradient), feature, output_shape in zip(
          nest.flatten_with_joined_string_paths(gradients),
          nest.flatten(self._feature_config), self._output_shapes):
        full_output_shape = list(output_shape) + [feature.table.dim]
        if gradient is not None and not isinstance(gradient, ops.Tensor):
          raise ValueError(
              f"found non-tensor type: {type(gradient)} at path {path}.")
        if gradient is not None:
          if gradient.shape != full_output_shape:
            raise ValueError("Found gradient of shape {} at path {}. Expected "
                             "shape {}.".format(gradient.shape, path,
                                                full_output_shape))
        else:
          # No gradient for this feature, since we must give a gradient for all
          # features, pass in a zero tensor here. Note that this is not correct
          # for all optimizers.
          logging.warning(
              "No gradient passed for feature %s, sending zero "
              "gradient. This may not be correct behavior for certain "
              "optimizers like Adam.", path)
          gradient = array_ops.zeros(full_output_shape, dtype=dtypes.float32)
        # Some gradients can be passed with op which shape is not correctly set.
        # This ensures that the shape of the gradient is correctly set.
        updated_gradients.append(
            array_ops.reshape(gradient, shape=gradient.shape))
      op = tpu_ops.send_tpu_embedding_gradients(
          inputs=updated_gradients,
          learning_rates=[
              math_ops.cast(fn(), dtype=dtypes.float32)
              for fn in self._dynamic_learning_rates
          ],
          config=self._config_proto.SerializeToString())
  
      # Apply the name tag to the op.
      if name is not None:
        _add_key_attr(op, name)
  
    def dequeue(self, name: Optional[Text] = None):
      """Get the embedding results.
  
      Returns a nested structure of `tf.Tensor` objects, matching the structure of
      the `feature_config` argument to the `TPUEmbedding` class. The output shape
      of the tensors is `(*output_shape, dim)`, `dim` is the dimension of the
      corresponding `TableConfig`. For output_shape, there are three places where
      it can be set.
        1. FeatureConfig provided in the __init__ function.
        2. Per_replica_output_shapes by directly calling the build method
             after initializing the tpu embedding class.
        3. Auto detected from the shapes of the input feature.
      The priority of these places is the exact same order.
  
      ```python
      strategy = tf.distribute.TPUStrategy(...)
      with strategy.scope():
        embedding = tf.tpu.experimental.embedding.TPUEmbedding(...)
  
      distributed_dataset = (
          strategy.distribute_datasets_from_function(
              dataset_fn=...,
              options=tf.distribute.InputOptions(
                  experimental_fetch_to_device=False))
      dataset_iterator = iter(distributed_dataset)
  
      @tf.function
      def training_step():
        def tpu_step(tpu_features):
          with tf.GradientTape() as tape:
            activations = embedding.dequeue()
            tape.watch(activations)
  
            loss = ... #  some computation involving activations
  
          embedding_gradients = tape.gradient(loss, activations)
          embedding.apply_gradients(embedding_gradients)
  
        embedding_features, tpu_features = next(dataset_iterator)
        embedding.enqueue(embedding_features, training=True)
        strategy.run(tpu_step, args=(tpu_features, ))
  
      training_step()
      ```
  
      Args:
        name: A name for the underlying op.
  
      Returns:
        A nested structure of tensors, with the same structure as `feature_config`
      passed to this instance of the `TPUEmbedding` object.
  
      Raises:
        RuntimeError: If called when object wasn't created under a `TPUStrategy`
          or if not built (either by manually calling build or calling enqueue).
      """
      if not self._using_tpu:
        raise RuntimeError("dequeue is not valid when TPUEmbedding object is not "
                           "created under a TPUStrategy.")
  
      if not self._built:
        raise RuntimeError("dequeue called on unbuilt TPUEmbedding object. "
                           "Please either call enqueue first or manually call "
                           "the build method.")
  
      # The activations returned by this op are per feature.
      activations = tpu_ops.recv_tpu_embedding_activations(
          num_outputs=len(self._config_proto.feature_descriptor),
          config=self._config_proto.SerializeToString())
  
      # Apply the name tag to the op.
      if name is not None:
        _add_key_attr(activations[0].op, name)
  
      # Pack the list back into the same nested structure as the features.
      return nest.pack_sequence_as(self._feature_config, activations)
  
    def _create_variables_and_slots(
        self
    ) -> Dict[Text, Dict[Text, tf_variables.Variable]]:
      """Create variables for TPU embeddings.
  
      Note under TPUStrategy this will ensure that all creations happen within a
      variable creation scope of the sharded variable creator.
  
      Returns:
        A dict of dicts. The outer dict is keyed by the table names and the inner
        dicts are keyed by 'parameters' and the slot variable names.
      """
  
      def create_variables(table):
        """Create all variables."""
        variable_shape = (table.vocabulary_size, table.dim)
  
        def getter(name, shape, dtype, initializer, trainable):
          del shape
          # _add_variable_with_custom_getter clears the shape sometimes, so we
          # take the global shape from outside the getter.
          initial_value = functools.partial(initializer, variable_shape,
                                            dtype=dtype)
          return tf_variables.Variable(
              name=name,
              initial_value=initial_value,
              shape=variable_shape,
              dtype=dtype,
              trainable=trainable)
  
        def variable_creator(name, initializer, trainable=True):
          # use add_variable_with_custom_getter here so that we take advantage of
          # the checkpoint loading to allow restore before the variables get
          # created which avoids double initialization.
          return self._add_variable_with_custom_getter(
              name=name,
              initializer=initializer,
              shape=variable_shape,
              dtype=dtypes.float32,
              getter=getter,
              trainable=trainable)
  
        parameters = variable_creator(table.name, table.initializer,
                                      trainable=not self._using_tpu)
  
        def slot_creator(name, initializer):
          return variable_creator(table.name + "/" + name,
                                  initializer,
                                  False)
  
        if table.optimizer is not None:
          slot_vars = table.optimizer._create_slots(parameters, slot_creator)  # pylint: disable=protected-access
        else:
          slot_vars = {}
        slot_vars["parameters"] = parameters
        return slot_vars
  
      # Store tables based on name rather than TableConfig as we can't track
      # through dicts with non-string keys, i.e. we won't be able to save.
      variables = {}
      for table in self._table_config:
        if not self._using_tpu:
          variables[table.name] = create_variables(table)
        else:
          with variable_scope.variable_creator_scope(
              make_sharded_variable_creator(self._hosts)):
            variables[table.name] = create_variables(table)
  
      return variables
  
    def _load_variables(self):
      # Only load the variables if we are:
      # 1) Using TPU
      # 2) Variables are created
      # 3) Not in save context (except if running eagerly)
      if self._using_tpu and self._built and not (
          not context.executing_eagerly() and save_context.in_save_context()):
        _load_variables_impl(self._config_proto.SerializeToString(),
                             self._hosts,
                             self._variables,
                             self._table_config)
  
    def _retrieve_variables(self):
      # Only retrieve the variables if we are:
      # 1) Using TPU
      # 2) Variables are created
      # 3) Not in save context (except if running eagerly)
      if self._using_tpu and self._built and not (
          not context.executing_eagerly() and save_context.in_save_context()):
        _retrieve_variables_impl(self._config_proto.SerializeToString(),
                                 self._hosts,
                                 self._variables,
                                 self._table_config)
  
    def _gather_saveables_for_checkpoint(
        self
    ) -> Dict[Text, Callable[[Text], "TPUEmbeddingSaveable"]]:
      """Overrides default Trackable implementation to add load/retrieve hook."""
      # This saveable should be here in both TPU and CPU checkpoints, so when on
      # CPU, we add the hook with no functions.
      # TODO(bfontain): Update restore logic in saver so that these hooks are
      # always executed. Once that is done, we can output an empty list when on
      # CPU.
  
      def factory(name=_HOOK_KEY):
        return TPUEmbeddingSaveable(name, self._load_variables,
                                    self._retrieve_variables)
      return {_HOOK_KEY: factory}
  
    # Some helper functions for the below enqueue function.
    def _add_data_for_tensor(self, tensor, weight, indices, values, weights,
                             int_zeros, float_zeros, path):
      if weight is not None:
        raise ValueError(
            "Weight specified for dense input {}, which is not allowed. "
            "Weight will always be 1 in this case.".format(path))
      # For tensors, there are no indices and no weights.
      indices.append(int_zeros)
      values.append(math_ops.cast(array_ops.reshape(tensor, [-1]), dtypes.int64))
      weights.append(float_zeros)
  
    def _add_data_for_sparse_tensor(self, tensor, weight, indices, values,
                                    weights, int_zeros, float_zeros, path,
                                    feature):
      sample_indices = math_ops.cast(tensor.indices, dtypes.int32)
      if tensor.shape.rank == 2:
        if not feature.output_shape and feature.max_sequence_length > 0:
          # Add one dimension to the last axis.
          sample_indices = array_ops.pad(
              sample_indices, paddings=[[0, 0], [0, 1]])
      indices.append(sample_indices)
      values.append(math_ops.cast(tensor.values, dtypes.int64))
      # If we have weights they must be a SparseTensor.
      if weight is not None:
        if not isinstance(weight, sparse_tensor.SparseTensor):
          raise ValueError("Weight for {} is type {} which does not match "
                           "type input which is SparseTensor.".format(
                               path, type(weight)))
        weights.append(math_ops.cast(weight.values, dtypes.float32))
      else:
        weights.append(float_zeros)
  
    def _add_data_for_ragged_tensor(self, tensor, weight, row_splits, values,
                                    weights, int_zeros, float_zeros, path,
                                    feature):
      row_splits.append(math_ops.cast(tensor.row_splits, dtypes.int32))
      values.append(math_ops.cast(tensor.values, dtypes.int64))
      # If we have weights they must be a RaggedTensor.
      if weight is not None:
        if not isinstance(weight, ragged_tensor.RaggedTensor):
          raise ValueError("Weight for {} is type {} which does not match "
                           "type input which is RaggedTensor.".format(
                               path, type(weight)))
        weights.append(math_ops.cast(weight.values, dtypes.float32))
      else:
        weights.append(float_zeros)
  
    def _generate_enqueue_op(
        self,
        flat_inputs: List[internal_types.NativeObject],
        flat_weights: List[Optional[internal_types.NativeObject]],
        flat_features: List[tpu_embedding_v2_utils.FeatureConfig],
        device_ordinal: int,
        mode_override: Text
    ) -> ops.Operation:
      """Outputs a the enqueue op given the inputs and weights.
  
      Args:
        flat_inputs: A list of input tensors.
        flat_weights: A list of input weights (or None) of the same length as
          flat_inputs.
        flat_features: A list of FeatureConfigs of the same length as flat_inputs.
        device_ordinal: The device to create the enqueue op for.
        mode_override: A tensor containing the string "train" or "inference".
  
      Returns:
        The enqueue op.
      """
      # Combiners are per table, list in the same order as the table order.
      combiners = [table.combiner for table in self._table_config]
  
      # These parallel arrays will be the inputs to the enqueue op.
      # sample_indices for sparse, row_splits for ragged.
      indices_or_row_splits = []
      values = []
      weights = []
  
      # We have to supply a empty/zero tensor in a list position where we don't
      # have data (e.g. indices for standard Tensor input, weight when no weight
      # is specified). We create one op here per call, so that we reduce the
      # graph size.
      int_zeros = array_ops.zeros((0,), dtype=dtypes.int32)
      float_zeros = array_ops.zeros((0,), dtype=dtypes.float32)
  
      # In the following loop we insert casts so that everything is either int32
      # or float32. This is because op inputs which are lists of tensors must be
      # of the same type within the list. Moreover the CPU implementations of
      # these ops cast to these types anyway, so we don't lose any data by casting
      # early.
      for inp, weight, (path, feature) in zip(
          flat_inputs, flat_weights, flat_features):
        if isinstance(inp, ops.Tensor):
          self._add_data_for_tensor(inp, weight, indices_or_row_splits, values,
                                    weights, int_zeros, float_zeros, path)
        elif isinstance(inp, sparse_tensor.SparseTensor):
          self._add_data_for_sparse_tensor(inp, weight, indices_or_row_splits,
                                           values, weights, int_zeros,
                                           float_zeros, path, feature)
        elif isinstance(inp, ragged_tensor.RaggedTensor):
          self._add_data_for_ragged_tensor(inp, weight, indices_or_row_splits,
                                           values, weights, int_zeros,
                                           float_zeros, path, feature)
        else:
          raise ValueError("Input {} is of unknown type {}. Please only pass "
                           "Tensor, SparseTensor or RaggedTensor as input to "
                           "enqueue.".format(path, type(inp)))
  
      return tpu_ops.enqueue_tpu_embedding_arbitrary_tensor_batch(
          sample_indices_or_row_splits=indices_or_row_splits,
          embedding_indices=values,
          aggregation_weights=weights,
          mode_override=mode_override,
          device_ordinal=device_ordinal,
          combiners=combiners)
  
    def _raise_error_for_incorrect_control_flow_context(self):
      """Raises an error if we are not in the TPUReplicateContext."""
      # Do not allow any XLA control flow (i.e. control flow in between a
      # TPUStrategy's run call and the call to this function), as we can't
      # extract the enqueue from the head when in XLA control flow.
      graph = ops.get_default_graph()
      in_tpu_ctx = False
      while graph is not None:
        ctx = graph._get_control_flow_context()  # pylint: disable=protected-access
        while ctx is not None:
          if isinstance(ctx, tpu.TPUReplicateContext):
            in_tpu_ctx = True
            break
          ctx = ctx.outer_context
        if in_tpu_ctx:
          break
        graph = getattr(graph, "outer_graph", None)
      if graph != ops.get_default_graph() and in_tpu_ctx:
        raise RuntimeError(
            "Current graph {} does not match graph which contains "
            "TPUReplicateContext {}. This is most likely due to the fact that "
            "enqueueing embedding data is called inside control flow or a "
            "nested function inside `strategy.run`. This is not supported "
            "because outside compilation fails to extract the enqueue ops as "
            "head of computation.".format(ops.get_default_graph(), graph))
      return in_tpu_ctx
  
    def _raise_error_for_non_direct_inputs(self, features):
      """Checks all tensors in features to see if they are a direct input."""
  
      # expand_composites here is important: as composite tensors pass through
      # tpu.replicate, they get 'flattened' into their component tensors and then
      # repacked before being passed to the tpu function. In means that it is the
      # component tensors which are produced by an op with the
      # "_tpu_input_identity" attribute.
      for path, input_tensor in nest.flatten_with_joined_string_paths(
          features, expand_composites=True):
        if input_tensor.op.type == "Placeholder":
          continue
        try:
          is_input = input_tensor.op.get_attr("_tpu_input_identity")
        except ValueError:
          is_input = False
        if not is_input:
          raise ValueError(
              "Received input tensor {} which is the output of op {} (type {}) "
              "which does not have the `_tpu_input_identity` attr. Please "
              "ensure that the inputs to this layer are taken directly from "
              "the arguments of the function called by "
              "strategy.run. Two possible causes are: dynamic batch size "
              "support or you are using a keras layer and are not passing "
              "tensors which match the dtype of the `tf.keras.Input`s."
              "If you are triggering dynamic batch size support, you can "
              "disable it by passing tf.distribute.RunOptions("
              "experimental_enable_dynamic_batch_size=False) to the options "
              "argument of strategy.run().".format(path,
                                                   input_tensor.op.name,
                                                   input_tensor.op.type))
  
    def _raise_error_for_inputs_not_on_cpu(self, flat_inputs, flat_paths):
      """Checks all tensors in features to see are placed on the CPU."""
  
      def check_device(path, device_string):
        spec = tf_device.DeviceSpec.from_string(device_string)
        if spec.device_type == "TPU":
          raise ValueError(
              "Received input tensor {} which is on a TPU input device {}. Input "
              "tensors for TPU embeddings must be placed on the CPU. Please "
              "ensure that your dataset is prefetching tensors to the host by "
              "setting the 'experimental_fetch_to_device' option of the "
              "dataset distribution function. See the documentation of the "
              "enqueue method for an example.".format(path, device_string))
  
      # expand_composites here is important, we need to check the device of each
      # underlying tensor.
      for input_tensor, input_path in zip(flat_inputs, flat_paths):
        if nest.is_nested_or_composite(input_tensor):
          input_tensors = nest.flatten(input_tensor, expand_composites=True)
        else:
          input_tensors = [input_tensor]
        for t in input_tensors:
          if (t.op.type == "Identity" and
              t.op.inputs[0].op.type == "TPUReplicatedInput"):
            for tensor in t.op.inputs[0].op.inputs:
              check_device(input_path, tensor.device)
          else:
            check_device(input_path, t.device)
  
    def enqueue(
        self,
        features,
        weights=None,
        training: bool = True,
        name: Optional[Text] = None,
        device: Optional[Text] = None):
      """Enqueues id tensors for embedding lookup.
  
      This function enqueues a structure of features to be looked up in the
      embedding tables. We expect that the input shapes of each of the tensors in
      features matches the output shapes set via FeatureConfig or build method
      (if any). the output shapes will be auto detected based on the input shapes
      with the max_sequence_length or output shape setting in the FeatureConfig.
      Note that the output shapes is based on per replica batch size.
      If your input dataset is batched to the global batch size and you use
      `tf.distribute.TPUStrategy`'s `experimental_distribute_dataset`
      or if you use `distribute_datasets_from_function` and batch
      to the per core batch size computed by the context passed to your input
      function, the output shapes should match automatically.
  
      The auto detected the output shapes:
        1. For dense tensor, if rank 2 or above, make sure the tensor has last
           dimension as 1. The output shape will be the input shape excluding
           the last dimension.
        2. For sparse tensor, make sure the tensor has rank 2 and above.
             a. If feature config has max_sequence_length equals 0 or output shape
                set (the max_sequence_length setting will be ignored), the
                output shape will be the input shape excluding the last dimension.
             b. Otherwize if the tensor is rank 2, the output shape will be input
                shape  with last dimension set as max_sequence_length. If the
                tensor is above rank 2, the output shape will be the input shape
                excluding the last dimension and the last dimension of the output
                shape will be set to max_sequence_length.
        3. For ragged tensor, make sure the tensor has rank 2.
             a. If feature config has max_sequence_length equals 0 or output shape
                set (the max_sequence_length setting will be ignored), the
                output shape will be the input shape excluding the last dimension.
             b. Otherwise, the output shape will be the input shape excluding the
                last dimension and the last dimension of the output shape will be
                set to max_sequence_length.
  
      ```python
      strategy = tf.distribute.TPUStrategy(...)
      with strategy.scope():
        embedding = tf.tpu.experimental.embedding.TPUEmbedding(...)
  
      distributed_dataset = (
          strategy.distribute_datasets_from_function(
              dataset_fn=...,
              options=tf.distribute.InputOptions(
                  experimental_fetch_to_device=False))
      dataset_iterator = iter(distributed_dataset)
  
      @tf.function
      def training_step():
        def tpu_step(tpu_features):
          with tf.GradientTape() as tape:
            activations = embedding.dequeue()
            tape.watch(activations)
  
            loss = ... #  some computation involving activations
  
          embedding_gradients = tape.gradient(loss, activations)
          embedding.apply_gradients(embedding_gradients)
  
        embedding_features, tpu_features = next(dataset_iterator)
        embedding.enqueue(embedding_features, training=True)
        strategy.run(tpu_step, args=(tpu_features,))
  
      training_step()
      ```
  
      NOTE: You should specify `training=True` when using
      `embedding.apply_gradients` as above and `training=False` when not using
      `embedding.apply_gradients` (e.g. for frozen embeddings or when doing
      evaluation).
  
      For finer grained control, in the above example the line
  
      ```
        embedding.enqueue(embedding_features, training=True)
      ```
  
      may be replaced with
  
      ```
        per_core_embedding_features = self.strategy.experimental_local_results(
            embedding_features)
  
        def per_core_enqueue(ctx):
          core_id = ctx.replica_id_in_sync_group
          device = strategy.extended.worker_devices[core_id]
          embedding.enqueue(per_core_embedding_features[core_id],
                            device=device)
  
        strategy.experimental_distribute_values_from_function(
            per_core_queue_inputs)
      ```
  
      Args:
        features: A nested structure of `tf.Tensor`s, `tf.SparseTensor`s or
          `tf.RaggedTensor`s, with the same structure as `feature_config`. Inputs
          will be downcast to `tf.int32`. Only one type out of `tf.SparseTensor`
          or `tf.RaggedTensor` is supported per call.
        weights: If not `None`, a nested structure of `tf.Tensor`s,
          `tf.SparseTensor`s or `tf.RaggedTensor`s, matching the above, except
          that the tensors should be of float type (and they will be downcast to
          `tf.float32`). For `tf.SparseTensor`s we assume the `indices` are the
          same for the parallel entries from `features` and similarly for
          `tf.RaggedTensor`s we assume the row_splits are the same.
        training: Defaults to `True`. If `False`, enqueue the batch as inference
          batch (forward pass only). Do not call `apply_gradients` when this is
          `False` as this may lead to a deadlock.
         name: A name for the underlying op.
         device: The device name (e.g. '/task:0/device:TPU:2') where this batch
           should be enqueued. This should be set if and only if features is not a
           `tf.distribute.DistributedValues` and enqueue is not being called
           inside a TPU context (e.g. inside `TPUStrategy.run`).
  
      Raises:
        ValueError: When called inside a strategy.run call and input is not
          directly taken from the args of the `strategy.run` call. Also if
          the size of any sequence in `features` does not match corresponding
          sequence in `feature_config`. Similarly for `weights`, if not `None`.
          If input shapes of features is unequal or different from a previous
          call.
        RuntimeError: When called inside a strategy.run call and inside XLA
          control flow. If batch_size is not able to be determined and build was
          not called.
        TypeError: If the type of any sequence in `features` does not match
          corresponding sequence in `feature_config`. Similarly for `weights`, if
          not `None`.
      """
      if not self._using_tpu:
        raise RuntimeError("enqueue is not valid when TPUEmbedding object is not "
                           "created under a TPUStrategy.")
  
      in_tpu_context = self._raise_error_for_incorrect_control_flow_context()
  
      nest.assert_same_structure(self._feature_config, features)
  
      if not self._verify_output_shapes_on_enqueue:
        if not self._output_shapes or not self._built:
          raise ValueError(
              "Configured not to check output shapes on each enqueue() call; please "
              "ensure build() was called with output shapes to initialize "
              "the TPU for embeddings.")
      else:
        input_shapes = self._get_input_shapes(features, in_tpu_context)
  
        self._maybe_build(input_shapes)
        # If is already built, we still need to check if the output shapes matches
        # with the previous ones.
        self._check_output_shapes(
            self._get_output_shapes_from_input_shapes(input_shapes))
  
      flat_inputs = nest.flatten(features)
      flat_weights = [None] * len(flat_inputs)
      if weights is not None:
        nest.assert_same_structure(self._feature_config, weights)
        flat_weights = nest.flatten(weights)
      flat_features = nest.flatten_with_joined_string_paths(self._feature_config)
      flat_paths, _ = zip(*flat_features)
  
      self._raise_error_for_inputs_not_on_cpu(flat_inputs, flat_paths)
      # If we are in a tpu_context, automatically apply outside compilation.
      if in_tpu_context:
        self._raise_error_for_non_direct_inputs(features)
  
        def generate_enqueue_ops():
          """Generate enqueue ops for outside compilation."""
          # Note that we put array_ops.where_v2 rather than a python if so that
          # the op is explicitly create and the constant ops are both in the graph
          # even though we don't expect training to be a tensor (and thus generate
          # control flow automatically). This need to make it easier to re-write
          # the graph later if we need to fix which mode needs to be used.
          mode_override = array_ops.where_v2(training,
                                             constant_op.constant("train"),
                                             constant_op.constant("inference"))
          # Device ordinal is -1 here, a later rewrite will fix this once the op
          # is expanded by outside compilation.
          enqueue_op = self._generate_enqueue_op(
              flat_inputs, flat_weights, flat_features, device_ordinal=-1,
              mode_override=mode_override)
  
          # Apply the name tag to the op.
          if name is not None:
            _add_key_attr(enqueue_op, name)
  
          # Ensure that this op has outbound control flow, otherwise it won't be
          # executed.
          ops.get_default_graph().control_outputs.append(enqueue_op)
  
        tpu.outside_compilation(generate_enqueue_ops)
  
      elif device is None:
        mode_override = "train" if training else "inference"
        # We generate enqueue ops per device, so we need to gather the all
        # features for a single device in to a dict.
        # We rely here on the fact that the devices in the PerReplica value occur
        # in the same (standard) order as self._strategy.extended.worker_devices.
        enqueue_ops = []
        for replica_id in range(self._strategy.num_replicas_in_sync):
          replica_inputs = distribute_utils.select_replica(replica_id,
                                                           flat_inputs)
          replica_weights = distribute_utils.select_replica(replica_id,
                                                            flat_weights)
          tpu_device = self._strategy.extended.worker_devices[replica_id]
          # TPU devices string are like /job:worker/replica:0/task:0/device:TPU:0
          # the device ordinal is the last number
          device_ordinal = (
              tf_device.DeviceSpec.from_string(tpu_device).device_index)
  
          with ops.device(device_util.get_host_for_device(tpu_device)):
            enqueue_op = self._generate_enqueue_op(
                replica_inputs, replica_weights, flat_features,
                device_ordinal=device_ordinal, mode_override=mode_override)
  
            # Apply the name tag to the op.
            if name is not None:
              _add_key_attr(enqueue_op, name)
            enqueue_ops.append(enqueue_op)
        ops.get_default_graph().control_outputs.extend(enqueue_ops)
      else:
        mode_override = "train" if training else "inference"
        device_spec = tf_device.DeviceSpec.from_string(device)
        if device_spec.device_type != "TPU":
          raise ValueError(
              "Non-TPU device {} passed to enqueue.".format(device))
  
        with ops.device(device_util.get_host_for_device(device)):
          enqueue_op = self._generate_enqueue_op(
              flat_inputs, flat_weights, flat_features,
              device_ordinal=device_spec.device_index,
              mode_override=mode_override)
  
          # Apply the name tag to the op.
          if name is not None:
            _add_key_attr(enqueue_op, name)
          ops.get_default_graph().control_outputs.append(enqueue_op)
  
    def _get_input_shapes(self, tensors,
                          in_tpu_context: bool) -> List[TensorShape]:
      """Get the input shapes from the input tensor."""
      input_shapes = []
      for (path, maybe_tensor), feature in zip(
          nest.flatten_with_joined_string_paths(tensors),
          nest.flatten(self._feature_config)):
        if not in_tpu_context:
          tensor = distribute_utils.select_replica(0, maybe_tensor)
        else:
          tensor = maybe_tensor
  
        if isinstance(tensor, ops.Tensor):
          input_shapes.append(
              self._get_input_shape_for_tensor(tensor, feature, path))
        elif isinstance(tensor, sparse_tensor.SparseTensor):
          input_shapes.append(
              self._get_input_shape_for_sparse_tensor(tensor, feature, path))
        elif isinstance(tensor, ragged_tensor.RaggedTensor):
          input_shapes.append(
              self._get_input_shape_for_ragged_tensor(tensor, feature, path))
      return input_shapes
  
    def _get_input_shape_for_tensor(self, tensor, feature, path) -> TensorShape:
      """Get the input shape for the dense tensor."""
      shape = tensor.shape.as_list()
      if len(shape) < 1:
        raise ValueError("Only rank 1 and above dense tensor is supported,"
                         " find rank {} sparse tensor for input {}".format(
                             len(shape), path))
      if len(shape) > 1 and shape[-1] != 1:
        raise ValueError(
            "Rank 2 or above dense tensor should have last dimension as 1 "
            "as the last dimension will always be reduced. "
            "Instead got dense tensor as shape {}".format(shape))
      return TensorShape(shape)
  
    def _get_input_shape_for_sparse_tensor(self, tensor, feature,
                                           path) -> TensorShape:
      """Get the input shape for the sparse tensor."""
      shape = tensor.shape.as_list()
      # Only 2 and above rank sparse tensor is supported.
      if len(shape) < 2:
        raise ValueError("Only rank 2 and above sparse tensor is supported,"
                         " find rank {} sparse tensor for input {}".format(
                             len(shape), path))
      if not feature.output_shape and feature.max_sequence_length > 0:
        # If the max_sequence_length is set and the output shape for FeatureConfig
        # is not set, we modify the shape of the input feature. Only rank 2
        # feature output shape is modified
        if len(shape) == 2:
          # If the sparse tensor is 2D and max_sequence_length is set,
          # we need to add one dimension to the input feature.
          shape.insert(len(shape) - 1, feature.max_sequence_length)
  
      return TensorShape(shape)
  
    def _get_input_shape_for_ragged_tensor(self, tensor, feature,
                                           path) -> TensorShape:
      """Get the input shape for the ragged tensor."""
      shape = tensor.shape.as_list()
      # Only rank 2 ragged tensor is supported.
      if len(shape) != 2:
        raise ValueError("Only rank 2 ragged tensor is supported,"
                         " find rank {} ragged tensor for input {}".format(
                             len(shape), path))
      if not feature.output_shape and feature.max_sequence_length > 0:
        # If the max_sequence_length is set and the output shape for FeatureConfig
        # is not set, add the sequence length as second last dimension of
        # the ragged tensor.
        shape.insert(len(shape) - 1, feature.max_sequence_length)
  
      return TensorShape(shape)
  
    def _update_output_shapes(self, incoming_output_shapes: List[TensorShape]):
      """Update the existing output shapes based on the new output shapes.
  
      The existing output shapes always have higher piority than the new incoming
      output shapes.
      Args:
        incoming_output_shapes: nested structure of TensorShape to override the
          existing output shapes.
      """
      nest.assert_same_structure(self._output_shapes, incoming_output_shapes)
      updated_output_shapes = []
      for old_output_shape, incoming_output_shape in zip(self._output_shapes,
                                                         incoming_output_shapes):
        if old_output_shape:
          updated_output_shapes.append(old_output_shape)
        else:
          updated_output_shapes.append(incoming_output_shape)
      self._output_shapes = updated_output_shapes
  
    def _check_output_shapes(self, incoming_output_shapes: List[TensorShape]):
      """Check the incoming output shapes against the output shapes stored."""
      # The incoming output shape should have the same structure with the existing
      # output shapes.
      nest.assert_same_structure(self._output_shapes, incoming_output_shapes)
  
      for (path, _), old_output_shape, incoming_output_shape in zip(
          nest.flatten_with_joined_string_paths(self._feature_config),
          self._output_shapes, incoming_output_shapes):
        # First check if both shapes are not None.
        if old_output_shape and incoming_output_shape:
          # We skip the check when the incoming output shape is rank 1 or 2 and
          # rank of the old output shape is larger. This can happen for
          # (sequence) ragged tensor, we push the check down to the enqueue op.
          if (len(incoming_output_shape) == 1 or len(incoming_output_shape)
              == 2) and len(old_output_shape) > len(incoming_output_shape):
            continue
          if len(old_output_shape) != len(
              incoming_output_shape) or not self._is_tensor_shape_match(
                  old_output_shape, incoming_output_shape):
            raise ValueError(
                f"Inconsistent shape founded for input feature {path}, "
                f"Output shape is set to be {old_output_shape}, "
                f"But got incoming output shape {incoming_output_shape}")
  
    def _check_output_shapes_fully_defined(self):
      """Check if the output shape is fully defined."""
      for (path, _), output_shape in zip(
          nest.flatten_with_joined_string_paths(self._feature_config),
          self._output_shapes):
        if not output_shape.is_fully_defined():
          raise ValueError(
              f"Input Feature {path} has output shape set as "
              f"{output_shape} which is not fully defined. "
              "Please specify the fully defined shape in either FeatureConfig "
              "or for the build method.")
  
    def _is_tensor_shape_match(self, shape_a: TensorShape,
                               shape_b: TensorShape) -> bool:
      """Check if shape b matches with shape a."""
      for s_a, s_b in zip(shape_a.as_list(), shape_b.as_list()):
        if s_a and s_b and s_a != s_b:
          return False
      return True
  
    def _get_output_shapes_from_batch_size(self, per_replica_batch_size):
      """Get the output shapes from the batch size."""
      output_shapes = []
      for feature in nest.flatten(self._feature_config):
        if not feature.output_shape and feature.max_sequence_length > 0:
          output_shapes.append(
              TensorShape([per_replica_batch_size, feature.max_sequence_length]))
        else:
          output_shapes.append(TensorShape(per_replica_batch_size))
      return output_shapes
  
  
  @def_function.function
  def _load_variables_impl(
      config: Text,
      hosts: List[Tuple[int, Text]],
      variables: Dict[Text, Dict[Text, tf_variables.Variable]],
      table_config: tpu_embedding_v2_utils.TableConfig):
    """Load embedding tables to onto TPU for each table and host.
  
    Args:
      config: A serialized TPUEmbeddingConfiguration proto.
      hosts: A list of CPU devices, on per host.
      variables: A dictionary of dictionaries of TPUShardedVariables. First key is
        the table name, second key is 'parameters' or the optimizer slot name.
      table_config: A list of tf.tpu.experimental.embedding.TableConfig objects.
    """
    def select_fn(host_id):
  
      def select_or_zeros(x):
        if host_id >= len(x.variables):
          # In the edge case where we have more hosts than variables, due to using
          # a small number of rows, we load zeros for the later hosts. We copy
          # the shape of the first host's variables, which we assume is defined
          # because TableConfig guarantees at least one row.
          return array_ops.zeros_like(x.variables[0])
        return x.variables[host_id]
  
      return select_or_zeros
  
    for host_id, host in enumerate(hosts):
      with ops.device(host):
        host_variables = nest.map_structure(select_fn(host_id), variables)
        for table in table_config:
          table.optimizer._load()(  # pylint: disable=protected-access
              table_name=table.name,
              num_shards=len(hosts),
              shard_id=host_id,
              config=config,
              **host_variables[table.name])
          # Ensure that only the first table/first host gets a config so that we
          # don't bloat graph by attaching this large string to each op.
          # We have num tables * num hosts of these so for models with a large
          # number of tables training on a large slice, this can be an issue.
          config = None
  
  
  @def_function.function
  def _retrieve_variables_impl(
      config: Text,
      hosts: List[Tuple[int, Text]],
      variables: Dict[Text, Dict[Text, tf_variables.Variable]],
      table_config: tpu_embedding_v2_utils.TableConfig):
    """Retrieve embedding tables from TPU to host memory.
  
    Args:
      config: A serialized TPUEmbeddingConfiguration proto.
      hosts: A list of all the host CPU devices.
      variables: A dictionary of dictionaries of TPUShardedVariables. First key is
        the table name, second key is 'parameters' or the optimizer slot name.
      table_config: A list of tf.tpu.experimental.embedding.TableConfig objects.
    """
    for host_id, host in enumerate(hosts):
      with ops.device(host):
        for table in table_config:
          retrieved = table.optimizer._retrieve()(  # pylint: disable=protected-access
              table_name=table.name,
              num_shards=len(hosts),
              shard_id=host_id,
              config=config)
          # When there are no slot variables (e.g with SGD) this returns a
          # single tensor rather than a tuple. In this case we put the tensor in
          # a list to make the following code easier to write.
          if not isinstance(retrieved, tuple):
            retrieved = (retrieved,)
  
          for i, slot in enumerate(["parameters"] +
                                   table.optimizer._slot_names()):  # pylint: disable=protected-access
            # We must assign the CPU variables the values of tensors that were
            # returned from the TPU.
            sharded_var = variables[table.name][slot]
            if host_id < len(sharded_var.variables):
              # In the edge case where we have more hosts than variables, due to
              # using a small number of rows, we skip the later hosts.
              sharded_var.variables[host_id].assign(retrieved[i])
          # Ensure that only the first table/first host gets a config so that we
          # don't bloat graph by attaching this large string to each op.
          # We have num tables * num hosts of these so for models with a large
          # number of tables training on a large slice, this can be an issue.
          config = None
  
  
  class TPUEmbeddingSaveable(saveable_hook.SaveableHook):
    """Save/Restore hook to Retrieve/Load TPUEmbedding variables."""
  
    def __init__(
        self,
        name: Text,
        load: Callable[[], Any],
        retrieve: Callable[[], Any]):
      self._load = load
      self._retrieve = retrieve
      super(TPUEmbeddingSaveable, self).__init__(name=name)
  
    def before_save(self):
      if self._retrieve is not None:
        self._retrieve()
  
    def after_restore(self):
      if self._load is not None:
        self._load()
  
  
  def _ragged_embedding_lookup_with_reduce(
      table: tf_variables.Variable,
      ragged: ragged_tensor.RaggedTensor,
      weights: ragged_tensor.RaggedTensor,
      combiner: Text) -> core.Tensor:
    """Compute a ragged lookup followed by a reduce on axis 1.
  
    Args:
      table: The embedding table.
      ragged: A RaggedTensor of ids to look up.
      weights: A RaggedTensor of weights (or None).
      combiner: One of "mean", "sum", "sqrtn".
  
    Returns:
      A Tensor.
    """
    if weights is None:
      weights = array_ops.ones_like(ragged, dtype=table.dtype)
    weights = array_ops.expand_dims(weights, axis=2)
    ragged_result = embedding_ops.embedding_lookup_ragged(table, ragged)
    ragged_result = math_ops.reduce_sum(ragged_result * weights, axis=1)
    if combiner == "mean":
      ragged_result = math_ops.div_no_nan(ragged_result,
                                          math_ops.reduce_sum(weights, axis=1))
    elif combiner == "sqrtn":
      ragged_result = math_ops.div_no_nan(
          ragged_result,
          math_ops.sqrt(math_ops.reduce_sum(weights * weights, axis=1)))
    return ragged_result
  
  
  @tf_export("tpu.experimental.embedding.serving_embedding_lookup")
  def cpu_embedding_lookup(inputs, weights, tables, feature_config):
    """Apply standard lookup ops with `tf.tpu.experimental.embedding` configs.
  
    This function is a utility which allows using the
    `tf.tpu.experimental.embedding` config objects with standard lookup functions.
    This can be used when exporting a model which uses
    `tf.tpu.experimental.embedding.TPUEmbedding` for serving on CPU. In particular
    `tf.tpu.experimental.embedding.TPUEmbedding` only supports lookups on TPUs and
    should not be part of your serving graph.
  
    Note that TPU specific options (such as `max_sequence_length`) in the
    configuration objects will be ignored.
  
    In the following example we take a trained model (see the documentation for
    `tf.tpu.experimental.embedding.TPUEmbedding` for the context) and create a
    saved model with a serving function that will perform the embedding lookup and
    pass the results to your model:
  
    ```python
    model = model_fn(...)
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=feature_config,
        batch_size=1024,
        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
    checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
    checkpoint.restore(...)
  
    @tf.function(input_signature=[{'feature_one': tf.TensorSpec(...),
                                   'feature_two': tf.TensorSpec(...),
                                   'feature_three': tf.TensorSpec(...)}])
    def serve_tensors(embedding_features):
      embedded_features = tf.tpu.experimental.embedding.serving_embedding_lookup(
          embedding_features, None, embedding.embedding_tables,
          feature_config)
      return model(embedded_features)
  
    model.embedding_api = embedding
    tf.saved_model.save(model,
                        export_dir=...,
                        signatures={'serving_default': serve_tensors})
  
    ```
  
    NOTE: Its important to assign the embedding api object to a member of your
    model as `tf.saved_model.save` only supports saving variables one `Trackable`
    object. Since the model's weights are in `model` and the embedding table are
    managed by `embedding`, we assign `embedding` to and attribute of `model` so
    that tf.saved_model.save can find the embedding variables.
  
    NOTE: The same `serve_tensors` function and `tf.saved_model.save` call will
    work directly from training.
  
    Args:
      inputs: a nested structure of Tensors, SparseTensors or RaggedTensors.
      weights: a nested structure of Tensors, SparseTensors or RaggedTensors or
        None for no weights. If not None, structure must match that of inputs, but
        entries are allowed to be None.
      tables: a dict of mapping TableConfig objects to Variables.
      feature_config: a nested structure of FeatureConfig objects with the same
        structure as inputs.
  
    Returns:
      A nested structure of Tensors with the same structure as inputs.
    """
  
    nest.assert_same_structure(inputs, feature_config)
  
    flat_inputs = nest.flatten(inputs)
    flat_weights = [None] * len(flat_inputs)
    if weights is not None:
      nest.assert_same_structure(inputs, weights)
      flat_weights = nest.flatten(weights)
    flat_features = nest.flatten_with_joined_string_paths(feature_config)
  
    outputs = []
    for inp, weight, (path, feature) in zip(
        flat_inputs, flat_weights, flat_features):
      table = tables[feature.table]
  
      if weight is not None:
        if isinstance(inp, ops.Tensor):
          raise ValueError(
              "Weight specified for {}, but input is dense.".format(path))
        elif type(weight) is not type(inp):
          raise ValueError(
              "Weight for {} is of type {} but it does not match type of the "
              "input which is {}.".format(path, type(weight), type(inp)))
        elif feature.max_sequence_length > 0:
          raise ValueError("Weight specified for {}, but this is a sequence "
                           "feature.".format(path))
  
      if isinstance(inp, ops.Tensor):
        if feature.max_sequence_length > 0:
          raise ValueError("Feature {} is a sequence feature but a dense tensor "
                           "was passed.".format(path))
        outputs.append(embedding_ops.embedding_lookup_v2(table, inp))
  
      elif isinstance(inp, sparse_tensor.SparseTensor):
        if not feature.output_shape and feature.max_sequence_length > 0:
          batch_size = math_ops.cast(array_ops.shape(inp)[0], dtype=dtypes.int64)
          sparse_shape = array_ops.stack(
              [batch_size, feature.max_sequence_length], axis=0)
          # TPU Embedding truncates sequences to max_sequence_length, and if we
          # don't truncate, scatter_nd will error out if the index was out of
          # bounds.
          truncated_inp = sparse_ops.sparse_slice(inp, start=[0, 0],
                                                  size=sparse_shape)
  
          dense_output_shape = array_ops.stack(
              [batch_size, feature.max_sequence_length, feature.table.dim],
              axis=0)
          outputs.append(
              array_ops.scatter_nd(
                  truncated_inp.indices,
                  array_ops.gather(table.read_value(), truncated_inp.values),
                  dense_output_shape))
        else:
          inp_rank = inp.dense_shape.get_shape()[0]
          if (not feature.validate_weights_and_indices and
              inp_rank is not None and inp_rank <= 2):
            outputs.append(
                embedding_ops.embedding_lookup_sparse_v2(
                    table,
                    inp,
                    sp_weights=weight,
                    combiner=feature.table.combiner))
          else:
            outputs.append(
                embedding_ops.safe_embedding_lookup_sparse_v2(
                    table,
                    inp,
                    sparse_weights=weight,
                    combiner=feature.table.combiner))
  
      elif isinstance(inp, ragged_tensor.RaggedTensor):
        if inp.shape.rank != 2:
          raise ValueError(
              "Only rank 2 ragged tensor is supported, but got rank {}".format(
                  inp.shape.rank))
        batch_size = inp.shape[0]
        if feature.output_shape:
          output_batch_size = math_ops.reduce_prod(feature.output_shape)
          # If the output batch size matches the data batch size, treat it as
          # normal ragged input.
          if output_batch_size == batch_size:
            ragged_output = _ragged_embedding_lookup_with_reduce(
                table, inp, weight, feature.table.combiner)
            ragged_output = array_ops.reshape(
                ragged_output, shape=feature.output_shape + [feature.table.dim])
          # If the data batch size is a factor of the output batch size, the
          # divide result will be the sequence length. Ignore the weights and
          # combiner.
          elif output_batch_size > batch_size and output_batch_size % batch_size == 0:
            ragged_output = embedding_ops.embedding_lookup_v2(table, inp)
            # Pad or truncate in the sequence dimension
            ragged_output = ragged_output.to_tensor(shape=[
                batch_size, output_batch_size // batch_size, feature.table.dim
            ])
            # Reshape to desire output shape.
            ragged_output = array_ops.reshape(
                ragged_output, feature.output_shape + [feature.table.dim])
          else:
            raise ValueError(
                "Output shape set in the FeatureConfig should be the factor of "
                "the input data batch size. But instead got output shape {}, "
                "input data batch size {}".format(feature.output_shape,
                                                  batch_size))
        else:
          if feature.max_sequence_length > 0:
            output_shape = [
                batch_size, feature.max_sequence_length, feature.table.dim
            ]
            ragged_lookup = embedding_ops.embedding_lookup_v2(table, inp)
            # Unlike scatter_nd, RaggedTensor.to_tensor truncates to the given
            # shape.
            ragged_output = ragged_lookup.to_tensor(shape=output_shape)
          else:
            ragged_output = _ragged_embedding_lookup_with_reduce(
                table, inp, weight, feature.table.combiner)
        outputs.append(ragged_output)
      else:
        raise ValueError("Input {} is type {}. Tensor, SparseTensor or "
                         "RaggedTensor expected.".format(path, type(inp)))
    return nest.pack_sequence_as(feature_config, outputs)
  
  
  def get_list_of_hosts(strategy: tpu_strategy.TPUStrategy) -> List[Text]:
    """Returns a sorted list of CPU devices for the remote jobs.
  
    Args:
      strategy: A TPUStrategy object.
  
    Returns:
      A sort list of device strings.
    """
    list_of_hosts = []
    # Assume this is sorted by task
    for tpu_device in strategy.extended.worker_devices:
      host = device_util.get_host_for_device(tpu_device)
      if host not in list_of_hosts:
        list_of_hosts.append(host)
    assert len(list_of_hosts) == strategy.extended.num_hosts
    return list_of_hosts
  
  
  def extract_variable_info(
      kwargs) -> Tuple[Text, Tuple[int, ...], dtypes.DType, Callable[[], Any]]:
    """Extracts the variable creation attributes from the kwargs.
  
    Args:
      kwargs: a dict of keyword arguments that were passed to a variable creator
        scope.
  
    Returns:
      A tuple of variable name, shape, dtype, initialization function.
    """
    if (isinstance(kwargs["initial_value"], functools.partial) and (
        "shape" in kwargs["initial_value"].keywords or
        kwargs["initial_value"].args)):
      # Sometimes shape is passed positionally, sometimes it's passed as a kwarg.
      if "shape" in kwargs["initial_value"].keywords:
        shape = kwargs["initial_value"].keywords["shape"]
      else:
        shape = kwargs["initial_value"].args[0]
      return (kwargs["name"], shape,
              kwargs["initial_value"].keywords.get("dtype", kwargs["dtype"]),
              kwargs["initial_value"].func)
    elif "shape" not in kwargs or kwargs["shape"] is None or not callable(
        kwargs["initial_value"]):
      raise ValueError(
          "Unable to extract initializer function and shape from {}. Please "
          "either pass a function that expects a shape and dtype as the "
          "initial value for your variable or functools.partial object with "
          "the shape and dtype kwargs set. This is needed so that we can "
          "initialize the shards of the ShardedVariable locally.".format(
              kwargs["initial_value"]))
    else:
      return (kwargs["name"], kwargs["shape"], kwargs["dtype"],
              kwargs["initial_value"])
  
  
  def make_sharded_variable_creator(
      hosts: List[Text]) -> Callable[..., TPUShardedVariable]:
    """Makes a sharded variable creator given a list of hosts.
  
    Args:
      hosts: a list of tensorflow devices on which to shard the tensors.
  
    Returns:
      A variable creator function.
    """
  
    def sharded_variable_creator(
        next_creator: Callable[..., tf_variables.Variable], *args, **kwargs):
      """The sharded variable creator."""
      kwargs["skip_mirrored_creator"] = True
  
      num_hosts = len(hosts)
      name, shape, dtype, unwrapped_initial_value = extract_variable_info(kwargs)
      initial_value = kwargs["initial_value"]
      rows = shape[0]
      cols = shape[1]
      partial_partition = rows % num_hosts
      full_rows_per_host = rows // num_hosts
      # We partition as if we were using MOD sharding: at least
      # `full_rows_per_host` rows to `num_hosts` hosts, where the first
      # `partial_partition` hosts get an additional row when the number of rows
      # is not cleanly divisible. Note that `full_rows_per_host` may be zero.
      partitions = (
          [full_rows_per_host + 1] * partial_partition
          + [full_rows_per_host] * (num_hosts - partial_partition))
      variables = []
      sharding_aware = "shard_info" in tf_inspect.getargspec(initial_value).args
  
      # Keep track of offset for sharding aware initializers.
      offset = 0
      kwargs["dtype"] = dtype
      for i, p in enumerate(partitions):
        if p == 0:
          # Skip variable creation for empty partitions, resulting from the edge
          # case of 'rows < num_hosts'. This is safe because both load/restore
          # can handle the missing values.
          continue
        with ops.device(hosts[i]):
          kwargs["name"] = "{}_{}".format(name, i)
          kwargs["shape"] = (p, cols)
          if sharding_aware:
            shard_info = base.ShardInfo(kwargs["shape"], (offset, 0))
            kwargs["initial_value"] = functools.partial(
                initial_value, shard_info=shard_info)
            offset += p
          else:
            kwargs["initial_value"] = functools.partial(
                unwrapped_initial_value, kwargs["shape"], dtype=dtype)
          variables.append(next_creator(*args, **kwargs))
      return TPUShardedVariable(variables, name=name)
    return sharded_variable_creator
   # Copyright 2015 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 tensorflow.python.ops.op_def_library."""
  
  from tensorflow.core.framework import tensor_shape_pb2
  from tensorflow.python.eager import function as eager_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import function
  from tensorflow.python.framework import op_def_library
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.framework import test_util
  from tensorflow.python.platform import googletest
  from tensorflow.python.util import compat
  
  
  class OpDefLibraryTest(test_util.TensorFlowTestCase):
  
    def Tensor(self, t, name="in"):
      return op_def_library.apply_op("OutT", T=t, name=name)
  
    def testNoRegisteredOpFails(self):
      with self.assertRaises(RuntimeError) as cm:
        op_def_library.apply_op("unknown")
      self.assertEqual(str(cm.exception), "Unrecognized Op name unknown")
  
    def testSimple(self):
      with ops.Graph().as_default():
        out = op_def_library.apply_op("Simple", a=3)
        self.assertEqual(dtypes.float32, out.dtype)
        self.assertProtoEquals("""
          name: 'Simple' op: 'Simple' input: 'Simple/a'
          """, out.op.node_def)
  
        out = op_def_library.apply_op("Simple", a=4)
        self.assertProtoEquals("""
          name: 'Simple_1' op: 'Simple' input: 'Simple_1/a'
          """, out.op.node_def)
  
        out = op_def_library.apply_op("Simple", a=5, name="named")
        self.assertProtoEquals("""
          name: 'named' op: 'Simple' input: 'named/a'
          """, out.op.node_def)
  
        out = op_def_library.apply_op(
            "Simple", a=[[1, 2, 3], [4, 5, 6]], name="two_d")
        self.assertProtoEquals("""
          name: 'two_d' op: 'Simple' input: 'two_d/a'
          """, out.op.node_def)
  
    def testSimpleFailures(self):
      with ops.Graph().as_default():
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Simple", a="Bad string")
        self.assertIn(
            "Expected int32 passed to parameter 'a' of op 'Simple', "
            "got 'Bad string' of type 'str' instead.", str(cm.exception))
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Simple", a=self.Tensor(dtypes.string))
        self.assertIn(
            "Input 'a' of 'Simple' Op has type string "
            "that does not match expected type of int32.", str(cm.exception))
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Simple", a=6, extra="bogus")
        self.assertIn("Simple got unexpected keyword arguments: extra",
                      str(cm.exception))
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "Simple", a=6, extra1="bogus", extra2="also_bogus")
        self.assertIn(
            "Simple got unexpected keyword arguments: extra1, "
            "extra2", str(cm.exception))
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Simple")
        self.assertIn("No argument for input a", str(cm.exception))
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Simple", wrong=7)
        self.assertIn("No argument for input a", str(cm.exception))
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Simple", a={"label": 1})
        self.assertIn(
            "Expected int32 passed to parameter 'a' of op 'Simple', "
            "got {'label': 1} of type 'dict' instead.", str(cm.exception))
  
    def testReservedInput(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("ReservedInput", input_=7, name="x")
        self.assertProtoEquals("""
          name: 'x' op: 'ReservedInput' input: 'x/input'
          """, op.node_def)
  
    def testPolymorphic(self):
      with ops.Graph().as_default():
        out = op_def_library.apply_op("Polymorphic", a=7, name="p")
        self.assertEqual(dtypes.int32, out.dtype)
        self.assertProtoEquals("""
          name: 'p' op: 'Polymorphic' input: 'p/a'
          attr { key: 'T' value { type: DT_INT32 } }
          """, out.op.node_def)
  
        out = op_def_library.apply_op("Polymorphic", a="s", name="q")
        self.assertEqual(dtypes.string, out.dtype)
        self.assertProtoEquals("""
          name: 'q' op: 'Polymorphic' input: 'q/a'
          attr { key: 'T' value { type: DT_STRING } }
          """, out.op.node_def)
  
        out = op_def_library.apply_op("Polymorphic", a=["s", "t", "u"], name="r")
        self.assertEqual(dtypes.string, out.dtype)
        self.assertProtoEquals("""
          name: 'r' op: 'Polymorphic' input: 'r/a'
          attr { key: 'T' value { type: DT_STRING } }
          """, out.op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Polymorphic", a="s", T=dtypes.string)
        self.assertEqual(
            str(cm.exception),
            "Should not specify value for inferred attr 'T' for "
            "Polymorphic.")
  
    def testPolymorphicOut(self):
      with ops.Graph().as_default():
        out = op_def_library.apply_op("PolymorphicOut", T=dtypes.int32, name="p")
        self.assertEqual(dtypes.int32, out.dtype)
        self.assertProtoEquals("""
          name: 'p' op: 'PolymorphicOut'
          attr { key: 'T' value { type: DT_INT32 } }
          """, out.op.node_def)
  
        out = op_def_library.apply_op("PolymorphicOut", T=dtypes.bool, name="q")
        self.assertEqual(dtypes.bool, out.dtype)
        self.assertProtoEquals("""
          name: 'q' op: 'PolymorphicOut'
          attr { key: 'T' value { type: DT_BOOL } }
          """, out.op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("PolymorphicOut")
        self.assertEqual(
            str(cm.exception), "No argument found for attr T for PolymorphicOut")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("PolymorphicOut", T=None)
        self.assertEqual(str(cm.exception),
                         "Expected DataType for argument 'T' not None.")
  
    def testPolymorphicDefaultOut(self):
      with ops.Graph().as_default():
        out = op_def_library.apply_op("PolymorphicDefaultOut", T=None, name="p")
        self.assertEqual(dtypes.string, out.dtype)
        self.assertProtoEquals("""
          name: 'p' op: 'PolymorphicDefaultOut'
          attr { key: 'T' value { type: DT_STRING } }
          """, out.op.node_def)
  
        out = op_def_library.apply_op(
            "PolymorphicDefaultOut", T=dtypes.bool, name="q")
        self.assertEqual(dtypes.bool, out.dtype)
        self.assertProtoEquals("""
          name: 'q' op: 'PolymorphicDefaultOut'
          attr { key: 'T' value { type: DT_BOOL } }
          """, out.op.node_def)
  
    def testBinary(self):
      with ops.Graph().as_default():
        out = op_def_library.apply_op("Binary", a=8, b=9, name="b")
        self.assertEqual(dtypes.int32, out.dtype)
        self.assertProtoEquals("""
          name: 'b' op: 'Binary' input: 'b/a' input: 'b/b'
          attr { key: 'T' value { type: DT_INT32 } }
          """, out.op.node_def)
  
        out = op_def_library.apply_op("Binary", a="left", b="right", name="c")
        self.assertEqual(dtypes.string, out.dtype)
        self.assertProtoEquals("""
          name: 'c' op: 'Binary' input: 'c/a' input: 'c/b'
          attr { key: 'T' value { type: DT_STRING } }
          """, out.op.node_def)
  
        with self.assertRaises(TypeError):
          op_def_library.apply_op("Binary", a="left", b=12)
  
        with self.assertRaises(TypeError):
          op_def_library.apply_op(
              "Binary", a=self.Tensor(dtypes.string), b=self.Tensor(dtypes.int32))
  
    def testRestrict(self):
      with ops.Graph().as_default():
        out = op_def_library.apply_op("Restrict", a="foo", name="g")
        self.assertEqual(dtypes.string, out.dtype)
        self.assertProtoEquals("""
          name: 'g' op: 'Restrict' input: 'g/a'
          attr { key: 'T' value { type: DT_STRING } }
          """, out.op.node_def)
  
        out = op_def_library.apply_op("Restrict", a=True, name="h")
        self.assertEqual(dtypes.bool, out.dtype)
        self.assertProtoEquals("""
          name: 'h' op: 'Restrict' input: 'h/a'
          attr { key: 'T' value { type: DT_BOOL } }
          """, out.op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Restrict", a=17)
        self.assertEqual(str(cm.exception),
                         "Value passed to parameter 'a' has DataType int32 "
                         "not in list of allowed values: string, bool")
  
    def testTypeList(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("TypeList", a=["foo"], name="z")
        self.assertProtoEquals("""
          name: 'z' op: 'TypeList' input: 'z/a_0'
          attr { key: 'T' value { list { type: DT_STRING } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("TypeList", a=[True, 12], name="y")
        self.assertProtoEquals("""
          name: 'y' op: 'TypeList' input: 'y/a_0' input: 'y/a_1'
          attr { key: 'T' value { list { type: DT_BOOL type: DT_INT32 } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("TypeList", a=[], name="empty")
        self.assertProtoEquals("""
          name: 'empty' op: 'TypeList' attr { key: 'T' value { list { } } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("TypeList", a=17)
        self.assertStartsWith(str(cm.exception),
                              "Expected list for 'a' "
                              "argument to 'TypeList' Op, not ")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("TypeList", a=[self.Tensor(dtypes.int32), None])
        self.assertStartsWith(str(cm.exception),
                              "Tensors in list passed to 'a' of 'TypeList' Op "
                              "have types [int32, <NOT CONVERTIBLE TO TENSOR>]")
  
    def testTypeListTwice(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "TypeListTwice", a=["foo", True], b=["bar", False], name="z")
        self.assertProtoEquals("""
          name: 'z' op: 'TypeListTwice'
          input: 'z/a_0' input: 'z/a_1' input: 'z/b_0' input: 'z/b_1'
          attr { key: 'T' value { list { type: DT_STRING type: DT_BOOL } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("TypeListTwice", a=[], b=[], name="empty")
        self.assertProtoEquals("""
          name: 'empty' op: 'TypeListTwice' attr { key: 'T' value { list { } } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("TypeListTwice", a=["foo", True], b=["bar", 6])
        self.assertEqual(str(cm.exception),
                         "Input 'b' of 'TypeListTwice' Op has type list of "
                         "string, int32 that does not match type list "
                         "string, bool of argument 'a'.")
  
    def testOutTypeList(self):
      with ops.Graph().as_default():
        out, = op_def_library.apply_op(
            "OutTypeList", T=[dtypes.float32], name="x")
        self.assertEqual(dtypes.float32, out.dtype)
        self.assertProtoEquals("""
          name: 'x' op: 'OutTypeList'
          attr { key: 'T' value { list { type: DT_FLOAT } } }
          """, out.op.node_def)
  
        out1, out2 = op_def_library.apply_op(
            "OutTypeList", T=[dtypes.int32, dtypes.bool], name="w")
        self.assertEqual(dtypes.int32, out1.dtype)
        self.assertEqual(dtypes.bool, out2.dtype)
        self.assertProtoEquals("""
          name: 'w' op: 'OutTypeList'
          attr { key: 'T' value { list { type: DT_INT32 type: DT_BOOL } } }
          """, out1.op.node_def)
  
        out = op_def_library.apply_op("OutTypeList", T=[], name="empty")
        self.assertEqual([], out)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("OutTypeList", T=dtypes.int32)
        self.assertEqual(
            str(cm.exception), "Expected list for attr T, obtained "
            "DType instead.")
  
    def testTypeListRestrict(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "TypeListRestrict", a=["foo", False], name="v")
        self.assertProtoEquals("""
          name: 'v' op: 'TypeListRestrict' input: 'v/a_0' input: 'v/a_1'
          attr { key: 'T' value { list { type: DT_STRING type: DT_BOOL } } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("TypeListRestrict", a=[True, 12])
        self.assertEqual(str(cm.exception),
                         "Value passed to parameter 'a' has DataType int32 "
                         "not in list of allowed values: string, bool")
  
    def testOutTypeListRestrict(self):
      with ops.Graph().as_default():
        out1, out2 = op_def_library.apply_op(
            "OutTypeListRestrict", t=[dtypes.bool, dtypes.string], name="u")
        self.assertEqual(dtypes.bool, out1.dtype)
        self.assertEqual(dtypes.string, out2.dtype)
        self.assertProtoEquals("""
          name: 'u' op: 'OutTypeListRestrict'
          attr { key: 't' value { list { type: DT_BOOL type: DT_STRING } } }
          """, out1.op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "OutTypeListRestrict", t=[dtypes.string, dtypes.int32])
        self.assertEqual(str(cm.exception),
                         "Value passed to parameter 't' has DataType int32 "
                         "not in list of allowed values: string, bool")
  
    def testAttr(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("Attr", a=12, name="t")
        self.assertProtoEquals("""
          name: 't' op: 'Attr' attr { key: 'a' value { i: 12 } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "Attr", a=tensor_shape.Dimension(13), name="u")
        self.assertProtoEquals("""
          name: 'u' op: 'Attr' attr { key: 'a' value { i: 13 } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Attr", a="bad")
        self.assertEqual(str(cm.exception),
                         "Expected int for argument 'a' not 'bad'.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Attr", a=[12])
        self.assertEqual(str(cm.exception),
                         "Expected int for argument 'a' not [12].")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Attr", a=None)
        self.assertEqual(str(cm.exception),
                         "Expected int for argument 'a' not None.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("Attr")
        self.assertEqual(
            str(cm.exception), "No argument found for attr a for "
            "Attr")
  
    def testAttrFloat(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrFloat", a=1.2, name="t")
        self.assertProtoEquals("""
          name: 't' op: 'AttrFloat' attr { key: 'a' value { f: 1.2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrFloat", a=12, name="u")
        self.assertProtoEquals("""
          name: 'u' op: 'AttrFloat' attr { key: 'a' value { f: 12 } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("AttrFloat", a="bad")
        self.assertEqual(str(cm.exception),
                         "Expected float for argument 'a' not 'bad'.")
  
    def testAttrFunc(self):
      with ops.Graph().as_default():
        @function.Defun(dtypes.float32, func_name="MyFn")
        def fn(x):
          return 2 + x
  
        op = op_def_library.apply_op("FuncAttr", f=fn, name="t")
        self.assertProtoEquals("""
          name: 't' op: 'FuncAttr' attr { key: 'f'
                                          value { func { name: 'MyFn' } } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("FuncAttr", f=3)
        self.assertEqual(str(cm.exception),
                         "Don't know how to convert 3 to a func for argument f")
  
    def testAttrFuncWithFuncWithAttrs(self):
      with ops.Graph().as_default():
        @eager_function.defun_with_attributes(
            input_signature=(tensor_spec.TensorSpec(None, dtypes.float32),),
            autograph=False,
            attributes={"_dummy_attr": 15})
        def fn(x):
          return 2 + x
  
        concrete_fn = fn.get_concrete_function()
  
        op = op_def_library.apply_op("FuncAttr", f=concrete_fn, name="t")
        self.assertProtoEquals("""
          name: 't' op: 'FuncAttr'
          attr {
            key: 'f'
            value {
              func {
                name: '%s'
                attr { key: "_dummy_attr" value { i: 15 } }
              }
            }
          }
          """ % compat.as_str(concrete_fn.name), op.node_def)
  
    def testAttrFuncList(self):
      with ops.Graph().as_default():
        @function.Defun(dtypes.float32, func_name="MyFn")
        def fn1(x):
          return 2 + x
        @function.Defun(dtypes.int32, dtypes.float32, func_name="MyFn2")
        def fn2(x, y):
          return 2 + x, y * 3
        @function.Defun(dtypes.int32, func_name="MyFn3")
        def fn3(y):
          return 2 + y
  
        op = op_def_library.apply_op("FuncListAttr", f=[fn1, fn2, fn3], name="t")
        self.assertProtoEquals("""
          name: 't' op: 'FuncListAttr'
          attr { key: 'f' value { list { func { name: 'MyFn' }
                                         func { name: 'MyFn2' }
                                         func { name: 'MyFn3' } } } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("FuncListAttr", f=[fn1, 3, fn2])
        self.assertEqual(str(cm.exception),
                         "Don't know how to convert 3 to a func for argument f")
  
    def testAttrBool(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrBool", a=True, name="t")
        self.assertProtoEquals("""
          name: 't' op: 'AttrBool' attr { key: 'a' value { b: true } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrBool", a=False, name="u")
        self.assertProtoEquals("""
          name: 'u' op: 'AttrBool' attr { key: 'a' value { b: false } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("AttrBool", a=0)
        self.assertEqual(str(cm.exception),
                         "Expected bool for argument 'a' not 0.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("AttrBool", a=1)
        self.assertEqual(str(cm.exception),
                         "Expected bool for argument 'a' not 1.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("AttrBool", a=[])
        self.assertEqual(str(cm.exception),
                         "Expected bool for argument 'a' not [].")
  
    def testAttrBoolList(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "AttrBoolList", a=[True, False, True], name="t")
        self.assertProtoEquals("""
          name: 't' op: 'AttrBoolList'
          attr { key: 'a' value { list { b: true b: false b:true } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrBoolList", a=[], name="u")
        self.assertProtoEquals("""
          name: 'u' op: 'AttrBoolList' attr { key: 'a' value { list { } } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("AttrBoolList", a=[0])
        self.assertEqual(str(cm.exception),
                         "Expected bool for argument 'a' not 0.")
  
    def testAttrMin(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrMin", a=12, name="s")
        self.assertProtoEquals("""
          name: 's' op: 'AttrMin' attr { key: 'a' value { i: 12 } }
          """, op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("AttrMin", a=2)
        self.assertEqual(str(cm.exception),
                         "Attr 'a' of 'AttrMin' Op passed 2 less than minimum 5.")
  
    def testAttrListMin(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrListMin", a=[1, 2], name="r")
        self.assertProtoEquals("""
          name: 'r' op: 'AttrListMin'
          attr { key: 'a' value { list { i: 1 i: 2 } } }
          """, op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("AttrListMin", a=[17])
        self.assertEqual(str(cm.exception),
                         "Attr 'a' of 'AttrListMin' Op "
                         "passed list of length 1 less than minimum 2.")
  
    def testAttrEnum(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrEnum", a="oranges", name="e")
        self.assertProtoEquals("""
          name: 'e' op: 'AttrEnum' attr { key: 'a' value { s: 'oranges' } }
          """, op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("AttrEnum", a="invalid")
        self.assertEqual(str(cm.exception),
                         'Attr \'a\' of \'AttrEnum\' Op '
                         'passed string \'invalid\' not in: '
                         '"apples", "oranges".')
  
    def testAttrEnumList(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "AttrEnumList", a=["oranges", "apples"], name="f")
        self.assertProtoEquals("""
          name: 'f' op: 'AttrEnumList'
          attr { key: 'a' value { list { s: 'oranges' s: 'apples' } } }
          """, op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op(
              "AttrEnumList", a=["apples", "invalid", "oranges"])
        self.assertEqual(str(cm.exception),
                         'Attr \'a\' of \'AttrEnumList\' Op '
                         'passed string \'invalid\' not '
                         'in: "apples", "oranges".')
  
    def testAttrShape(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrShape", a=[5], name="s1")
        self.assertProtoEquals("""
          name: 's1' op: 'AttrShape'
          attr { key: 'a' value { shape { dim { size: 5 } } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrShape", a=(4, 3, 2), name="s2")
        self.assertProtoEquals("""
          name: 's2' op: 'AttrShape'
          attr { key: 'a' value {
            shape { dim { size: 4 } dim { size: 3 } dim { size: 2 } } } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "AttrShape", a=tensor_shape.TensorShape([3, 2]), name="s3")
        self.assertProtoEquals("""
          name: 's3' op: 'AttrShape'
          attr { key: 'a' value {
            shape { dim { size: 3 } dim { size: 2 } } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrShape", a=[], name="s4")
        self.assertProtoEquals("""
          name: 's4' op: 'AttrShape' attr { key: 'a' value { shape { } } }
          """, op.node_def)
  
        shape = tensor_shape_pb2.TensorShapeProto()
        shape.dim.add().size = 6
        shape.dim.add().size = 3
        op = op_def_library.apply_op("AttrShape", a=shape, name="s5")
        self.assertProtoEquals("""
          name: 's5' op: 'AttrShape'
          attr { key: 'a' value { shape { dim { size: 6 } dim { size: 3 } } } }
          """, op.node_def)
  
        # TODO(josh11b): Re-enable this test once we stop promoting scalars to
        # shapes.
        # with self.assertRaises(TypeError) as cm:
        #   op_def_library.apply_op("AttrShape", a=5)
        # self.assertEqual(str(cm.exception),
        #                  "Don't know how to convert 5 to a TensorShapeProto for"
        #                  " argument 'a'")
  
        with self.assertRaises(TypeError):
          op_def_library.apply_op("AttrShape", a="ABC")
  
    def testAttrShapeList(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "AttrShapeList", a=[[3, 2], [6, 5, 4]], name="sl")
        self.assertProtoEquals("""
          name: 'sl' op: 'AttrShapeList'
          attr { key: 'a' value { list {
            shape { dim { size: 3 } dim { size: 2 } }
            shape { dim { size: 6 } dim { size: 5 } dim { size: 4 } } } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrShapeList", a=[], name="esl")
        self.assertProtoEquals("""
          name: 'esl' op: 'AttrShapeList' attr { key: 'a' value { list { } } }
          """, op.node_def)
  
    def testAttrPartialShape(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrPartialShape", a=[5], name="s1")
        self.assertProtoEquals("""
          name: 's1' op: 'AttrPartialShape'
          attr { key: 'a' value { shape { dim { size: 5 } } } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "AttrPartialShape", a=(4, None, 2), name="s2")
        self.assertProtoEquals("""
          name: 's2' op: 'AttrPartialShape'
          attr { key: 'a' value {
            shape { dim { size: 4 } dim { size: -1 } dim { size: 2 } } } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "AttrPartialShape", a=tensor_shape.TensorShape([3, None]), name="s3")
        self.assertProtoEquals("""
          name: 's3' op: 'AttrPartialShape'
          attr { key: 'a' value {
            shape { dim { size: 3 } dim { size: -1 } } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrPartialShape", a=[], name="s4")
        self.assertProtoEquals("""
          name: 's4' op: 'AttrPartialShape'
          attr { key: 'a' value { shape { } } }
          """, op.node_def)
  
        shape = tensor_shape_pb2.TensorShapeProto()
        shape.dim.add().size = -1
        shape.dim.add().size = 3
        op = op_def_library.apply_op("AttrPartialShape", a=shape, name="s5")
        self.assertProtoEquals("""
          name: 's5' op: 'AttrPartialShape'
          attr { key: 'a' value {
            shape { dim { size: -1 } dim { size: 3 } } } }
          """, op.node_def)
  
        # TODO(ebrevdo): Re-enable once we stop promoting scalars to shapes.
        # with self.assertRaises(TypeError) as cm:
        #   op_def_library.apply_op("AttrPartialShape", a=5)
        # self.assertEqual(str(cm.exception),
        #                  "Don't know how to convert 5 to a TensorShapeProto for"
        #                  " argument 'a'")
  
        with self.assertRaises(TypeError):
          op_def_library.apply_op("AttrPartialShape", a="ABC")
  
    def testAttrPartialShapeList(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "AttrPartialShapeList", a=[[3, 2], [6, None, 4]], name="sl")
        self.assertProtoEquals("""
          name: 'sl' op: 'AttrPartialShapeList'
          attr { key: 'a' value { list {
            shape { dim { size: 3 } dim { size: 2 } }
            shape { dim { size: 6 } dim { size: -1 } dim { size: 4 } } } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrPartialShapeList", a=[], name="esl")
        self.assertProtoEquals("""
          name: 'esl' op: 'AttrPartialShapeList' attr {
            key: 'a' value { list { } } }
          """, op.node_def)
  
    def testAttrDefault(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrDefault", a=None, name="d")
        self.assertProtoEquals("""
          name: 'd' op: 'AttrDefault' attr { key: 'a' value { s: 'banana' } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrDefault", a="kiwi", name="c")
        self.assertProtoEquals("""
          name: 'c' op: 'AttrDefault' attr { key: 'a' value { s: 'kiwi' } }
          """, op.node_def)
  
    def testAttrListDefault(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrListDefault", a=None, name="b")
        self.assertProtoEquals("""
          name: 'b' op: 'AttrListDefault'
          attr { key: 'a' value { list { i: 5 i: 15 } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrListDefault", a=[3], name="a")
        self.assertProtoEquals("""
          name: 'a' op: 'AttrListDefault'
          attr { key: 'a' value { list { i: 3 } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrListDefault", a=[], name="empty")
        self.assertProtoEquals("""
          name: 'empty' op: 'AttrListDefault'
          attr { key: 'a' value { list { } } }
          """, op.node_def)
  
    def testAttrEmptyListDefault(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("AttrEmptyListDefault", a=None, name="b")
        self.assertProtoEquals("""
          name: 'b' op: 'AttrEmptyListDefault'
          attr { key: 'a' value { list { } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrEmptyListDefault", a=[3], name="a")
        self.assertProtoEquals("""
          name: 'a' op: 'AttrEmptyListDefault'
          attr { key: 'a' value { list { f: 3 } } }
          """, op.node_def)
  
        op = op_def_library.apply_op("AttrEmptyListDefault", a=[], name="empty")
        self.assertProtoEquals("""
          name: 'empty' op: 'AttrEmptyListDefault'
          attr { key: 'a' value { list { } } }
          """, op.node_def)
  
    def testReservedAttr(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("ReservedAttr", range_=7, name="x")
        self.assertProtoEquals("""
          name: 'x' op: 'ReservedAttr' attr { key: 'range' value { i: 7 } }
          """, op.node_def)
  
    def testDefaultAttrType(self):
      with ops.Graph().as_default():
        # Give an input whose type has no obvious output type.
        op = op_def_library.apply_op("AttrTypeDefault", a=[], name="n")
        self.assertProtoEquals("""
          name: 'n' op: 'AttrTypeDefault' input: 'n/a'
          attr { key: 'T' value { type: DT_INT32 } }
          """, op.node_def)
  
        # Give an input whose type can be inferred as different
        # than the default.
        op = op_def_library.apply_op("AttrTypeDefault", a=[1.0], name="f")
        self.assertProtoEquals("""
          name: 'f' op: 'AttrTypeDefault' input: 'f/a'
          attr { key: 'T' value { type: DT_FLOAT } }
          """, op.node_def)
  
    def testDefaultListAttrType(self):
      with ops.Graph().as_default():
        # Give an input whose type can be inferred as different
        # than the default.
        op = op_def_library.apply_op(
            "AttrListTypeDefault", a=[1.0], b=[2.0], name="n")
        self.assertProtoEquals("""
          name: 'n' op: 'AttrListTypeDefault' input: 'n/a_0' input: 'n/b_0'
          attr { key: 'T' value { type: DT_FLOAT } }
          attr { key: 'N' value { i: 1 } }
          """, op.node_def)
  
    def testNIntsIn(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("NIntsIn", a=[1, 2], name="n")
        self.assertProtoEquals("""
          name: 'n' op: 'NIntsIn' input: 'n/a_0' input: 'n/a_1'
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op("NIntsIn", a=[5, 4, 3, 2, 1], name="o")
        self.assertProtoEquals("""
          name: 'o' op: 'NIntsIn'
          input: 'o/a_0' input: 'o/a_1' input: 'o/a_2' input: 'o/a_3' input: 'o/a_4'
          attr { key: 'N' value { i: 5 } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NIntsIn", a=["foo", "bar"])
        self.assertEqual(
            str(cm.exception),
            "Tensors in list passed to 'a' of 'NIntsIn' Op have types "
            "[string, string] that do not match expected type int32.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "NIntsIn",
              a=[self.Tensor(dtypes.string),
                 self.Tensor(dtypes.string)])
        self.assertEqual(str(cm.exception),
                         "Tensors in list passed to 'a' of 'NIntsIn' Op have "
                         "types [string, string] that do not match expected type "
                         "int32.")
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("NIntsIn", a=[99])
        self.assertEqual(str(cm.exception),
                         "List argument 'a' to 'NIntsIn' Op "
                         "with length 1 shorter than "
                         "minimum length 2.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NIntsIn", a=[38, "bar"])
        self.assertEqual(
            str(cm.exception),
            "Tensors in list passed to 'a' of 'NIntsIn' Op have types "
            "[int32, string] that do not match expected type int32.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "NIntsIn",
              a=[self.Tensor(dtypes.int32),
                 self.Tensor(dtypes.string)])
        self.assertEqual(str(cm.exception),
                         "Tensors in list passed to 'a' of 'NIntsIn' Op "
                         "have types [int32, string] that do not match expected "
                         "type int32.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NIntsIn", a=17)
        self.assertStartsWith(str(cm.exception),
                              "Expected list for 'a' argument "
                              "to 'NIntsIn' Op, not ")
  
    def testNPolymorphicIn(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op("NPolymorphicIn", a=[1, 2], name="n")
        self.assertProtoEquals("""
          name: 'n' op: 'NPolymorphicIn' input: 'n/a_0' input: 'n/a_1'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "NPolymorphicIn", a=[5, 4, 3, 2, 1], name="o")
        self.assertProtoEquals("""
          name: 'o' op: 'NPolymorphicIn'
          input: 'o/a_0' input: 'o/a_1' input: 'o/a_2' input: 'o/a_3' input: 'o/a_4'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 5 } }
          """, op.node_def)
  
        op = op_def_library.apply_op("NPolymorphicIn", a=["foo", "bar"], name="p")
        self.assertProtoEquals("""
          name: 'p' op: 'NPolymorphicIn' input: 'p/a_0' input: 'p/a_1'
          attr { key: 'T' value { type: DT_STRING } }
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "NPolymorphicIn",
            a=[1, self.Tensor(dtypes.float32, name="x")],
            name="q")
        self.assertProtoEquals("""
          name: 'q' op: 'NPolymorphicIn' input: 'q/a_0' input: 'x'
          attr { key: 'T' value { type: DT_FLOAT } }
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "NPolymorphicIn",
            a=[
                self.Tensor(dtypes.float32, name="y"),
                self.Tensor(dtypes.float32_ref, name="z")
            ],
            name="r")
        self.assertProtoEquals("""
          name: 'r' op: 'NPolymorphicIn' input: 'y' input: 'z'
          attr { key: 'T' value { type: DT_FLOAT } }
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("NPolymorphicIn", a=[99])
        self.assertEqual(str(cm.exception),
                         "List argument 'a' to 'NPolymorphicIn' Op with length 1 "
                         "shorter than minimum length 2.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NPolymorphicIn", a=[38, "bar"])
        self.assertEqual(str(cm.exception),
                         "Tensors in list passed to 'a' of 'NPolymorphicIn' Op "
                         "have types [int32, string] that don't all match.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "NPolymorphicIn", a=[38, self.Tensor(dtypes.string)])
        self.assertEqual(str(cm.exception),
                         "Tensors in list passed to 'a' of 'NPolymorphicIn' Op "
                         "have types [int32, string] that don't all match.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NPolymorphicIn", a=[38, None])
        self.assertEqual(str(cm.exception),
                         "Tensors in list passed to 'a' of 'NPolymorphicIn' Op "
                         "have types [int32, <NOT CONVERTIBLE TO TENSOR>] that "
                         "don't all match.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "NPolymorphicIn", a=["abcd", self.Tensor(dtypes.int32)])
        self.assertEqual(str(cm.exception),
                         "Tensors in list passed to 'a' of 'NPolymorphicIn' Op "
                         "have types [string, int32] that don't all match.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NPolymorphicIn", a=17)
        self.assertStartsWith(str(cm.exception),
                              "Expected list for 'a' argument "
                              "to 'NPolymorphicIn' Op, not ")
  
    def testNPolymorphicRestrictIn(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "NPolymorphicRestrictIn", a=["foo", "bar"], name="p")
        self.assertProtoEquals("""
          name: 'p' op: 'NPolymorphicRestrictIn' input: 'p/a_0' input: 'p/a_1'
          attr { key: 'T' value { type: DT_STRING } }
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "NPolymorphicRestrictIn", a=[False, True, False], name="b")
        self.assertProtoEquals("""
          name: 'b' op: 'NPolymorphicRestrictIn'
          input: 'b/a_0' input: 'b/a_1' input: 'b/a_2'
          attr { key: 'T' value { type: DT_BOOL } }
          attr { key: 'N' value { i: 3 } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NPolymorphicRestrictIn", a=[1, 2])
        self.assertEqual(
            str(cm.exception),
            "Value passed to parameter 'a' has DataType int32 not in "
            "list of allowed values: string, bool")
  
    def testNInTwice(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "NInTwice", a=[1, 2], b=["one", "two"], name="n")
        self.assertProtoEquals("""
          name: 'n' op: 'NInTwice'
          input: 'n/a_0' input: 'n/a_1' input: 'n/b_0' input: 'n/b_1'
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op("NInTwice", a=[], b=[], name="o")
        self.assertProtoEquals("""
          name: 'o' op: 'NInTwice' attr { key: 'N' value { i: 0 } }
          """, op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("NInTwice", a=[1, 2, 3], b=["too short"])
        self.assertEqual(str(cm.exception),
                         "List argument 'b' to 'NInTwice' Op "
                         "with length 1 must match "
                         "length 3 of argument 'a'.")
  
    def testNInPolymorphicTwice(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "NInPolymorphicTwice", a=[1, 2], b=[3, 4], name="n")
        self.assertProtoEquals("""
          name: 'n' op: 'NInPolymorphicTwice'
          input: 'n/a_0' input: 'n/a_1' input: 'n/b_0' input: 'n/b_1'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("NInPolymorphicTwice", a=[1, 2, 3], b=[5])
        self.assertEqual(str(cm.exception),
                         "List argument 'b' to 'NInPolymorphicTwice' Op "
                         "with length 1 "
                         "must match length 3 of argument 'a'.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "NInPolymorphicTwice", a=[1, 2], b=["one", "two"])
        self.assertEqual(str(cm.exception),
                         "Tensors in list passed to 'b' of 'NInPolymorphicTwice' "
                         "Op have types [string, string] that do not match type "
                         "int32 inferred from earlier arguments.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "NInPolymorphicTwice",
              a=[self.Tensor(dtypes.int32)],
              b=[self.Tensor(dtypes.string)])
        self.assertEqual(str(cm.exception),
                         "Tensors in list passed to 'b' of "
                         "'NInPolymorphicTwice' Op have types [string] that do "
                         "not match type int32 inferred from earlier arguments.")
  
    def testNInTwoTypeVariables(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "NInTwoTypeVariables", a=[1, 2], b=[True, False], name="n")
        self.assertProtoEquals("""
          name: 'n' op: 'NInTwoTypeVariables'
          input: 'n/a_0' input: 'n/a_1' input: 'n/b_0' input: 'n/b_1'
          attr { key: 'S' value { type: DT_INT32 } }
          attr { key: 'T' value { type: DT_BOOL } }
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "NInTwoTypeVariables", a=[1, 2], b=[3, 4], name="o")
        self.assertProtoEquals("""
          name: 'o' op: 'NInTwoTypeVariables'
          input: 'o/a_0' input: 'o/a_1' input: 'o/b_0' input: 'o/b_1'
          attr { key: 'S' value { type: DT_INT32 } }
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "NInTwoTypeVariables",
            a=[self.Tensor(dtypes.int32, name="q")],
            b=[self.Tensor(dtypes.string, name="r")],
            name="p")
        self.assertProtoEquals("""
          name: 'p' op: 'NInTwoTypeVariables' input: 'q' input: 'r'
          attr { key: 'S' value { type: DT_INT32 } }
          attr { key: 'T' value { type: DT_STRING } }
          attr { key: 'N' value { i: 1 } }
          """, op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("NInTwoTypeVariables", a=[1, 2, 3], b=["5"])
        self.assertEqual(str(cm.exception),
                         "List argument 'b' to 'NInTwoTypeVariables' Op "
                         "with length 1 "
                         "must match length 3 of argument 'a'.")
  
    def testInPolymorphicTwice(self):
      with ops.Graph().as_default():
        op = op_def_library.apply_op(
            "InPolymorphicTwice", a=[8], b=[3, 4, 5], name="n")
        self.assertProtoEquals("""
          name: 'n' op: 'InPolymorphicTwice'
          input: 'n/a_0' input: 'n/b_0' input: 'n/b_1' input: 'n/b_2'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 1 } }
          attr { key: 'M' value { i: 3 } }
          """, op.node_def)
  
        op = op_def_library.apply_op("InPolymorphicTwice", a=[8], b=[], name="o")
        self.assertProtoEquals("""
          name: 'o' op: 'InPolymorphicTwice' input: 'o/a_0'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 1 } }
          attr { key: 'M' value { i: 0 } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "InPolymorphicTwice", a=[], b=[3, 4], name="p")
        self.assertProtoEquals("""
          name: 'p' op: 'InPolymorphicTwice' input: 'p/b_0' input: 'p/b_1'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 0 } }
          attr { key: 'M' value { i: 2 } }
          """, op.node_def)
  
        op = op_def_library.apply_op(
            "InPolymorphicTwice", a=[], b=[3.0, 4.0], name="q")
        self.assertProtoEquals("""
          name: 'q' op: 'InPolymorphicTwice' input: 'q/b_0' input: 'q/b_1'
          attr { key: 'T' value { type: DT_FLOAT } }
          attr { key: 'N' value { i: 0 } }
          attr { key: 'M' value { i: 2 } }
          """, op.node_def)
  
        # Empty input lists: assume default type for T.
        op = op_def_library.apply_op(
            "InPolymorphicTwice", a=[], b=[], name="r")
        self.assertProtoEquals("""
          name: 'r' op: 'InPolymorphicTwice'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 0 } }
          attr { key: 'M' value { i: 0 } }
          """, op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "InPolymorphicTwice", a=[1, 2], b=["one", "two"])
        self.assertEqual(
            str(cm.exception),
            "Tensors in list passed to 'b' of 'InPolymorphicTwice' Op "
            "have types [string, string] that do not match type int32 "
            "inferred from earlier arguments.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op(
              "InPolymorphicTwice",
              a=[self.Tensor(dtypes.int32)],
              b=[self.Tensor(dtypes.string)])
        self.assertEqual(str(cm.exception),
                         "Tensors in list passed to 'b' of 'InPolymorphicTwice' "
                         "Op have types [string] that do not match type int32 "
                         "inferred from earlier arguments.")
  
    def testNIntsOut(self):
      with ops.Graph().as_default():
        out1, out2 = op_def_library.apply_op("NIntsOut", N=2, name="n")
        self.assertEqual(dtypes.int32, out1.dtype)
        self.assertEqual(dtypes.int32, out2.dtype)
        self.assertProtoEquals("""
          name: 'n' op: 'NIntsOut' attr { key: 'N' value { i: 2 } }
          """, out1.op.node_def)
  
        out1, out2, out3, out4, out5 = op_def_library.apply_op(
            "NIntsOut", N=5, name="o")
        self.assertEqual(dtypes.int32, out1.dtype)
        self.assertEqual(dtypes.int32, out2.dtype)
        self.assertEqual(dtypes.int32, out3.dtype)
        self.assertEqual(dtypes.int32, out4.dtype)
        self.assertEqual(dtypes.int32, out5.dtype)
        self.assertProtoEquals("""
          name: 'o' op: 'NIntsOut' attr { key: 'N' value { i: 5 } }
          """, out5.op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("NIntsOut", N=1)
        self.assertEqual(
            str(cm.exception),
            "Attr 'N' of 'NIntsOut' Op passed 1 less than minimum 2.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NIntsOut", N=[3])
        self.assertEqual(str(cm.exception),
                         "Expected int for argument 'N' not [3].")
  
    def testNIntsOutDefault(self):
      with ops.Graph().as_default():
        out1, out2, out3 = op_def_library.apply_op(
            "NIntsOutDefault", N=None, name="z")
        self.assertEqual(dtypes.int32, out1.dtype)
        self.assertEqual(dtypes.int32, out2.dtype)
        self.assertEqual(dtypes.int32, out3.dtype)
        self.assertProtoEquals("""
          name: 'z' op: 'NIntsOutDefault' attr { key: 'N' value { i: 3 } }
          """, out1.op.node_def)
  
        out1, out2 = op_def_library.apply_op("NIntsOutDefault", N=2, name="y")
        self.assertEqual(dtypes.int32, out1.dtype)
        self.assertEqual(dtypes.int32, out2.dtype)
        self.assertProtoEquals("""
          name: 'y' op: 'NIntsOutDefault' attr { key: 'N' value { i: 2 } }
          """, out2.op.node_def)
  
    def testNPolymorphicOut(self):
      with ops.Graph().as_default():
        out1, out2 = op_def_library.apply_op(
            "NPolymorphicOut", N=2, T=dtypes.int32, name="n")
        self.assertEqual(dtypes.int32, out1.dtype)
        self.assertEqual(dtypes.int32, out2.dtype)
        self.assertProtoEquals("""
          name: 'n' op: 'NPolymorphicOut'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 2 } }
          """, out1.op.node_def)
  
        out1, out2, out3 = op_def_library.apply_op(
            "NPolymorphicOut", T=dtypes.string, N=3, name="o")
        self.assertEqual(dtypes.string, out1.dtype)
        self.assertEqual(dtypes.string, out2.dtype)
        self.assertEqual(dtypes.string, out3.dtype)
        self.assertProtoEquals("""
          name: 'o' op: 'NPolymorphicOut'
          attr { key: 'T' value { type: DT_STRING } }
          attr { key: 'N' value { i: 3 } }
          """, out3.op.node_def)
  
        with self.assertRaises(ValueError) as cm:
          op_def_library.apply_op("NPolymorphicOut", N=1, T=dtypes.string)
        self.assertEqual(str(cm.exception),
                         "Attr 'N' of 'NPolymorphicOut' Op "
                         "passed 1 less than minimum 2.")
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NPolymorphicOut", N=3, T=[dtypes.string])
        self.assertEqual(
            str(cm.exception),
            "Expected DataType for argument 'T' not [tf.string].")
  
    def testNPolymorphicOutDefault(self):
      with ops.Graph().as_default():
        out1, out2 = op_def_library.apply_op(
            "NPolymorphicOutDefault", N=None, T=None, name="r")
        self.assertEqual(dtypes.bool, out1.dtype)
        self.assertEqual(dtypes.bool, out2.dtype)
        self.assertProtoEquals("""
          name: 'r' op: 'NPolymorphicOutDefault'
          attr { key: 'T' value { type: DT_BOOL } }
          attr { key: 'N' value { i: 2 } }
          """, out1.op.node_def)
  
        out1, out2, out3 = op_def_library.apply_op(
            "NPolymorphicOutDefault", N=3, T=None, name="s")
        self.assertEqual(dtypes.bool, out1.dtype)
        self.assertEqual(dtypes.bool, out2.dtype)
        self.assertEqual(dtypes.bool, out3.dtype)
        self.assertProtoEquals("""
          name: 's' op: 'NPolymorphicOutDefault'
          attr { key: 'T' value { type: DT_BOOL } }
          attr { key: 'N' value { i: 3 } }
          """, out1.op.node_def)
  
        out1, out2 = op_def_library.apply_op(
            "NPolymorphicOutDefault", N=None, T=dtypes.int32, name="t")
        self.assertEqual(dtypes.int32, out1.dtype)
        self.assertEqual(dtypes.int32, out2.dtype)
        self.assertProtoEquals("""
          name: 't' op: 'NPolymorphicOutDefault'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 2 } }
          """, out1.op.node_def)
  
        out1, out2, out3 = op_def_library.apply_op(
            "NPolymorphicOutDefault", N=3, T=dtypes.int32, name="u")
        self.assertEqual(dtypes.int32, out1.dtype)
        self.assertEqual(dtypes.int32, out2.dtype)
        self.assertEqual(dtypes.int32, out3.dtype)
        self.assertProtoEquals("""
          name: 'u' op: 'NPolymorphicOutDefault'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: 'N' value { i: 3 } }
          """, out1.op.node_def)
  
    def testNPolymorphicRestrictOut(self):
      with ops.Graph().as_default():
        out1, out2, out3 = op_def_library.apply_op(
            "NPolymorphicRestrictOut", N=3, T=dtypes.bool, name="u")
        self.assertEqual(dtypes.bool, out1.dtype)
        self.assertEqual(dtypes.bool, out2.dtype)
        self.assertEqual(dtypes.bool, out3.dtype)
        self.assertProtoEquals("""
          name: 'u' op: 'NPolymorphicRestrictOut'
          attr { key: 'T' value { type: DT_BOOL } }
          attr { key: 'N' value { i: 3 } }
          """, out1.op.node_def)
  
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("NPolymorphicRestrictOut", N=2, T=dtypes.int32)
        self.assertEqual(str(cm.exception),
                         "Value passed to parameter 'T' has DataType int32 "
                         "not in list of allowed values: string, bool")
  
    def testRef(self):
      with ops.Graph().as_default():
        out = op_def_library.apply_op("RefOut", T=dtypes.bool, name="o")
        self.assertEqual(dtypes.bool_ref, out.dtype)
        self.assertProtoEquals("""
          name: 'o' op: 'RefOut'
          attr { key: 'T' value { type: DT_BOOL } }
          """, out.op.node_def)
  
        op = op_def_library.apply_op("RefIn", a=out, name="i")
        self.assertProtoEquals("""
          name: 'i' op: 'RefIn' input: 'o'
          attr { key: 'T' value { type: DT_BOOL } }
          attr { key: "_class" value { list { s: "loc:@o" } } }
          """, op.node_def)
  
        # Can pass ref to non-ref input.
        out = op_def_library.apply_op("RefOut", T=dtypes.int32, name="r")
        out = op_def_library.apply_op("Simple", a=out, name="s")
        self.assertProtoEquals("""
          name: 's' op: 'Simple' input: 'r'
          """, out.op.node_def)
  
        # Can't pass non-ref to ref input.
        with self.assertRaises(TypeError) as cm:
          op_def_library.apply_op("RefIn", a=2)
        self.assertEqual(
            str(cm.exception),
            "'RefIn' Op requires that input 'a' be a mutable tensor " +
            "(e.g.: a tf.Variable)")
  
        input_a = op_def_library.apply_op("RefOut", T=dtypes.int32, name="t")
        input_b = op_def_library.apply_op("RefOut", T=dtypes.int32, name="u")
        op = op_def_library.apply_op("TwoRefsIn", a=input_a, b=input_b, name="v")
        # NOTE(mrry): The order of colocation constraints is an implementation
        # detail.
        self.assertProtoEquals("""
          name: 'v' op: 'TwoRefsIn' input: 't' input: 'u'
          attr { key: 'T' value { type: DT_INT32 } }
          attr { key: "_class" value { list { s: "loc:@t" s: "loc:@u" } } }
          """, op.node_def)
  
    def testSpecifyDevice(self):
      graph = ops.Graph()
      with graph.as_default():
        with graph.device("/job:ADevice"):
          op_def_library.apply_op("Simple", a=3)
        # We look at the whole graph here to make sure the Const op is also given
        # the specified device.
        graph_def = graph.as_graph_def()
        self.assertEqual(len(graph_def.node), 2)
        for node in graph_def.node:
          self.assertDeviceEqual(node.device, "/job:ADevice")
  
    def testStructuredOutputSingleList(self):
      with ops.Graph().as_default():
        for n_a in [0, 1, 3]:
          a = op_def_library.apply_op("SimpleStruct", n_a=n_a)
          self.assertIsInstance(a, list)
          self.assertEqual(n_a, len(a))
  
    def testStructuredOutputListAndSingle(self):
      with ops.Graph().as_default():
        for n_a in [0, 1, 3]:
          a, b = op_def_library.apply_op("MixedStruct", n_a=n_a)
          self.assertIsInstance(a, list)
          self.assertEqual(n_a, len(a))
          self.assertTrue(all(x.dtype == dtypes.int32 for x in a))
          self.assertIsInstance(b, ops.Tensor)
          self.assertEqual(dtypes.float32, b.dtype)
  
    def testStructuredOutputMultipleLists(self):
      with ops.Graph().as_default():
        for n_a in [0, 1, 3]:
          for n_b in [0, 1, 3]:
            for t_c in [[],
                        [dtypes.int32],
                        [dtypes.int32, dtypes.float32]]:
              a, b, c = op_def_library.apply_op(
                  "ComplexStruct", n_a=n_a, n_b=n_b, t_c=t_c)
  
              self.assertEqual(n_a, len(a))
              self.assertTrue(all(x.dtype == dtypes.int32 for x in a))
              self.assertEqual(n_b, len(b))
              self.assertTrue(all(x.dtype == dtypes.int64 for x in b))
              self.assertEqual(t_c, [x.dtype for x in c])
  
  
  class OpDefLibraryGraphTest(test_util.TensorFlowTestCase):
  
    def testNoGraph(self):
      out = op_def_library.apply_op("Simple", a=3)
      self.assertEqual(out.graph, ops.get_default_graph())
  
    def testDefaultGraph(self):
      graph = ops.Graph()
      with graph.as_default():
        out = op_def_library.apply_op("Simple", a=3)
        self.assertEqual(out.graph, graph)
  
    def testDifferentGraphFails(self):
      with ops.Graph().as_default():
        a = op_def_library.apply_op("Simple", a=3)
      with ops.Graph().as_default():
        b = op_def_library.apply_op("Simple", a=4)
      with self.assertRaises(ValueError) as cm:
        op_def_library.apply_op("Binary", a=a, b=b)
      self.assertIn("must be from the same graph", str(cm.exception))
  
  
  if __name__ == "__main__":
    googletest.main()
  
apis: tensorflow.python.ops.ragged.ragged_factory_ops.constant tensorflow.python.platform.googletest.main tensorflow.python.ops.ragged.ragged_tensor.RaggedTensor.from_row_splits tensorflow.python.ops.ragged.ragged_tensor.RaggedTensor.from_nested_row_splits tensorflow.python.ops.array_ops.placeholder_with_default tensorflow.python.eager.context.executing_eagerly tensorflow.python.framework.constant_op.constant
apis: tensorflow.python.framework.test_util.run_v1_only numpy.ones tensorflow.python.framework.tensor_spec.TensorSpec tensorflow.python.framework.tensor_shape.TensorShape tensorflow.python.framework.sparse_tensor.SparseTensorSpec tensorflow.python.framework.sparse_tensor.SparseTensor.from_value tensorflow.python.ops.sparse_ops.sparse_transpose tensorflow.python.framework.sparse_tensor.SparseTensor tensorflow.python.framework.sparse_tensor.convert_to_tensor_or_sparse_tensor tensorflow.python.ops.sparse_ops.sparse_tensor_dense_matmul tensorflow.python.ops.array_ops.shape tensorflow.python.framework.sparse_tensor.is_sparse tensorflow.python.ops.math_ops.cast tensorflow.python.framework.sparse_tensor.SparseTensorValue tensorflow.python.eager.context.graph_mode tensorflow.python.framework.ops.convert_to_tensor tensorflow.python.platform.googletest.main numpy.array tensorflow.python.ops.sparse_ops.sparse_tensor_to_dense
apis: tensorflow.python.util.tf_export.tf_export tensorflow.python.framework.ops.get_default_graph tensorflow.python.eager.context.executing_eagerly tensorflow.python.eager.def_function.function tensorflow.python.framework.ops.device
apis: tensorflow.python.util.tf_export.tf_export tensorflow.python.training.server_lib.ClusterSpec
apis: tensorflow.python.eager.context.executing_eagerly tensorflow.python.tpu.tpu.outside_compilation tensorflow.python.tpu.tpu.initialize_system_for_tpu_embedding tensorflow.python.ops.sparse_ops.sparse_slice tensorflow.python.ops.array_ops.pad tensorflow.python.util.nest.flatten tensorflow.python.ops.array_ops.expand_dims tensorflow.python.saved_model.save_context.in_save_context tensorflow.python.framework.tensor_shape.TensorShape tensorflow.python.util.nest.assert_same_structure tensorflow.python.tpu.ops.tpu_ops.is_tpu_embedding_initialized tensorflow.python.util.tf_inspect.getargspec tensorflow.python.distribute.distribute_utils.select_replica tensorflow.python.framework.constant_op.constant tensorflow.python.util.tf_export.tf_export tensorflow.python.training.tracking.base.ShardInfo tensorflow.python.framework.device.DeviceSpec.from_string tensorflow.python.ops.math_ops.reduce_prod tensorflow.python.util.nest.is_nested_or_composite tensorflow.python.ops.embedding_ops.embedding_lookup_sparse_v2 tensorflow.python.framework.ops.device tensorflow.python.ops.embedding_ops.safe_embedding_lookup_sparse_v2 tensorflow.python.util.nest.flatten_with_joined_string_paths tensorflow.python.ops.math_ops.reduce_sum tensorflow.python.ops.array_ops.zeros tensorflow.python.ops.array_ops.shape tensorflow.python.distribute.distribution_strategy_context.get_strategy tensorflow.python.ops.math_ops.cast tensorflow.python.util.compat.as_bytes tensorflow.python.ops.embedding_ops.embedding_lookup_v2 tensorflow.python.framework.ops.init_scope tensorflow.python.ops.variables.Variable tensorflow.python.ops.embedding_ops.embedding_lookup_ragged tensorflow.python.tpu.tpu_embedding_v2_utils.log_tpu_embedding_configuration tensorflow.python.framework.ops.get_default_graph tensorflow.python.ops.array_ops.zeros_like tensorflow.python.ops.array_ops.ones_like tensorflow.core.protobuf.tpu.tpu_embedding_configuration_pb2.TPUEmbeddingConfiguration tensorflow.python.ops.array_ops.stack tensorflow.python.distribute.device_util.get_host_for_device tensorflow.python.ops.array_ops.reshape tensorflow.python.tpu.ops.tpu_ops.enqueue_tpu_embedding_arbitrary_tensor_batch tensorflow.python.util.nest.pack_sequence_as
apis: tensorflow.python.framework.tensor_spec.TensorSpec tensorflow.python.framework.ops.Graph tensorflow.python.framework.ops.get_default_graph tensorflow.python.platform.googletest.main tensorflow.python.framework.tensor_shape.Dimension tensorflow.python.util.compat.as_str tensorflow.python.framework.tensor_shape.TensorShape tensorflow.python.framework.function.Defun tensorflow.python.framework.op_def_library.apply_op tensorflow.core.framework.tensor_shape_pb2.TensorShapeProto