tensorflow/python/tpu/tpu_test.py

# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================

"""Tests for tpu_function helpers."""

from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import importer
from tensorflow.python.framework import ops
from tensorflow.python.layers import convolutional
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import control_flow_util
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import special_math_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.platform import test
from tensorflow.python.tpu import tpu
from tensorflow.python.tpu import tpu_feed
from tensorflow.python.tpu import tpu_replication
from tensorflow.python.tpu import training_loop
from tensorflow.python.tpu.ops import tpu_ops


class TPUContextTest(test.TestCase):

  def testIsInContext(self):
    """Test that control_flow_util can check that we're in a TPU context."""
    with ops.Graph().as_default():
      z1 = array_ops.identity(1)
      pivot = control_flow_ops.no_op()
      context = tpu_replication.TPUReplicateContext(
          b"context", 1, pivot=pivot)
      context.Enter()
      z2 = array_ops.identity(1)
      context.Exit()
      self.assertFalse(control_flow_util.IsInXLAContext(z1.op))
      self.assertTrue(control_flow_util.IsInXLAContext(z2.op))

  def testHandlesNameCollision(self):
    """Test AddValue handles name collisions for ops from different graphs."""
    with ops.Graph().as_default():
      z = array_ops.zeros([2, 3], name="a")
      assert z.name == "a:0", "Expected: a:0, Found: %s" % z.name

      @def_function.function
      def f():
        pivot = control_flow_ops.no_op()
        context = tpu_replication.TPUReplicateContext(
            b"context", 1, pivot=pivot)
        context.Enter()
        array_ops.identity(z)  # Capture z.
        z1 = array_ops.zeros([3, 2], name="a")
        assert z1.name == "a:0", "Expected: a:0, Found: %s" % z1.name
        z2 = array_ops.zeros([3, 2], name="a")
        # Prior to fixing b/166794533 this would fail with a shape mismatch
        # because context.AddValue would have cached `z` by its name which
        # collides with z1's name.
        result = z1 + z2
        context.Exit()
        return result

      f.get_concrete_function()


class TPULayerRewriteTest(test.TestCase):

  def testUsingInfeedQueueWithRegularizer(self):
    """Test that Layer regularizers can reference data created in loops."""

    with ops.Graph().as_default():

      def make_regularizer(scale):
        def regularizer(inputs):
          return scale * math_ops.reduce_sum(math_ops.square(inputs))
        return regularizer

      def training_step(inputs, scale):
        outputs = convolutional.conv2d(
            inputs,
            filters=16,
            kernel_size=(3, 3),
            data_format="channels_first",
            kernel_regularizer=make_regularizer(scale))
        loss = math_ops.reduce_mean(math_ops.square(outputs))
        return loss.op

      inputs = array_ops.zeros(shape=(128, 32, 32, 16))
      scale = array_ops.ones(shape=())
      infeed = tpu_feed.InfeedQueue(
          tuple_types=[dtypes.float32, dtypes.float32],
          tuple_shapes=[inputs.shape, scale.shape])

      def loop():
        return training_loop.repeat(5, training_step, infeed_queue=infeed)

      # This should not throw an error.
      tpu.rewrite(loop)


class TPUGraphPruneTest(test.TestCase):

  def test_prune_unconnected_ops(self):
    with ops.Graph().as_default():
      a = array_ops.placeholder(dtype=dtypes.float32, name="a")
      b = array_ops.placeholder(dtype=dtypes.float32, name="b")
      constant_op.constant(1.0, name="constant")
      x = variable_scope.get_variable(
          name="x",
          dtype=dtypes.float32,
          shape=[],
          use_resource=True,
          initializer=init_ops.constant_initializer(2.0))
      y = variable_scope.get_variable(
          name="y",
          dtype=dtypes.float32,
          shape=[],
          use_resource=True,
          initializer=init_ops.constant_initializer(3.0))
      math_ops.add(a, b)
      math_ops.add(x, y)
      graph_def = ops.get_default_graph().as_graph_def()

      for node in graph_def.node:
        # Attach a TPU_REPLICATE_ATTR to each node.
        node.attr[tpu_replication._TPU_REPLICATE_ATTR].s = b"0"
        # Rewire placeholder "a" and variable "y" leaving them unconnected.
        for (input_index, node_input) in enumerate(node.input):
          if node_input == "b":
            node.input[input_index] = "constant"
          if node_input == "y":
            node.input[input_index] = "x"

    with ops.Graph().as_default() as graph:
      # Reimport the graph and prune unconnected ops.
      importer.import_graph_def(graph_def)
      tpu.prune_unconnected_ops_from_xla(ops.get_default_graph())

      # Verify that ops "a" and "x" still have TPU_REPLICATE_ATTR.
      a = graph.get_operation_by_name("import/a").get_attr(
          tpu_replication._TPU_REPLICATE_ATTR)
      self.assertEqual(b"0", a)
      x = graph.get_operation_by_name("import/x").get_attr(
          tpu_replication._TPU_REPLICATE_ATTR)
      self.assertEqual(b"0", x)
      # Verify that ops "b" and "y" have TPU_REPLICATE_ATTR removed.
      with self.assertRaisesRegex(
          ValueError,
          "Operation \'import/b\' has no attr named \'_tpu_replicate\'"):
        graph.get_operation_by_name("import/b").get_attr(
            tpu_replication._TPU_REPLICATE_ATTR)
      with self.assertRaisesRegex(
          ValueError,
          "Operation \'import/y\' has no attr named \'_tpu_replicate\'"):
        graph.get_operation_by_name("import/y").get_attr(
            tpu_replication._TPU_REPLICATE_ATTR)


class TPUOpsTest(test.TestCase):

  def test_all_to_all_zero_split_count(self):
    with self.assertRaisesRegex(
        ValueError, "split_count 0 must at least be one"):
      tpu_ops.all_to_all(
          x=[0.0, 0.1652, 0.6543],
          group_assignment=[1, -1],
          concat_dimension=0,
          split_dimension=0,
          split_count=0)

  def test_all_to_all_group_assignment_wrong_shape(self):
    with self.assertRaisesRegex(
        ValueError, "group_assignment must have rank 2"):
      tpu_ops.all_to_all(
          x=[0.0, 0.1652, 0.6543],
          group_assignment=[1, -1],
          concat_dimension=0,
          split_dimension=0,
          split_count=2)

  def test_all_to_all_split_count_not_equal_to_group_assignment_shape(self):
    with self.assertRaisesRegex(
        ValueError, "split_count 1 must equal the size of the second dimension "
        "of group_assignment 2"):
      tpu_ops.all_to_all(
          x=[0.0, 0.1652, 0.6543],
          group_assignment=[[0, 1], [2, 3]],
          concat_dimension=0,
          split_dimension=0,
          split_count=1)

  def test_all_to_all_split_count_not_divide_input_shape(self):
    with self.assertRaisesRegex(
        ValueError, "input dimension 3 not divisible by split_count 2"):
      tpu_ops.all_to_all(
          x=[[0.0], [0.1652], [0.6543]],
          group_assignment=[[0, 1], [2, 3]],
          concat_dimension=1,
          split_dimension=0,
          split_count=2)


def do_einsum():
  a = array_ops.placeholder(dtype=dtypes.float32, name="a", shape=[2, 3, 4])
  b = array_ops.placeholder(dtype=dtypes.float32, name="b", shape=[2, 4, 5])
  return special_math_ops.einsum("abc,acd->abd", a, b)


def find_einsum(g):
  graph_def = g.as_graph_def()
  for node in graph_def.node:
    if node.op == "Einsum":
      return True
  return False


def find_xla_einsum(g):
  graph_def = g.as_graph_def()
  for node in graph_def.node:
    if node.op == "XlaEinsum":
      return True
  return False


class TPUXlaEinsumTest(test.TestCase):

  def test_tpu_rewrite_uses_xla_einsum(self):
    with ops.Graph().as_default() as g:
      tpu.rewrite(do_einsum)
      self.assertTrue(find_einsum(g) or find_xla_einsum(g))

  def test_default_does_not_use_xla_einsum(self):
    with ops.Graph().as_default() as g:
      do_einsum()
      self.assertFalse(find_xla_einsum(g))


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