# Copyright 2019 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.
# ==============================================================================
"""Unit tests for op_callback."""

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

import threading

import numpy as np

from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import backprop
from tensorflow.python.eager import def_function
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import op_callbacks
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import script_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.ops import variables
from tensorflow.python.util import compat

# Keep all the hard-coded op type strings in one place so they are easy to
# change all at once in the face of any possible future op type name changes.
_ADD_OP = b"AddV2"
_ASSIGN_ADD_VARIABLE_OP = b"AssignAddVariableOp"
_CONSTANT_OP = b"Const"
_COS_OP = b"Cos"
_GREATER_OP = b"Greater"
_IDENTITY_OP = b"Identity"
_IF_OP = b"If"
_LESS_OP = b"Less"
_LOG_OP = b"Log"
_MATMUL_OP = b"MatMul"
_MUL_OP = b"Mul"
_PLACEHOLDER_OP = b"Placeholder"
_POW_OP = b"Pow"
_READ_VARIALBE_OP = b"ReadVariableOp"
_SIN_OP = b"Sin"
_SPARSE_TENSOR_DENSE_MATMUL_OP = b"SparseTensorDenseMatMul"
_SQRT_OP = b"Sqrt"
_SQUARE_OP = b"Square"
_STATELESS_IF_OP = b"StatelessIf"
_UNIQUE_OP = b"Unique"
_WHILE_OP = b"While"


class _NumpyFunctionCallback(object):

  def __init__(self, instrument_graph_ops=True):
    self.instrument_graph_ops = instrument_graph_ops
    self.reset()

  def callback(self, op_type, inputs, attrs, outputs, op_name=None, graph=None):
    is_eager = not graph
    if is_eager:
      self.eager_op_types.append(
          compat.as_bytes(op_type) if op_type else op_type)
      self.eager_op_names.append(
          compat.as_bytes(op_name) if op_name else op_name)
      self.eager_attrs.append(attrs)
      self.eager_graphs.append(graph)
      self.eager_inputs.append(inputs)
    else:
      self.graph_op_types.append(
          compat.as_bytes(op_type) if op_type else op_type)
      self.graph_op_names.append(
          compat.as_bytes(op_name) if op_name else op_name)
      self.graph_attrs.append(attrs)
      self.graph_graphs.append(graph)
      self.graph_graph_versions.append(graph.version)
      self.graph_inputs.append(inputs)

      if not self.instrument_graph_ops:
        return outputs

      # Instrument the graph with numpy_function.
      instrumented_outputs = []
      for output in outputs:
        if compat.as_bytes(op_type) in (
            _IF_OP, _STATELESS_IF_OP, _WHILE_OP, _IDENTITY_OP):
          # TODO(cais): Overriding the output of StatelessIf, If and While ops
          # currently fails with error. Investigate (b/139668453).
          # Avoid instrumenting Identity ops as well, as they are inserted
          # by tf.function/AutoGraph for marshalling outputs.
          instrumented_output = output
        else:

          def record(ndarray_value):
            if compat.as_bytes(op_name) not in self.graph_internal_ndarrays:
              self.graph_internal_ndarrays[compat.as_bytes(op_name)] = []
            self.graph_internal_ndarrays[compat.as_bytes(op_name)].append(
                ndarray_value)
            return ndarray_value

          instrumented_output = script_ops.numpy_function(
              record, [output], output.dtype)
          instrumented_output.set_shape(output.shape)
        instrumented_outputs.append(instrumented_output)

      return instrumented_outputs

  def reset(self):
    self.eager_op_types = []
    self.eager_op_names = []
    self.eager_attrs = []
    self.eager_graphs = []
    self.eager_inputs = []
    self.graph_op_types = []
    self.graph_op_names = []
    self.graph_attrs = []
    self.graph_graphs = []
    self.graph_graph_versions = []
    self.graph_inputs = []

    # A dict mapping tensor name (e.g., "MatMut_10") to a list of ndarrays.
    # The list is the history of the tensor's computation result inside
    # `tf.Graph`s (`FuncGraph`s).
    # For an op with multiple output tensors, the outputs are interleaved in
    # the list.
    self.graph_internal_ndarrays = {}


class OpCallbacksTest(test_util.TensorFlowTestCase):

  def testSingleThreadedStack(self):
    instrument_0 = _NumpyFunctionCallback()
    instrument_1 = _NumpyFunctionCallback()

    with op_callbacks.op_callback(instrument_0.callback):
      self.assertEqual(1, len(op_callbacks._state.callback_stack))
      self.assertEqual(instrument_0.callback,
                       op_callbacks._state.callback_stack[0])

      with op_callbacks.op_callback(instrument_1.callback):
        self.assertEqual(2, len(op_callbacks._state.callback_stack))
        self.assertEqual(instrument_0.callback,
                         op_callbacks._state.callback_stack[0])
        self.assertEqual(instrument_1.callback,
                         op_callbacks._state.callback_stack[1])

      self.assertEqual(1, len(op_callbacks._state.callback_stack))
      self.assertEqual(instrument_0.callback,
                       op_callbacks._state.callback_stack[0])

    self.assertEqual(0, len(op_callbacks._state.callback_stack))

  def testMultiThreadedStacks(self):
    # Instrument for the main thread.
    instrument_0 = _NumpyFunctionCallback()

    # Instrument for the to-be-created thread.
    instrument_1 = _NumpyFunctionCallback()

    def thread1_job():
      with op_callbacks.op_callback(instrument_1.callback):

        @def_function.function
        def func1(x):
          return math_ops.sqrt(math_ops.log(x))

        x = constant_op.constant(4.0)
        self.assertAllClose(func1(x), np.sqrt(np.log(4.0)))

    thread1 = threading.Thread(target=thread1_job)

    # Start job on separate thread.
    thread1.start()

    # Run something on the main thread.
    with op_callbacks.op_callback(instrument_0.callback):

      @def_function.function
      def func0(x):
        return math_ops.square(math_ops.sin(x))

      x = constant_op.constant(4.0)
      self.assertAllClose(func0(x), np.square(np.sin(4.0)))

    thread1.join()

    # Assert that there is no cross-talk between the main thread
    # and the created thread.
    self.assertIn(_LOG_OP, instrument_1.graph_op_types)
    self.assertIn(_SQRT_OP, instrument_1.graph_op_types)
    self.assertNotIn(_SIN_OP, instrument_1.graph_op_types)
    self.assertNotIn(_SQUARE_OP, instrument_1.graph_op_types)

    self.assertNotIn(_LOG_OP, instrument_0.graph_op_types)
    self.assertNotIn(_SQRT_OP, instrument_0.graph_op_types)
    self.assertIn(_SIN_OP, instrument_0.graph_op_types)
    self.assertIn(_SQUARE_OP, instrument_0.graph_op_types)

  def testEagerOpExecution(self):
    instrument = _NumpyFunctionCallback()

    with op_callbacks.op_callback(instrument.callback):
      x = constant_op.constant(6.0)
      y = math_ops.square(math_ops.log(x))
      self.assertAllClose(y, np.square(np.log(6.0)))

    self.assertEqual(instrument.eager_op_types, [_LOG_OP, _SQUARE_OP])
    # Op names are unavailable under eager mode.
    self.assertEqual(instrument.eager_op_names, [None, None])
    self.assertEqual(instrument.eager_graphs, [None, None])
    self.assertEqual(len(instrument.eager_inputs), 2)
    self.assertEqual(len(instrument.eager_inputs[0]), 1)
    self.assertIsInstance(instrument.eager_inputs[0], tuple)
    self.assertEqual(instrument.eager_inputs[0][0], x)
    self.assertEqual(len(instrument.eager_inputs[1]), 1)
    self.assertIsInstance(instrument.eager_inputs[1], tuple)
    self.assertAllClose(instrument.eager_inputs[1][0], np.log(6.0))
    self.assertFalse(instrument.graph_op_types)
    self.assertFalse(instrument.graph_op_names)
    self.assertFalse(instrument.graph_attrs)
    self.assertFalse(instrument.graph_graphs)
    self.assertFalse(instrument.graph_inputs)

  def testMultiThreadedEagerOpExecution(self):
    # Instrument for the main thread.
    instrument_0 = _NumpyFunctionCallback()

    # Instrument for the to-be-created thread.
    instrument_1 = _NumpyFunctionCallback()

    def thread_1_job():
      with op_callbacks.op_callback(instrument_1.callback):
        x = constant_op.constant(6.0)
        y = math_ops.square(math_ops.log(x))
        return y

    thread_1 = threading.Thread(target=thread_1_job)
    thread_1.start()

    # While thread_1 is ongoing, do something on the main thread.
    with op_callbacks.op_callback(instrument_0.callback):
      x = constant_op.constant(2.0)
      y = math_ops.cos(x)
      self.assertAllClose(y, np.cos(2.0))

    thread_1.join()

    self.assertEqual(instrument_0.eager_op_types, [_COS_OP])
    self.assertEqual(instrument_0.eager_op_names, [None])
    self.assertEqual(instrument_1.eager_op_types, [_LOG_OP, _SQUARE_OP])
    self.assertEqual(instrument_1.eager_op_names, [None, None])

  def testEagerFunctionExecution(self):
    instrument = _NumpyFunctionCallback()

    @def_function.function
    def square_log(x):
      return math_ops.square(math_ops.log(x))

    # Call the function once, so that the graph construction won't show up
    # in the callback.
    x_float32 = constant_op.constant(6.0, dtype=dtypes.float32)
    x_float64 = constant_op.constant(6.0, dtype=dtypes.float64)
    square_log(x_float32)
    square_log(x_float64)

    with op_callbacks.op_callback(instrument.callback):
      y = square_log(x_float32)
      self.assertAllClose(y, np.square(np.log(6.0)))
      y = square_log(x_float64)
      self.assertAllClose(y, np.square(np.log(6.0)))

    # Each of the two dtypes should be associated with its own FuncGraph.
    self.assertIn(
        square_log.get_concrete_function(x_float32).name,
        instrument.eager_op_types)
    self.assertIn(
        square_log.get_concrete_function(x_float64).name,
        instrument.eager_op_types)

    self.assertEqual(len(instrument.eager_inputs), 2)
    self.assertIsInstance(instrument.eager_inputs[0], tuple)
    self.assertEqual(instrument.eager_inputs[0][0], x_float32)
    self.assertIsInstance(instrument.eager_inputs[1], tuple)
    self.assertEqual(instrument.eager_inputs[1][0], x_float64)

    self.assertEqual(instrument.eager_op_names, [None, None])
    self.assertFalse(instrument.graph_op_types)
    self.assertFalse(instrument.graph_op_names)
    self.assertFalse(instrument.graph_inputs)

  def testMultiThreadedEagerFunctionExecution(self):
    # Instrument for the main thread.
    instrument_0 = _NumpyFunctionCallback()

    # Instrument for the to-be-created thread.
    instrument_1 = _NumpyFunctionCallback()

    @def_function.function
    def square_log(x):
      return math_ops.square(math_ops.log(x))

    # Call the function once, so that the graph construction won't show up
    # in the callback.
    x_float32 = constant_op.constant(6.0, dtype=dtypes.float32)
    x_float64 = constant_op.constant(6.0, dtype=dtypes.float64)
    square_log(x_float32)
    square_log(x_float64)

    def thread_1_job():
      with op_callbacks.op_callback(instrument_1.callback):
        square_log(x_float32)

    thread_1 = threading.Thread(target=thread_1_job)
    thread_1.start()

    # In the meantime, run some computation on the main thread.
    with op_callbacks.op_callback(instrument_0.callback):
      square_log(x_float64)

    thread_1.join()

    # Each of the two dtypes should be associated with its own FuncGraph.
    self.assertIn(
        square_log.get_concrete_function(x_float64).name,
        instrument_0.eager_op_types)
    self.assertEqual(instrument_0.eager_op_names, [None])
    self.assertFalse(instrument_0.graph_op_types)
    self.assertIn(
        square_log.get_concrete_function(x_float32).name,
        instrument_1.eager_op_types)
    self.assertEqual(instrument_1.eager_op_names, [None])
    self.assertFalse(instrument_1.graph_op_types)

  def testSimpleGraphConstructionScopeOutsideFunction(self):
    instrument = _NumpyFunctionCallback()

    with op_callbacks.op_callback(instrument.callback):

      @def_function.function
      def log_2plus_unique_x(x):
        unique_values, unique_pos = array_ops.unique(x)
        return math_ops.log(2.0 + unique_values), unique_pos

      x = constant_op.constant([-1.0, -1.0, 0.0], dtype=dtypes.float32)
      y1, y2 = log_2plus_unique_x(x)
      self.assertAllClose(y1, [0.0, np.log(2.0)])
      self.assertAllClose(y2, [0, 0, 1])
    self.assertIn(_UNIQUE_OP, instrument.graph_op_types)
    self.assertIn(_ADD_OP, instrument.graph_op_types)
    self.assertIn(_LOG_OP, instrument.graph_op_types)
    self.assertEqual(
        len(instrument.graph_op_names), len(instrument.graph_op_types))

    # Check the graph internal ndarrays recorded at runtime.
    unique_op_outputs = instrument.graph_internal_ndarrays[_UNIQUE_OP]
    self.assertEqual(len(unique_op_outputs), 2)
    self.assertAllClose(unique_op_outputs[0], [-1.0, 0.0])
    self.assertAllClose(unique_op_outputs[1], [0, 0, 1])
    add_op_outputs = instrument.graph_internal_ndarrays[b"add"]
    self.assertEqual(len(add_op_outputs), 1)
    self.assertAllClose(add_op_outputs[0], [1.0, 2.0])
    log_op_outputs = instrument.graph_internal_ndarrays[_LOG_OP]
    self.assertEqual(len(log_op_outputs), 1)
    self.assertAllClose(log_op_outputs[0], [0.0, np.log(2.0)])

  def testSimpleGraphConstructionWithCallbackReturningNone(self):
    """Test that callbacks that return None works."""
    op_types = []
    def no_return_callback(op_type,
                           inputs,
                           attrs,
                           outputs,
                           op_name=None,
                           graph=None):
      del inputs, attrs, outputs, op_name, graph  # Unused.
      op_types.append(compat.as_bytes(op_type))

    with op_callbacks.op_callback(no_return_callback):
      @def_function.function
      def log1p(x):
        return math_ops.log(1.0 + x)
      x = constant_op.constant(3.0)
      y = log1p(x)
      self.assertAllClose(y, np.log(4.0))
    self.assertIn(_ADD_OP, op_types)
    self.assertIn(_LOG_OP, op_types)

  def testGraphConstructionInputsAndGraphAreCapturedCorrectly(self):
    instrument = _NumpyFunctionCallback(instrument_graph_ops=False)

    with op_callbacks.op_callback(instrument.callback):

      @def_function.function
      def log_2plus_unique_x(x):
        unique_values, unique_pos = array_ops.unique(x)
        return math_ops.log(2.0 + unique_values), unique_pos

      x = constant_op.constant([-1.0, -1.0, 0.0], dtype=dtypes.float32)
      y1, y2 = log_2plus_unique_x(x)
      self.assertAllClose(y1, [0.0, np.log(2.0)])
      self.assertAllClose(y2, [0, 0, 1])

    # Check the recorded input tensors.
    self.assertEqual(
        len(instrument.graph_inputs), len(instrument.graph_op_types))
    unique_inputs = instrument.graph_inputs[instrument.graph_op_types.index(
        _UNIQUE_OP)]
    self.assertIsInstance(unique_inputs, tuple)
    self.assertEqual(len(unique_inputs), 1)
    self.assertEqual(
        compat.as_bytes(unique_inputs[0].op.op_def.name), _PLACEHOLDER_OP)

    add_inputs = instrument.graph_inputs[instrument.graph_op_types.index(
        _ADD_OP)]
    self.assertIsInstance(add_inputs, tuple)
    self.assertEqual(len(add_inputs), 2)
    self.assertEqual(
        compat.as_bytes(add_inputs[0].op.op_def.name), _CONSTANT_OP)
    self.assertEqual(compat.as_bytes(add_inputs[1].op.op_def.name), _UNIQUE_OP)

    log_inputs = instrument.graph_inputs[instrument.graph_op_types.index(
        _LOG_OP)]
    self.assertIsInstance(log_inputs, tuple)
    self.assertEqual(len(log_inputs), 1)
    self.assertEqual(compat.as_bytes(log_inputs[0].op.op_def.name), _ADD_OP)

    # Check the recorded graphs.
    self.assertEqual(
        len(instrument.graph_graphs), len(instrument.graph_op_types))
    self.assertGreater(len(instrument.graph_graph_versions), 1)
    for i in range(len(instrument.graph_graph_versions) - 1):
      self.assertGreater(instrument.graph_graph_versions[i + 1],
                         instrument.graph_graph_versions[i])

  def testEagerGraphOpConstructionSimpleGraphScopeInsideFunction(self):
    instrument = _NumpyFunctionCallback()

    @def_function.function
    def log_2plus_unique_x(x):
      with op_callbacks.op_callback(instrument.callback):
        unique_values, _ = array_ops.unique(x)
        y = math_ops.log(2.0 + unique_values)
      return math_ops.sin(y)

    x = constant_op.constant([-1.0, -1.0, 0.0], dtype=dtypes.float32)
    output = log_2plus_unique_x(x)
    self.assertAllClose(output, np.sin([0.0, np.log(2.0)]))

    # The following ops should have been captured by the callback
    # because they were constructed within the scope of `op_callback()`.
    self.assertIn(_UNIQUE_OP, instrument.graph_op_types)
    self.assertIn(_ADD_OP, instrument.graph_op_types)
    self.assertIn(_LOG_OP, instrument.graph_op_types)
    # The "Sin" op should not have been captured, because it was constructed
    # outside the scope of `op_callback()`.
    self.assertNotIn(_SIN_OP, instrument.graph_op_types)
    self.assertEqual(
        len(instrument.graph_op_names), len(instrument.graph_op_types))

    # Check the graph internal ndarrays recorded at runtime.
    unique_op_outputs = instrument.graph_internal_ndarrays[_UNIQUE_OP]
    self.assertEqual(len(unique_op_outputs), 2)
    self.assertAllClose(unique_op_outputs[0], [-1.0, 0.0])
    self.assertAllClose(unique_op_outputs[1], [0, 0, 1])
    add_op_outputs = instrument.graph_internal_ndarrays[b"add"]
    self.assertEqual(len(add_op_outputs), 1)
    self.assertAllClose(add_op_outputs[0], [1.0, 2.0])
    log_op_outputs = instrument.graph_internal_ndarrays[_LOG_OP]
    self.assertEqual(len(log_op_outputs), 1)
    self.assertAllClose(log_op_outputs[0], [0.0, np.log(2.0)])

  def testEagerOpAttributesAreCapture(self):
    instrument = _NumpyFunctionCallback()
    with op_callbacks.op_callback(instrument.callback):
      m = constant_op.constant([[1.0, -1.0], [0.0, 1.0]])
      x = constant_op.constant([[-2.0], [3.0]])
      y = math_ops.matmul(m, x, transpose_a=True, transpose_b=False)
      self.assertAllClose(y, [[-2.0], [5.0]])
    self.assertEqual(len(instrument.eager_attrs), 1)
    self.assertIsInstance(instrument.eager_attrs[0], tuple)
    self.assertEqual(
        instrument.eager_attrs[0][instrument.eager_attrs[0].index("transpose_a")
                                  + 1], True)
    self.assertEqual(
        instrument.eager_attrs[0][instrument.eager_attrs[0].index("transpose_b")
                                  + 1], False)
    self.assertEqual(len(instrument.graph_attrs), 0)

  def testGraphOpAttributesAreCapture(self):
    instrument = _NumpyFunctionCallback()
    with op_callbacks.op_callback(instrument.callback):

      @def_function.function
      def my_matmul(m, x):
        return math_ops.matmul(m, x, transpose_a=True, transpose_b=False)

      m = constant_op.constant([[1.0, -1.0], [0.0, 1.0]])
      x = constant_op.constant([[-2.0], [3.0]])
      y = my_matmul(m, x)
      self.assertAllClose(y, [[-2.0], [5.0]])
    index = instrument.graph_op_types.index(_MATMUL_OP)
    self.assertIsInstance(instrument.graph_attrs[index], tuple)
    self.assertEqual(
        instrument.graph_attrs[index][
            instrument.graph_attrs[index].index("transpose_a") + 1].b, True)
    self.assertEqual(
        instrument.graph_attrs[index][
            instrument.graph_attrs[index].index("transpose_b") + 1].b, False)
    self.assertEqual(len(instrument.eager_attrs), 1)
    self.assertIsInstance(instrument.eager_attrs[0], tuple)

  def testEagerGraphOpConstructionIfControlFlow(self):
    instrument = _NumpyFunctionCallback()

    with op_callbacks.op_callback(instrument.callback):

      @def_function.function
      def my_function_with_cond(x):
        if math_ops.greater(x, 0.0):
          return x**2.0
        else:
          return x**3.0

      x = constant_op.constant(-4.0)
      self.assertAllClose(my_function_with_cond(x), -64.0)

    self.assertIn(_IF_OP, instrument.graph_op_types)
    self.assertIn(_GREATER_OP, instrument.graph_op_types)
    self.assertIn(_POW_OP, instrument.graph_op_types)
    self.assertEqual(
        len(instrument.graph_op_names), len(instrument.graph_op_types))

    # Check the graph internal ndarrays recorded at runtime.
    greater_op_outputs = instrument.graph_internal_ndarrays[_GREATER_OP]
    self.assertEqual(len(greater_op_outputs), 1)
    self.assertAllClose(greater_op_outputs[0], False)
    pow_op_outputs = instrument.graph_internal_ndarrays[b"pow"]
    self.assertEqual(len(pow_op_outputs), 1)
    self.assertAllClose(pow_op_outputs[0], -64.0)

  def testEagerGraphOpConstructionWhileLoopControlFlow(self):
    instrument = _NumpyFunctionCallback()

    with op_callbacks.op_callback(instrument.callback):

      @def_function.function
      def my_function_with_while(counter, lim, accum):
        while math_ops.less(counter, lim):
          accum.assign_add(accum)
          counter.assign_add(1.0)

      counter = variables.Variable(0.0)
      lim = constant_op.constant(4.0, dtype=dtypes.float32)
      accum = variables.Variable(1.0)
      my_function_with_while(counter, lim, accum)
    self.assertAllClose(accum.read_value(), 16.0)
    self.assertIn(_WHILE_OP, instrument.graph_op_types)
    self.assertIn(_LESS_OP, instrument.graph_op_types)
    self.assertIn(_ASSIGN_ADD_VARIABLE_OP, instrument.graph_op_types)
    self.assertEqual(
        len(instrument.graph_op_names), len(instrument.graph_op_types))

    # Check the graph internal ndarrays recorded at runtime.
    read_variable_op_outputs = instrument.graph_internal_ndarrays[
        _READ_VARIALBE_OP]
    self.assertAllClose(read_variable_op_outputs, [1.0, 2.0, 4.0, 8.0])
    less_op_outputs = instrument.graph_internal_ndarrays[_LESS_OP]
    self.assertAllClose(less_op_outputs, [True, True, True, True, False])

  def testDatasetMapTest(self):
    instrument = _NumpyFunctionCallback()

    with op_callbacks.op_callback(instrument.callback):
      tensor = constant_op.constant(
          [0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0])

      def map_fn(x):
        return math_ops.log(math_ops.square(x) + 1)

      dataset = dataset_ops.Dataset.from_tensor_slices(tensor).batch(2).map(
          map_fn)
      iterator = dataset.make_one_shot_iterator()

      self.assertAllClose(iterator.next(), np.log([1.25, 2]))
      self.assertAllClose(iterator.next(), np.log([3.25, 5]))

      self.assertIn(_SQUARE_OP, instrument.graph_op_types)
      self.assertIn(_ADD_OP, instrument.graph_op_types)
      self.assertIn(_LOG_OP, instrument.graph_op_types)
      self.assertEqual(
          len(instrument.eager_op_types), len(instrument.eager_op_names))

  def testSparseTensorEagerExecution(self):
    instrument = _NumpyFunctionCallback()

    with op_callbacks.op_callback(instrument.callback):
      indices = [[1, 2], [2, 0], [3, 4]]
      values = [0.0, 8.0, -2.0]
      shape = [4, 5]
      sp = sparse_tensor.SparseTensorValue(indices, values, shape)
      w = ops.convert_to_tensor(np.ones([5, 1], np.float32))

      y = sparse_ops.sparse_tensor_dense_matmul(sp, w)
      self.assertAllClose(y, [[0.0], [0.0], [8.0], [-2.0]])
      self.assertIn(_SPARSE_TENSOR_DENSE_MATMUL_OP, instrument.eager_op_types)
      self.assertFalse(instrument.graph_op_types)

  def testSparseTensorFuncGraph(self):
    instrument = _NumpyFunctionCallback()

    with op_callbacks.op_callback(instrument.callback):

      @def_function.function
      def dense_matmul(sp, w):
        return sparse_ops.sparse_tensor_dense_matmul(sp, w)

      indices = [[1, 2], [2, 0], [3, 4]]
      values = [0.0, 8.0, -2.0]
      shape = [4, 5]
      sp = sparse_tensor.SparseTensorValue(indices, values, shape)
      w = ops.convert_to_tensor(np.ones([5, 1], np.float32))
      y = dense_matmul(sp, w)
      self.assertAllClose(y, [[0.0], [0.0], [8.0], [-2.0]])
      self.assertIn(_SPARSE_TENSOR_DENSE_MATMUL_OP, instrument.graph_op_types)
      self.assertIn(
          dense_matmul.get_concrete_function(sp, w).name,
          instrument.eager_op_types)

    # Check the graph internal ndarrays recorded at runtime.
    sparse_matmul_outputs = instrument.graph_internal_ndarrays[
        _SPARSE_TENSOR_DENSE_MATMUL_OP + b"/" + _SPARSE_TENSOR_DENSE_MATMUL_OP]
    self.assertEqual(len(sparse_matmul_outputs), 1)
    self.assertAllClose(sparse_matmul_outputs[0], [[0.0], [0.0], [8.0], [-2.0]])

  def testOverrideDTypeInFuncGraph(self):

    def to_float64(op_type, inputs, attrs, outputs, op_name=None, graph=None):
      del op_type, inputs, attrs, op_name, graph  # Unused.
      return [math_ops.cast(output, dtypes.float64) for output in outputs]

    with op_callbacks.op_callback(to_float64):

      @def_function.function
      def add_1_times_2(x):
        return (x + 1.0) * 2.0

      x = constant_op.constant(3.0, dtype=dtypes.float32)
      y = add_1_times_2(x)
      self.assertEqual(y.dtype, dtypes.float64)
      self.assertAllClose(y, 8.0)

  def testNoOutputOpUnderEagerExecution(self):
    instrument = _NumpyFunctionCallback()
    with op_callbacks.op_callback(instrument.callback):
      x = constant_op.constant(10.0)
      y = constant_op.constant(20.0)
      z = x + y
      w = control_flow_ops.group([z])
      self.assertIsNone(w)
    self.assertEqual(instrument.eager_op_types, [_ADD_OP])

  def testOpCallbackWorksWithGradientTape(self):
    instrument = _NumpyFunctionCallback()

    with op_callbacks.op_callback(instrument.callback):
      v = variables.Variable(3.0, dtype=dtypes.float32)
      @def_function.function
      def get_gradients():
        with backprop.GradientTape() as tape:
          loss = math_ops.sin(math_ops.square(v))
          gradients = tape.gradient(loss, v)
        return gradients

      gradients = get_gradients()
      # Applying the chain rule.
      self.assertAllClose(gradients, np.cos(3.0 * 3.0) * 3.0 * 2.0)
      self.assertIn(_SQUARE_OP, instrument.graph_op_types)
      self.assertIn(_SIN_OP, instrument.graph_op_types)
      # The mul and cos ops are created for backprop.
      self.assertIn(_MUL_OP, instrument.graph_op_types)
      self.assertIn(_COS_OP, instrument.graph_op_types)

      # Check the ndarrays from runtime.
      cos_op_outputs = instrument.graph_internal_ndarrays[_COS_OP]
      self.assertEqual(len(cos_op_outputs), 1)
      self.assertAllClose(cos_op_outputs[0], np.cos(3.0 * 3.0))

  def testKeraModelFit(self):
    # TODO(cais): The purely PyFunc (numpy_function) based instrumentation
    # doesn't work for the entire Keras model and its fit() call, due to some
    # shape inference limitations. Use tfdbg's gen_debug_ops for testing
    # instead (b/139668469).
    instrument = _NumpyFunctionCallback(instrument_graph_ops=False)

    with op_callbacks.op_callback(instrument.callback):
      model = keras.Sequential()
      model.add(keras.layers.Dense(10, input_shape=(8,), activation="relu"))
      model.add(keras.layers.BatchNormalization())
      model.add(keras.layers.Dense(1, activation="linear"))
      model.compile(loss="mse", optimizer="adam")

      batch_size = 4
      xs = random_ops.random_normal([batch_size, 8])
      ys = random_ops.random_normal([batch_size, 1])
      history = model.fit(xs, ys, epochs=2, verbose=0)

      # Simply assert that the training proceeded as expected and that
      # op callbacks are invoked. We prefer not to assert on the details of the
      # graph construction and the execution, in order to avoid future
      # maintenance cost.
      self.assertEqual(len(history.history["loss"]), 2)
      self.assertTrue(instrument.graph_op_types)
      self.assertEqual(len(instrument.graph_op_types),
                       len(instrument.graph_op_names))
      self.assertTrue(instrument.eager_op_types)


class OpCallbacksErrorConditionsTest(test_util.TensorFlowTestCase):

  def testNonCallableObjectArgErrors(self):
    with self.assertRaisesRegex(ValueError, r"is expected to be callable"):
      with op_callbacks.op_callback(1337):
        pass

  def testOverridingWithWrongNumberOfTensorOutputsErrors(self):
    def wrong_outputs_callback(op_type,
                               inputs,
                               attrs,
                               outputs,
                               op_name=None,
                               graph=None):
      del op_type, inputs, attrs, op_name, graph  # Unused.
      return outputs[0], math_ops.negative(outputs[0])

    with op_callbacks.op_callback(wrong_outputs_callback):

      @def_function.function
      def log1p(x):
        return math_ops.log(1.0 + x)

      x = constant_op.constant(3.0)
      with self.assertRaisesRegex(
          ValueError,
          r"returned 2 tensors, .* does not match .* \(1\)"):
        log1p(x)


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