# 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.
# ==============================================================================
"""Tests for Unified APIs' python bindings."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import timeit

from absl.testing import parameterized

from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework.experimental import _unified_api
from tensorflow.python.framework.experimental import context_stack as context_lib
from tensorflow.python.framework.experimental import def_function
from tensorflow.python.framework.experimental import math_ops as unified_math_ops
from tensorflow.python.framework.experimental import nn_ops as unified_nn_ops
from tensorflow.python.framework.experimental import tape as tape_lib
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_grad  # pylint: disable=unused-import
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.platform import test

SetTracingImplementation = _unified_api.SetTracingImplementation
TensorCastHelper = _unified_api.EagerTensorToImmediateExecutionTensorHandle


def get_immediate_execution_context():
  context.context().ensure_initialized()
  return _unified_api.EagerContextToImmediateExecutionContext(
      context.context()._handle)


def maybe_cast(t, perform_cast):
  if perform_cast:
    return TensorCastHelper(t)
  return t


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

  @parameterized.named_parameters([
      ("Graph", False),
      ("Mlir", True),
  ])
  def testAdd(self, use_mlir):
    if use_mlir:
      SetTracingImplementation("mlir")

    def model(a, b):
      return unified_math_ops.add(a, b)

    with context_lib.set_default(get_immediate_execution_context()):
      a = TensorCastHelper(constant_op.constant([1., 2.]))
      b = TensorCastHelper(constant_op.constant([3., 4.]))

      func_output = def_function.function(model)(a, b)
      self.assertAllEqual(func_output.numpy(), [4., 6.])

      eager_output = model(a, b)
      self.assertAllEqual(eager_output.numpy(), [4., 6.])

  @parameterized.named_parameters([
      ("Graph", False),
      ("Mlir", True),
  ])
  def testAddGrad(self, use_mlir):
    if use_mlir:
      SetTracingImplementation("mlir")

    def model(a, b):
      with tape_lib.GradientTape() as tape:
        tape.watch(a)
        tape.watch(b)
        result = unified_math_ops.add(a, b)
      grads = tape.gradient(result, [a, b])
      return grads

    with context_lib.set_default(get_immediate_execution_context()):
      a = TensorCastHelper(constant_op.constant([1., 2.]))
      b = TensorCastHelper(constant_op.constant([3., 4.]))

      func_outputs = def_function.function(model)(a, b)
      self.assertAllEqual(func_outputs[0].numpy(), [1.0, 1.0])
      self.assertAllEqual(func_outputs[1].numpy(), [1.0, 1.0])

      eager_outputs = model(a, b)
      self.assertAllEqual(eager_outputs[0].numpy(), [1.0, 1.0])
      self.assertAllEqual(eager_outputs[1].numpy(), [1.0, 1.0])

  @parameterized.named_parameters([
      ("Graph", False),
      ("Mlir", True),
  ])
  def testRelu(self, use_mlir):
    if use_mlir:
      SetTracingImplementation("mlir")

    def model(t):
      return unified_nn_ops.relu(t)

    with context_lib.set_default(get_immediate_execution_context()):
      positive = TensorCastHelper(constant_op.constant([1.]))
      negative = TensorCastHelper(constant_op.constant([-1.]))

      model_fn = def_function.function(model)
      func_output = model_fn(positive)
      self.assertAllEqual(func_output.numpy(), [1.])
      func_output = model_fn(negative)
      self.assertAllEqual(func_output.numpy(), [0.])

      eager_output = model(positive)
      self.assertAllEqual(eager_output.numpy(), [1.])
      eager_output = model(negative)
      self.assertAllEqual(eager_output.numpy(), [0.])

  @parameterized.named_parameters([
      ("Graph", False),
      ("Mlir", True),
  ])
  def testReluGrad(self, use_mlir):
    if use_mlir:
      SetTracingImplementation("mlir")

    def model(t):
      with tape_lib.GradientTape() as tape:
        tape.watch(t)
        result = unified_nn_ops.relu(t)
      grads = tape.gradient(result, t)
      return grads

    with context_lib.set_default(get_immediate_execution_context()):
      positive = TensorCastHelper(constant_op.constant([1.]))
      negative = TensorCastHelper(constant_op.constant([-1.]))

      model_fn = def_function.function(model)
      func_output = model_fn(positive)
      self.assertAllEqual(func_output.numpy(), [1.])
      func_output = model_fn(negative)
      self.assertAllEqual(func_output.numpy(), [0.])

      eager_output = model(positive)
      self.assertAllEqual(eager_output.numpy(), [1.])
      eager_output = model(negative)
      self.assertAllEqual(eager_output.numpy(), [0.])

  @parameterized.named_parameters([
      ("Graph", False),
      ("Mlir", True),
  ])
  def testNeg(self, use_mlir):
    if use_mlir:
      SetTracingImplementation("mlir")

    def model(a):
      return unified_math_ops.neg(a)

    with context_lib.set_default(get_immediate_execution_context()):
      a = TensorCastHelper(constant_op.constant([2.]))

      func_output = def_function.function(model)(a)
      self.assertAllEqual(func_output.numpy(), [-2.])

      eager_output = model(a)
      self.assertAllEqual(eager_output.numpy(), [-2.])

  @parameterized.named_parameters([
      ("Graph", False),
      ("Mlir", True),
  ])
  def testNegGrad(self, use_mlir):
    if use_mlir:
      SetTracingImplementation("mlir")

    def model(a):
      with tape_lib.GradientTape() as tape:
        tape.watch(a)
        result = unified_math_ops.neg(a)
      grads = tape.gradient(result, a)
      return grads

    with context_lib.set_default(get_immediate_execution_context()):
      a = TensorCastHelper(constant_op.constant([2.]))

      func_outputs = def_function.function(model)(a)
      self.assertAllEqual(func_outputs.numpy(), [-1.0])

      eager_outputs = model(a)
      self.assertAllEqual(eager_outputs.numpy(), [-1.0])

  @parameterized.named_parameters([
      ("Graph", False),
      ("Mlir", True),
  ])
  def testSub(self, use_mlir):
    if use_mlir:
      SetTracingImplementation("mlir")

    def model(a, b):
      return unified_math_ops.sub(a, b)

    with context_lib.set_default(get_immediate_execution_context()):
      a = TensorCastHelper(constant_op.constant([1., 2.]))
      b = TensorCastHelper(constant_op.constant([3., 4.]))

      func_output = def_function.function(model)(a, b)
      self.assertAllEqual(func_output.numpy(), [-2., -2.])

      eager_output = model(a, b)
      self.assertAllEqual(eager_output.numpy(), [-2., -2.])

  @parameterized.named_parameters([
      ("Graph", False),
      ("Mlir", True),
  ])
  def testSubGrad(self, use_mlir):
    if use_mlir:
      SetTracingImplementation("mlir")

    def model(a, b):
      with tape_lib.GradientTape() as tape:
        tape.watch(a)
        tape.watch(b)
        result = unified_math_ops.sub(a, b)
      grads = tape.gradient(result, [a, b])
      return grads

    with context_lib.set_default(get_immediate_execution_context()):
      a = TensorCastHelper(constant_op.constant([1., 2.]))
      b = TensorCastHelper(constant_op.constant([3., 4.]))

      func_outputs = def_function.function(model)(a, b)
      self.assertAllEqual(func_outputs[0].numpy(), [1.0, 1.0])
      self.assertAllEqual(func_outputs[1].numpy(), [-1.0, -1.0])

      eager_outputs = model(a, b)
      self.assertAllEqual(eager_outputs[0].numpy(), [1.0, 1.0])
      self.assertAllEqual(eager_outputs[1].numpy(), [-1.0, -1.0])


class UnifiedTapeBenchmark(test.Benchmark):

  def _computeMnistMlpGrads(self, math_ops_lib, nn_ops_lib, backprop_lib, cast,
                            num_iters, hidden_layers, hidden_size, batch_size):
    batch_size = 1
    image_size = 28 * 28
    num_classes = 10

    def model(x, hidden_weights, softmax_weight, labels):
      with backprop_lib.GradientTape() as tape:
        for weight in hidden_weights + [softmax_weight]:
          tape.watch(weight)
        for hidden_weight in hidden_weights:
          x = math_ops_lib.mat_mul(x, hidden_weight)
          x = nn_ops_lib.relu(x)
        logits = math_ops_lib.mat_mul(x, softmax_weight)
        loss = nn_ops_lib.sparse_softmax_cross_entropy_with_logits(
            logits=logits, labels=labels)

      grads = tape.gradient(loss, hidden_weights + [softmax_weight])
      return grads

    x = maybe_cast(array_ops.ones([batch_size, image_size]), cast)
    hidden_weights = []
    for i in range(hidden_layers):
      hidden_weights.append(
          maybe_cast(
              random_ops.random_uniform(
                  [hidden_size if i else image_size, hidden_size]), cast))
    softmax_weight = maybe_cast(
        random_ops.random_uniform([hidden_size, num_classes]), cast)
    labels = maybe_cast(array_ops.zeros([batch_size], dtype=dtypes.int32), cast)

    with context_lib.set_default(get_immediate_execution_context()):
      # Warm up.
      for _ in range(10):
        model(x, hidden_weights, softmax_weight, labels)
      runtimes = timeit.repeat(
          lambda: model(x, hidden_weights, softmax_weight, labels),
          repeat=num_iters,
          number=10)
    return min(runtimes) / 10

  def benchmarkTwoHiddenLayerMnistEagerUnified(self):
    num_iters = 100
    duration = self._computeMnistMlpGrads(
        unified_math_ops,
        unified_nn_ops,
        tape_lib,
        True,
        num_iters,
        hidden_layers=2,
        hidden_size=100,
        batch_size=1)
    self.report_benchmark(
        name="TwoHiddenLayerMnistEagerUnified",
        iters=num_iters,
        wall_time=duration)

  def benchmarkTwoHiddenLayerMnistEager(self):
    num_iters = 100
    duration = self._computeMnistMlpGrads(
        math_ops,
        nn_ops,
        backprop,
        False,
        num_iters,
        hidden_layers=2,
        hidden_size=100,
        batch_size=1)
    self.report_benchmark(
        name="TwoHiddenLayerMnistEager", iters=num_iters, wall_time=duration)

  def benchmarkTenHiddenLayerMnistEagerUnified(self):
    num_iters = 100
    duration = self._computeMnistMlpGrads(
        unified_math_ops,
        unified_nn_ops,
        tape_lib,
        True,
        num_iters,
        hidden_layers=10,
        hidden_size=100,
        batch_size=1)
    self.report_benchmark(
        name="TenHiddenLayerMnistEagerUnified",
        iters=num_iters,
        wall_time=duration)

  def benchmarkTenHiddenLayerMnistEager(self):
    num_iters = 100
    duration = self._computeMnistMlpGrads(
        math_ops,
        nn_ops,
        backprop,
        False,
        num_iters,
        hidden_layers=10,
        hidden_size=100,
        batch_size=1)
    self.report_benchmark(
        name="TenHiddenLayerMnistEager", iters=num_iters, wall_time=duration)


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