sharding_util_ops_test.py

# Copyright 2021 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 sharding util ops (XlaSplitND, XlaConcatND)."""

from typing import Any, List, Optional

from absl.testing import parameterized
import numpy as np

from tensorflow.compiler.tests import xla_test
from tensorflow.python.client.session import Session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework.ops import control_dependencies
from tensorflow.python.framework.ops import Tensor
from tensorflow.python.ops import gen_tpu_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test


def create_tensor_split_graph(
    sess: Session,
    input_value: Any,
    input_dtype: Any,
    num_outputs: int,
    num_splits: List[int],
    paddings: Optional[List[int]] = None) -> List[Tensor]:
  del sess

  const_input_op = constant_op.constant(input_value, dtype=input_dtype)
  return gen_tpu_ops.xla_split_nd(
      const_input_op, num_outputs, num_splits, paddings=paddings)


def create_resource_split_graph(
    sess: Session,
    input_value: Any,
    input_dtype: Any,
    num_outputs: int,
    num_splits: List[int],
    paddings: Optional[List[int]] = None) -> List[Tensor]:
  variable = resource_variable_ops.ResourceVariable(
      initial_value=input_value, dtype=input_dtype)
  sess.run(variables.variables_initializer([variable]))
  return gen_tpu_ops.read_variable_xla_split_nd(
      variable.handle, input_dtype, num_outputs, num_splits, paddings=paddings)


class XlaSplitNDOpTest(xla_test.XLATestCase, parameterized.TestCase):

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testSplitDimensionZero(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[[0]]],
            input_dtype=dtype,
            num_outputs=1,
            num_splits=[1, 1, 0])
        with self.assertRaisesOpError('index 2 must be positive, but got 0'):
          sess.run(split)

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testSplitDimensionNegative(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[[0]]],
            input_dtype=dtype,
            num_outputs=1,
            num_splits=[1, -1, 1])
        with self.assertRaisesOpError('index 1 must be positive, but got -1'):
          sess.run(split)

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testNumOutputsMismatch(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[0, 1],
            input_dtype=dtype,
            num_outputs=1,
            num_splits=[2])
        with self.assertRaisesOpError('\'N\' must match number of slices 2'):
          sess.run(split)

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testPaddingsLengthMismatch(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[0, 1], [2, 3]],
            input_dtype=dtype,
            num_outputs=4,
            num_splits=[2, 2],
            paddings=[0])
        with self.assertRaisesOpError('length 2, but got 1'):
          sess.run(split)

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testPaddingsNegative(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[0, 1], [2, 3]],
            input_dtype=dtype,
            num_outputs=4,
            num_splits=[2, 2],
            paddings=[0, -1])
        with self.assertRaisesOpError('non-negative, but got -1 at index 1'):
          sess.run(split)

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testInputRankSplitMismatch(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[0, 1], [2, 3]],
            input_dtype=dtype,
            num_outputs=8,
            num_splits=[2, 2, 2])
        with self.assertRaisesOpError(
            '\'num_splits\' length 3, but got rank 2'):
          sess.run(split)

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testDimNotEvenlySplit(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[0, 1], [2, 3], [4, 5], [6, 7]],
            input_dtype=dtype,
            num_outputs=6,
            num_splits=[3, 2])
        with self.assertRaisesOpError('divisible by \'num_splits\' 3'):
          sess.run(split)

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testDimWithPaddingNotEvenlySplit(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[0, 1], [2, 3], [4, 5], [6, 7]],
            input_dtype=dtype,
            num_outputs=4,
            num_splits=[2, 2],
            paddings=[0, 1])
        with self.assertRaisesOpError('divisible by \'num_splits\' 2'):
          sess.run(split)

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testNoSplits(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
            input_dtype=dtype,
            num_outputs=1,
            num_splits=[1, 1, 1])
        results = sess.run(split)
      self.assertLen(results, 1)
      self.assertAllClose(results[0], [[[0, 1], [2, 3]], [[4, 5], [6, 7]]])

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testNoSplitsWithPadding(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[[0]], [[1]]],
            input_dtype=dtype,
            num_outputs=1,
            num_splits=[1, 1, 1],
            paddings=[0, 1, 1])
        results = sess.run(split)
      self.assertLen(results, 1)
      self.assertAllClose(results[0], [[[0, 0], [0, 0]], [[1, 0], [0, 0]]])

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testSplitNoPadding(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[
                [0, 1, 2, 3],
                [4, 5, 6, 7],
                [8, 9, 10, 11],
                [12, 13, 14, 15],
            ],
            input_dtype=dtype,
            num_outputs=4,
            num_splits=[2, 2])
        results = sess.run(split)
      self.assertLen(results, 4)
      self.assertAllClose(results[0], [[0, 1], [4, 5]])
      self.assertAllClose(results[1], [[2, 3], [6, 7]])
      self.assertAllClose(results[2], [[8, 9], [12, 13]])
      self.assertAllClose(results[3], [[10, 11], [14, 15]])

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testSplitPartialPadding(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[
                [0, 1, 2],
                [3, 4, 5],
                [6, 7, 8],
            ],
            input_dtype=dtype,
            num_outputs=4,
            num_splits=[2, 2],
            paddings=[1, 1])
        results = sess.run(split)
      self.assertLen(results, 4)
      self.assertAllClose(results[0], [[0, 1], [3, 4]])
      self.assertAllClose(results[1], [[2, 0], [5, 0]])
      self.assertAllClose(results[2], [[6, 7], [0, 0]])
      self.assertAllClose(results[3], [[8, 0], [0, 0]])

  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
                                  ('Resource', create_resource_split_graph))
  def testSplitCompletePadding(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=[[0], [1]],
            input_dtype=dtype,
            num_outputs=4,
            num_splits=[2, 2],
            paddings=[2, 3])
        results = sess.run(split)
      self.assertLen(results, 4)
      self.assertAllClose(results[0], [[0, 0], [1, 0]])
      self.assertAllClose(results[1], [[0, 0], [0, 0]])
      self.assertAllClose(results[2], [[0, 0], [0, 0]])
      self.assertAllClose(results[3], [[0, 0], [0, 0]])

  @parameterized.named_parameters(
      ('1Tensor', create_tensor_split_graph, 1),
      ('2Tensor', create_tensor_split_graph, 2),
      ('3Tensor', create_tensor_split_graph, 3),
      ('4Tensor', create_tensor_split_graph, 4),
      ('5Tensor', create_tensor_split_graph, 5),
      ('6Tensor', create_tensor_split_graph, 6),
      ('7Tensor', create_tensor_split_graph, 7),
      ('8Tensor', create_tensor_split_graph, 8),
      ('1Resource', create_resource_split_graph, 1),
      ('2Resource', create_resource_split_graph, 2),
      ('3Resource', create_resource_split_graph, 3),
      ('4Resource', create_resource_split_graph, 4),
      ('5Resource', create_resource_split_graph, 5),
      ('6Resource', create_resource_split_graph, 6),
      ('7Resource', create_resource_split_graph, 7),
      ('8Resource', create_resource_split_graph, 8),
  )
  def testRanked(self, graph_fn, rank):
    num_splits = [2] * rank
    num_outputs = 2 << (rank - 1)
    input_value = np.reshape(np.arange(np.prod(num_splits)), num_splits)
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        split = graph_fn(
            sess,
            input_value=input_value,
            input_dtype=dtype,
            num_outputs=num_outputs,
            num_splits=num_splits)
        results = sess.run(split)
      self.assertLen(results, num_outputs)
      for i, result in enumerate(results):
        expected_output = np.reshape(i, [1] * rank).astype(dtype)
        self.assertAllClose(result, expected_output)


def create_tensor_concat_graph(
    sess: Session,
    input_values: List[Any],
    input_dtype: Any,
    num_concats: List[int],
    paddings: Optional[List[int]] = None,
    output_shape: Optional[List[int]] = None) -> Tensor:
  del sess
  del output_shape

  const_input_ops = [
      constant_op.constant(i, dtype=input_dtype) for i in input_values
  ]
  return gen_tpu_ops.xla_concat_nd(const_input_ops, num_concats, paddings)


def create_resource_concat_graph(
    sess: Session,
    input_values: List[Any],
    input_dtype: Any,
    num_concats: List[int],
    paddings: Optional[List[int]] = None,
    output_shape: Optional[List[int]] = None) -> Tensor:
  variable_shape = [] if output_shape is None else output_shape
  variable = resource_variable_ops.ResourceVariable(
      initial_value=np.zeros(variable_shape, dtype=input_dtype),
      dtype=input_dtype)
  sess.run(variables.variables_initializer([variable]))
  const_input_ops = [
      constant_op.constant(i, dtype=input_dtype) for i in input_values
  ]
  concat = gen_tpu_ops.assign_variable_xla_concat_nd(variable.handle,
                                                     const_input_ops,
                                                     num_concats, paddings)
  with control_dependencies([concat]):
    return variable.read_value()


class XlaConcatNDOpTest(xla_test.XLATestCase, parameterized.TestCase):

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testConcatDimensionZero(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[[[[0]]]],
            input_dtype=dtype,
            num_concats=[1, 1, 0])
        with self.assertRaisesOpError('index 2 must be positive, but got 0'):
          sess.run(concat)

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testConcatDimensionNegative(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[[[[0]]]],
            input_dtype=dtype,
            num_concats=[1, -1, 1])
        with self.assertRaisesOpError('index 1 must be positive, but got -1'):
          sess.run(concat)

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testNumInputsMismatch(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess, input_values=[[0, 1]], input_dtype=dtype, num_concats=[2])
        with self.assertRaisesOpError('\'N\' must match number of slices 2'):
          sess.run(concat)

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testPaddingsLengthMismatch(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[[[0, 1], [2, 3]]],
            input_dtype=dtype,
            num_concats=[1, 1],
            paddings=[0])
        with self.assertRaisesOpError('length 2, but got 1'):
          sess.run(concat)

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testPaddingsNegative(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[[[0, 1], [2, 3]]],
            input_dtype=dtype,
            num_concats=[1, 1],
            paddings=[0, -1])
        with self.assertRaisesOpError('non-negative, but got -1 at index 1'):
          sess.run(concat)

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testInputRankConcatMismatch(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess, input_values=[[0]], input_dtype=dtype, num_concats=[1, 1])
        with self.assertRaisesOpError(
            '\'num_concats\' length 2, but got rank 1'):
          sess.run(concat)

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testDifferentShapedInputs(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[[0], [1, 2]],
            input_dtype=dtype,
            num_concats=[2])
        with self.assertRaisesOpError(
            r'same expected shape \[1\], but got \[2\] at index 1'):
          sess.run(concat)

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testPaddingExceedsOutputDimSize(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[[0]],
            input_dtype=dtype,
            num_concats=[1],
            paddings=[2])
        with self.assertRaisesOpError(
            'exceed expected output shape dimension 1 at index 0, but got 2'):
          sess.run(concat)

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testNoConcats(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]],
            input_dtype=dtype,
            num_concats=[1, 1, 1],
            output_shape=[2, 2, 2])
        result = sess.run(concat)
      self.assertAllClose(result, [[[0, 1], [2, 3]], [[4, 5], [6, 7]]])

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testNoConcatsWithPadding(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]],
            input_dtype=dtype,
            num_concats=[1, 1, 1],
            output_shape=[1, 1, 1],
            paddings=[1, 1, 1])
        result = sess.run(concat)
      self.assertAllClose(result, [[[0]]])

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testConcatNoPadding(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[
                [[0, 1], [2, 3]],
                [[4, 5], [6, 7]],
                [[8, 9], [10, 11]],
                [[12, 13], [14, 15]],
            ],
            input_dtype=dtype,
            num_concats=[2, 2],
            output_shape=[4, 4])
        result = sess.run(concat)
      self.assertAllClose(
          result,
          [[0, 1, 4, 5], [2, 3, 6, 7], [8, 9, 12, 13], [10, 11, 14, 15]])

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testConcatPartialPadding(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[
                [[0, 1], [2, 3]],
                [[4, 5], [6, 7]],
                [[8, 9], [10, 11]],
                [[12, 13], [14, 15]],
            ],
            input_dtype=dtype,
            num_concats=[2, 2],
            output_shape=[3, 3],
            paddings=[1, 1])
        result = sess.run(concat)
      self.assertAllClose(result, [[0, 1, 4], [2, 3, 6], [8, 9, 12]])

  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
                                  ('Resource', create_resource_concat_graph))
  def testConcatCompletePadding(self, graph_fn):
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=[
                [[0, 1], [2, 3]],
                [[4, 5], [6, 7]],
                [[8, 9], [10, 11]],
                [[12, 13], [14, 15]],
            ],
            input_dtype=dtype,
            num_concats=[2, 2],
            output_shape=[2, 2],
            paddings=[2, 2])
        result = sess.run(concat)
      self.assertAllClose(result, [[0, 1], [2, 3]])

  @parameterized.named_parameters(
      ('1Tensor', create_tensor_concat_graph, 1),
      ('2Tensor', create_tensor_concat_graph, 2),
      ('3Tensor', create_tensor_concat_graph, 3),
      ('4Tensor', create_tensor_concat_graph, 4),
      ('5Tensor', create_tensor_concat_graph, 5),
      ('6Tensor', create_tensor_concat_graph, 6),
      ('7Tensor', create_tensor_concat_graph, 7),
      ('8Tensor', create_tensor_concat_graph, 8),
      ('1Resource', create_resource_concat_graph, 1),
      ('2Resource', create_resource_concat_graph, 2),
      ('3Resource', create_resource_concat_graph, 3),
      ('4Resource', create_resource_concat_graph, 4),
      ('5Resource', create_resource_concat_graph, 5),
      ('6Resource', create_resource_concat_graph, 6),
      ('7Resource', create_resource_concat_graph, 7),
      ('8Resource', create_resource_concat_graph, 8),
  )
  def testRanked(self, graph_fn, rank):
    num_concats = [2] * rank
    num_inputs = 2 << (rank - 1)
    input_values = [np.reshape(i, [1] * rank) for i in range(num_inputs)]
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        concat = graph_fn(
            sess,
            input_values=input_values,
            input_dtype=dtype,
            num_concats=num_concats,
            output_shape=num_concats)
        result = sess.run(concat)
      expected_output = np.arange(0,
                                  num_inputs).reshape(num_concats).astype(dtype)
      self.assertAllClose(result, expected_output)


def create_tensor_roundtrip_graph(
    sess: Session,
    value: Any,
    dtype: Any,
    num_partitions: List[int],
    paddings: Optional[List[int]] = None) -> Tensor:
  del sess

  const_input_op = constant_op.constant(value, dtype=dtype)
  split = gen_tpu_ops.xla_split_nd(
      const_input_op,
      np.prod(num_partitions),
      num_partitions,
      paddings=paddings)
  concat = gen_tpu_ops.xla_concat_nd(split, num_partitions, paddings)
  return math_ops.equal(const_input_op, concat)


def create_resource_roundtrip_graph(
    sess: Session,
    value: Any,
    dtype: Any,
    num_partitions: List[int],
    paddings: Optional[List[int]] = None) -> Tensor:
  variable = resource_variable_ops.ResourceVariable(
      initial_value=value, dtype=dtype)
  sess.run(variables.variables_initializer([variable]))
  split = gen_tpu_ops.read_variable_xla_split_nd(
      variable.handle,
      dtype,
      np.prod(num_partitions),
      num_partitions,
      paddings=paddings)
  concat = gen_tpu_ops.assign_variable_xla_concat_nd(variable.handle, split,
                                                     num_partitions, paddings)
  with control_dependencies([concat]):
    return math_ops.equal(variable.read_value(),
                          constant_op.constant(value, dtype=dtype))


class XlaSplitConcatNDTest(xla_test.XLATestCase, parameterized.TestCase):

  @parameterized.named_parameters(
      ('1Tensor', create_tensor_roundtrip_graph, 1),
      ('2Tensor', create_tensor_roundtrip_graph, 2),
      ('3Tensor', create_tensor_roundtrip_graph, 3),
      ('4Tensor', create_tensor_roundtrip_graph, 4),
      ('5Tensor', create_tensor_roundtrip_graph, 5),
      ('6Tensor', create_tensor_roundtrip_graph, 6),
      ('7Tensor', create_tensor_roundtrip_graph, 7),
      ('8Tensor', create_tensor_roundtrip_graph, 8),
      ('1Resource', create_resource_roundtrip_graph, 1),
      ('2Resource', create_resource_roundtrip_graph, 2),
      ('3Resource', create_resource_roundtrip_graph, 3),
      ('4Resource', create_resource_roundtrip_graph, 4),
      ('5Resource', create_resource_roundtrip_graph, 5),
      ('6Resource', create_resource_roundtrip_graph, 6),
      ('7Resource', create_resource_roundtrip_graph, 7),
      ('8Resource', create_resource_roundtrip_graph, 8),
  )
  def testNoPadding(self, graph_fn, rank):
    num_partitions = [2] * rank
    shape = [4] * rank
    value = np.arange(0, np.prod(shape)).reshape(shape)
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        validate = graph_fn(sess, value, dtype, num_partitions)
        result = sess.run(validate)
      self.assertAllEqual(result, np.broadcast_to(True, shape))

  @parameterized.named_parameters(
      ('1Tensor', create_tensor_roundtrip_graph, 1),
      ('2Tensor', create_tensor_roundtrip_graph, 2),
      ('3Tensor', create_tensor_roundtrip_graph, 3),
      ('4Tensor', create_tensor_roundtrip_graph, 4),
      ('5Tensor', create_tensor_roundtrip_graph, 5),
      ('6Tensor', create_tensor_roundtrip_graph, 6),
      ('7Tensor', create_tensor_roundtrip_graph, 7),
      ('8Tensor', create_tensor_roundtrip_graph, 8),
      ('1Resource', create_resource_roundtrip_graph, 1),
      ('2Resource', create_resource_roundtrip_graph, 2),
      ('3Resource', create_resource_roundtrip_graph, 3),
      ('4Resource', create_resource_roundtrip_graph, 4),
      ('5Resource', create_resource_roundtrip_graph, 5),
      ('6Resource', create_resource_roundtrip_graph, 6),
      ('7Resource', create_resource_roundtrip_graph, 7),
      ('8Resource', create_resource_roundtrip_graph, 8),
  )
  def testPartialPadding(self, graph_fn, rank):
    num_partitions = [2] * rank
    shape = [4] * rank
    value = np.arange(0, np.prod(shape)).reshape(shape)
    paddings = [2] * rank
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        validate = graph_fn(sess, value, dtype, num_partitions, paddings)
        result = sess.run(validate)
      self.assertAllEqual(result, np.broadcast_to(True, shape))

  @parameterized.named_parameters(
      ('1Tensor', create_tensor_roundtrip_graph, 1),
      ('2Tensor', create_tensor_roundtrip_graph, 2),
      ('3Tensor', create_tensor_roundtrip_graph, 3),
      ('4Tensor', create_tensor_roundtrip_graph, 4),
      ('5Tensor', create_tensor_roundtrip_graph, 5),
      ('6Tensor', create_tensor_roundtrip_graph, 6),
      ('7Tensor', create_tensor_roundtrip_graph, 7),
      ('8Tensor', create_tensor_roundtrip_graph, 8),
      ('1Resource', create_resource_roundtrip_graph, 1),
      ('2Resource', create_resource_roundtrip_graph, 2),
      ('3Resource', create_resource_roundtrip_graph, 3),
      ('4Resource', create_resource_roundtrip_graph, 4),
      ('5Resource', create_resource_roundtrip_graph, 5),
      ('6Resource', create_resource_roundtrip_graph, 6),
      ('7Resource', create_resource_roundtrip_graph, 7),
      ('8Resource', create_resource_roundtrip_graph, 8),
  )
  def testCompletePadding(self, graph_fn, rank):
    num_partitions = [2] * rank
    shape = [4] * rank
    value = np.arange(0, np.prod(shape)).reshape(shape)
    paddings = [4] * rank
    for dtype in self.numeric_types:
      with self.session() as sess, self.device_scope():
        validate = graph_fn(sess, value, dtype, num_partitions, paddings)
        result = sess.run(validate)
      self.assertAllEqual(result, np.broadcast_to(True, shape))


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