import functools
import sys

from absl.testing import parameterized
import numpy as np

from tensorflow.python import pywrap_tfe
from tensorflow.python.eager import backprop
from tensorflow.python.eager import backprop_util
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import record
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
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_util
from tensorflow.python.framework import test_util
from tensorflow.python.framework.memory_checker import MemoryChecker
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import cond as tf_cond
from tensorflow.python.ops import custom_gradient
from tensorflow.python.ops import embedding_ops
from tensorflow.python.ops import functional_ops
from tensorflow.python.ops import gradient_checker_v2
from tensorflow.python.ops import gradients
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
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.ops import resource_variable_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops import while_loop
from tensorflow.python.training import training

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

  @test_util.run_in_graph_and_eager_modes
  def testAggregateGradients(self):

    def fn(x):
      ind1 = constant_op.constant(np.array([0, 1]))
      ind2 = constant_op.constant(np.array([2, 3]))
      ind3 = constant_op.constant(np.array([1, 3]))
      g1 = embedding_ops.embedding_lookup(x, ind1)
      g2 = embedding_ops.embedding_lookup(x, ind2)
      g3 = embedding_ops.embedding_lookup(x, ind3)
      return g1 * g2 * g3

    var_np = np.random.rand(4, 2).astype(np.float32)
    var = constant_op.constant(var_np)
    grad = backprop.gradients_function(fn, [0])(var)[0]
    grad = self.evaluate(ops.convert_to_tensor(grad))

    if not context.executing_eagerly():
      tf_var = array_ops.constant(var_np, dtypes.float32)
      tf_ind1 = array_ops.constant([0, 1])
      tf_ind2 = array_ops.constant([2, 3])
      tf_ind3 = array_ops.constant([1, 3])
      tf_g1 = embedding_ops.embedding_lookup(tf_var, tf_ind1)
      tf_g2 = embedding_ops.embedding_lookup(tf_var, tf_ind2)
      tf_g3 = embedding_ops.embedding_lookup(tf_var, tf_ind3)
      tf_y = tf_g1 * tf_g2 * tf_g3
      tf_grad = gradients.gradients(tf_y, [tf_var])[0]

      tf_dense_grad = math_ops.unsorted_segment_sum(tf_grad.values,
                                                    tf_grad.indices,
                                                    tf_grad.dense_shape[0])

      self.assertAllClose(grad, self.evaluate(tf_dense_grad))

  @test_util.run_in_graph_and_eager_modes
  def testAggregateGradientsWithTensor(self):

    def fn(x):
      ind1 = constant_op.constant(np.array([0, 1]))
      # A mixture of IndexedSlices and dense tensor to aggregate.
      g1 = embedding_ops.embedding_lookup(x, ind1)
      g2 = math_ops.reduce_sum(x * constant_op.constant(2.0))
      return g1 * g2

    var_np = np.random.rand(4, 2).astype(np.float32)
    var = constant_op.constant(var_np)
    grad = backprop.gradients_function(fn, [0])(var)[0]
    grad = self.evaluate(ops.convert_to_tensor(grad))

    if not context.executing_eagerly():
      tf_var = array_ops.constant(var_np, dtypes.float32)
      tf_ind1 = array_ops.constant([0, 1])
      tf_g1 = embedding_ops.embedding_lookup(tf_var, tf_ind1)
      tf_g2 = math_ops.reduce_sum(tf_var * 2.0, axis=(0, 1))
      tf_y = tf_g1 * tf_g2
      tf_grad = gradients.gradients(tf_y, [tf_var])[0]

      self.assertAllClose(grad, tf_grad)

  def testImplicitGradWithResourceVariable(self):
    x = resource_variable_ops.ResourceVariable(
        initial_value=constant_op.constant(1.0), name='x')

    def fn():
      b = constant_op.constant(2.0)
      c = math_ops.add(x.value(), b)
      return math_ops.add(c, constant_op.constant(3.0))

    grads_and_vars = backprop.implicit_grad(fn)()
    self.assertAllEqual(grads_and_vars[0][0], 1.0)
    self.assertAllEqual(id(grads_and_vars[0][1]), id(x))

  @parameterized.named_parameters([('Function', def_function.function),
                                   ('NoFunction', lambda f: f)])
  def testNoOpBehaviorConsistent(self, decorator):

    @decorator
    def f(x):
      # Test all different types of no-ops
      x1 = array_ops.identity(x)
      x2 = math_ops.add_v2(x, 0)
      x3 = math_ops.subtract(x, 0)
      x4 = math_ops.multiply(x, 1)
      with backprop.GradientTape() as t:
        t.watch(x)
        t.watch(x1)
        t.watch(x2)
        t.watch(x3)
        t.watch(x4)
        y1 = x * 2.
        y2 = x1 * 3.
        y3 = x2 * 3.
        y4 = x3 * 3.
        y5 = x4 * 3.
        loss = y1 + y2 + y3 + y4 + y5
      return t.gradient(loss, [x, x1, x2, x3, x4])

    self.assertAllClose([2., 3., 3., 3., 3.], f(constant_op.constant(10.)))

  def testResourceHandleOutputWithoutHandleData(self):
    # This is a bit of a weird thing to test since we try to maintain handle
    # data. But users do create their own resources, and those often do not have
    # any handle data.
    h = resource_variable_ops.var_handle_op(
        shape=[], dtype=dtypes.float32, shared_name='abc')

    with backprop.GradientTape() as tape:
      x = constant_op.constant(1.)
      tape.watch(x)
      tape.watch(h)
      y, h = array_ops.identity_n([x, h])

    self.assertAllClose(1., tape.gradient(y, x))

  def testGradientInsideLoop(self):
    with ops.Graph().as_default():
      v = resource_variable_ops.ResourceVariable(1.0)

      def body(_):
        _ = v + 1.0  # This reads the variable inside the loop context
        with backprop.GradientTape() as t:
          result = v * 2
        self.assertIsNotNone(t.gradient(result, v))
        return 1.0

      while_loop.while_loop(lambda i: False, body, [1.0])

  def testWhereGradient(self):
    # Note: where is special because only some of its arguments are of
    # differentiable dtypes.

    def f(x):
      return array_ops.where(x < 10, x, x * x)

    g = backprop.gradients_function(f)

    self.assertAllEqual(g(5.)[0], 1.0)
    self.assertAllEqual(g(50.)[0], 100.0)

  def testTwoTargets(self):
    with backprop.GradientTape() as t:
      x = constant_op.constant(3.0)
      y = constant_op.constant(2.0)
      t.watch([x, y])
      xx = 2 * x
      yy = 3 * y
    dx, dy = t.gradient([xx, yy], [x, y])
    self.assertAllEqual(dx, 2.0)
    self.assertAllEqual(dy, 3.0)

  def testCustomGradientEmptyError(self):

    @custom_gradient.custom_gradient
    def identity(x):

      def grad(_):
        return []  # This return value is wrong!

      return x, grad

    x = variables.Variable(1.0)
    with backprop.GradientTape() as t:
      y = identity(x)
    with self.assertRaises(ValueError):
      t.gradient(y, [x])

  def test_stop_gradient_hides_downstream_ops(self):

    @custom_gradient.custom_gradient
    def _backward_pass_error(x):

      def _grad(_):
        raise AssertionError(
            'Unexpectedly ran the backward function. This probably means that '
            'tf.GradientTape is not properly ignoring tensors downstream of '
            'tf.stop_gradient.')

      return x, _grad

    @def_function.function
    def f(x):
      return _backward_pass_error(x)

    x = constant_op.constant(1.)
    with backprop.GradientTape() as tape:
      tape.watch(x)
      y = f(array_ops.stop_gradient(x))

    self.assertIsNone(tape.gradient(y, x))

  def testOutputGradUsedInComputation(self):
    with backprop.GradientTape() as t:
      x = constant_op.constant(3.0)
      y = constant_op.constant(2.0)
      t.watch([x, y])
      loss = x * y
    dx, = t.gradient([loss, x], [x], output_gradients=[1.0, 2.0])
    self.assertAllEqual(dx, 4.0)

  def testDy(self):

    def f(x):
      return x

    grad_fn = backprop.gradients_function(f)
    self.assertAllEqual(2., grad_fn(1., dy=2.)[0])

  def testGradientInteger(self):

    def f(x):
      return x + x

    int_tensor = constant_op.constant(1)
    self.assertEqual(backprop.gradients_function(f)(int_tensor)[0], None)

  def testErrors(self):

    @custom_gradient.custom_gradient
    def f(x):

      def grad(_):
        raise RuntimeError('x')

      return x, grad

    # TODO(apassos) raise the right error here
    with self.assertRaises(RuntimeError):
      backprop.gradients_function(f)(constant_op.constant(1.0))

  def testGradientsFunctionInCustomGradient(self):

    @custom_gradient.custom_gradient
    def f(x):
      (y,) = backprop.gradients_function(lambda x: x * x)(x)

      def grad(dy):
        return [2 * dy]

      return y, grad

    self.assertAllEqual(f(1.0), 2.0)

  def testImplicitGradOverEmbeddingLookup(self):
    batch_size = 8
    embedding_size = 512
    vocab_size = 1000
    lrn_rate = 0.1
    random_init = random_ops.random_uniform([vocab_size, embedding_size])

    x = array_ops.ones((batch_size), dtypes.int64)
    embedding = resource_variable_ops.ResourceVariable(
        initial_value=random_init, dtype=dtypes.float32, name='embedding')

    def f():
      embedded_x = embedding_ops.embedding_lookup(embedding, x)
      return constant_op.constant(1.0, dtypes.float32) - embedded_x

    grad = backprop.implicit_grad(f)()[0][0]
    opt = training.GradientDescentOptimizer(lrn_rate)

    with ops.Graph().as_default(), self.cached_session():
      tf_x = array_ops.ones((batch_size), dtypes.int64)
      # TODO(ashankar,apassos): Change to ResourceVariable.
      tf_embedding = variables.Variable(
          random_init.numpy(), name='tf_embedding')
      tf_embedded_x = embedding_ops.embedding_lookup(tf_embedding, tf_x)
      tf_y = 1.0 - tf_embedded_x
      tf_grad = gradients.gradients(tf_y, [tf_embedding])[0]
      tf_opt = training.GradientDescentOptimizer(0.1)
      tf_embedding.initializer.run()

      self.assertAllClose(tf_grad.indices, grad.indices)
      self.assertAllClose(tf_grad.values, grad.values)

      tf_opt.apply_gradients([(tf_grad, tf_embedding)]).run()
      expected = self.evaluate(tf_embedding)
    opt.apply_gradients([(grad, embedding)])
    self.assertAllClose(expected, embedding.read_value())

  def testImplicitGradOrdering(self):
    v0 = resource_variable_ops.ResourceVariable(1.0)
    v1 = resource_variable_ops.ResourceVariable(2.0)

    def f():
      x = v1 * v1
      y = v0 * v0
      return x + y

    grads = backprop.implicit_grad(f)()
    ordered_variables = [x[1] for x in grads]
    self.assertIs(ordered_variables[0], v0)
    self.assertIs(ordered_variables[1], v1)

  def testTapeNoOpGradient(self):
    x = constant_op.constant(3.0)
    with backprop.GradientTape() as t:
      t.watch(x)
      y = x
    self.assertEqual(t.gradient(y, x).numpy(), 1.0)

  def testTapeIdentityGradientIsIdentity(self):
    x = constant_op.constant(3.0)
    with backprop.GradientTape() as t:
      t.watch(x)
      y = array_ops.identity(x)
    self.assertEqual(t.gradient(y, x).numpy(), 1.0)

  def testFunctionIndexedSlicesGradient(self):

    @def_function.function
    def f(x):
      return x + 1

    with backprop.GradientTape() as t:
      x = constant_op.constant([1.0])
      t.watch(x)
      y = f(x)
      y = array_ops.gather(y, [0])
    self.assertAllEqual(t.gradient(y, x), [1.0])

  def testTapeGradientMultiTargetOneIsSource(self):
    x = constant_op.constant(2.0)
    with backprop.GradientTape() as t:
      t.watch(x)
      y = x * x
    self.assertEqual(t.gradient([x, y], x).numpy(), 5.0)

  def testTapeNoOpGradientWithMultiTargetAllSource(self):
    x = constant_op.constant(3.0)
    with backprop.GradientTape() as t:
      t.watch(x)
      y = x
    self.assertEqual(t.gradient([y, y], x).numpy(), 2.0)

  def testTapeNoOpGradientWithMultiTargetMultiSource(self):
    x = constant_op.constant(3.0)
    y = constant_op.constant(5.0)
    with backprop.GradientTape() as t:
      t.watch(x)
      t.watch(y)
      z = y * y
    self.assertAllEqual(t.gradient([x, y, z], [x, y]), [1.0, 11.0])

  def testTapeGradientStringTarget(self):
    s = constant_op.constant('unknown', dtype=dtypes.string)
    x = constant_op.constant(3.0)

    with backprop.GradientTape() as t:
      t.watch(x)
      t.watch(s)
    grads = t.gradient(s, x)
    self.assertEqual(grads, None)

  def testTapeNoOpGradientStringSourceAndTarget(self):
    s = constant_op.constant('unknown', dtype=dtypes.string)

    with backprop.GradientTape() as t:
      t.watch(s)
    grads = t.gradient(s, s)
    self.assertEqual(grads, None)

  def testTapeNoOpGradientWithMultiTargetMultiSourceIncludeString(self):
    x = constant_op.constant(3.0)
    y = constant_op.constant(5.0)
    s = constant_op.constant('unknown', dtype=dtypes.string)

    with backprop.GradientTape() as t:
      t.watch(x)
      t.watch(y)
      t.watch(s)
      z = y * y
    grads = t.gradient([x, y, z, s], [x, y, s])
    self.assertAllEqual(grads[:2], [1.0, 11.0])
    self.assertEqual(grads[2], None)

  def testTapeNoOpOnVariableIsIdentity(self):
    v0 = resource_variable_ops.ResourceVariable(1.0)
    with backprop.GradientTape() as t:
      y = v0.read_value()
    self.assertEqual(t.gradient(y, v0).numpy(), 1.0)

  @test_util.assert_no_new_tensors
  @test_util.assert_no_garbage_created
  def testTapeNoOpGradient2By2(self):
    a_2_by_2 = constant_op.constant(2.0, shape=[2, 2])
    with backprop.GradientTape(persistent=True) as tape:
      tape.watch(a_2_by_2)
    dy_dy = tape.gradient(a_2_by_2, [a_2_by_2])[0]
    self.assertAllEqual(dy_dy.numpy(),
                        constant_op.constant(1.0, shape=[2, 2]).numpy())

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testTapeNoOpGradientMultiTarget2By2(self):
    a_2_by_2 = constant_op.constant(2.0, shape=[2, 2])
    with backprop.GradientTape(persistent=True) as tape:
      tape.watch(a_2_by_2)
    dy_dy = tape.gradient([a_2_by_2, a_2_by_2], [a_2_by_2])[0]
    self.assertAllEqual(dy_dy.numpy(),
                        constant_op.constant(2.0, shape=[2, 2]).numpy())

  def testTapeStopRecording(self):
    with backprop.GradientTape() as t:
      x = resource_variable_ops.ResourceVariable(1.0)
      with t.stop_recording():
        y = x * x
    self.assertEqual(t.gradient(y, x), None)

  def testTapeStopStartRecording(self):
    with backprop.GradientTape(persistent=True) as t:
      x = resource_variable_ops.ResourceVariable(1.0)
      x2 = x * 2  # This should be differentiated through.
      with t.stop_recording():
        y = x2 * x2
      z = x2 * x2
    self.assertEqual(t.gradient(y, x2), None)

    # If the x*2 was not differentiated through, this would be 2.0, not 4.0
    self.assertEqual(t.gradient(z, x2).numpy(), 4.0)

  def testTapeReset(self):
    with backprop.GradientTape() as t:
      v = resource_variable_ops.ResourceVariable(1.0)
      loss = v * v
      t.reset()
      loss += v * v
    self.assertAllEqual(t.gradient(loss, v), 2.0)

  def testPythonMax(self):
    x = [
        resource_variable_ops.ResourceVariable(2.),
        resource_variable_ops.ResourceVariable(3.),
        resource_variable_ops.ResourceVariable(5.)
    ]
    with backprop.GradientTape() as t:
      f = max(x)
    grad = t.gradient(f, x)
    self.assertAllEqual(self.evaluate(f), 5.)
    self.assertAllEqual(self.evaluate(grad), [None, None, 1.0])

  def testAutomaticWatchedVariables(self):
    with backprop.GradientTape() as t:
      self.assertEqual(0, len(t.watched_variables()))
      v = resource_variable_ops.ResourceVariable(1.0)
      loss = v * v
      self.assertAllEqual([v], t.watched_variables())

      t.reset()
      self.assertEqual(0, len(t.watched_variables()))
      loss += v * v
      self.assertAllEqual([v], t.watched_variables())

  def testExplicitWatchedVariables(self):
    with backprop.GradientTape() as t:
      self.assertEqual(0, len(t.watched_variables()))
      v = resource_variable_ops.ResourceVariable(1.0)
      t.watch(v)
      self.assertAllEqual([v], t.watched_variables())

      t.reset()
      self.assertEqual(0, len(t.watched_variables()))
      t.watch(v)
      self.assertAllEqual([v], t.watched_variables())

  @test_util.assert_no_new_tensors
  def testGradientNone(self):

    def loss(x, l):
      return math_ops.reduce_mean(
          nn_ops.softmax_cross_entropy_with_logits(logits=x, labels=l),
          constant_op.constant([0]))

    logits = constant_op.constant([[0.0, 0.0]])
    labels = constant_op.constant([[1.0, 0.0]])
    # softmax_cross_entropy_with_logits returns two outputs and in this case the
    # gradient wrt the second is None.
    g, = backprop.gradients_function(loss, [0])(logits, labels)
    self.assertAllEqual(g.numpy(), [[-0.5, 0.5]])

  @test_util.run_in_graph_and_eager_modes
  def testGradientWithinTapeBlock(self):
    v1 = resource_variable_ops.ResourceVariable(1.)
    self.evaluate(v1.initializer)
    with backprop.GradientTape() as t:
      loss = 2 * v1
      grad = t.gradient(loss, v1)
    self.assertAllEqual(self.evaluate(grad), 2.0)

    with backprop.GradientTape(persistent=True) as t:
      loss = 2 * v1
      grad = t.gradient(loss, v1)
    self.assertAllEqual(self.evaluate(grad), 2.0)

  @test_util.run_in_graph_and_eager_modes
  def testNestedSelfContexts(self):
    v1 = resource_variable_ops.ResourceVariable(1.)
    self.evaluate(v1.initializer)
    with backprop.GradientTape() as t:
      with self.assertRaises(ValueError):
        with t:
          pass

  @test_util.assert_no_new_tensors
  def testSecondGrad(self):

    def first(x):
      l = constant_op.constant([[0.0]])
      x = nn_ops.softmax_cross_entropy_with_logits(labels=l, logits=x)
      x = math_ops.reduce_sum(x, constant_op.constant([0]))
      return x

    def second(x):
      grad = backprop.gradients_function(first, [0])(x)[0]
      return math_ops.reduce_sum(grad, constant_op.constant([0]))

    f = constant_op.constant([[0.1]])
    grad = backprop.gradients_function(second, [0])(f)[0]
    self.assertAllEqual([[0.0]], grad)

  @test_util.run_in_graph_and_eager_modes
  def testWatchingIsTapeLocal(self):
    x1 = resource_variable_ops.ResourceVariable(2.0, trainable=False)
    x2 = resource_variable_ops.ResourceVariable(2.0, trainable=False)

    with backprop.GradientTape() as tape1:
      with backprop.GradientTape() as tape2:
        tape1.watch(x1)
        tape2.watch([x1, x2])
        y = x1**3
        z = x2**2
        dy, dz = tape2.gradient([y, z], [x1, x2])
      d2y, d2z = tape1.gradient([dy, dz], [x1, x2])

    self.evaluate([x1.initializer, x2.initializer])
    self.assertEqual(self.evaluate(d2y), 12.0)
    self.assertIsNone(d2z)

  @test_util.assert_no_new_tensors
  def testMakeVJP(self):

    def f(x):
      return x * x

    wrapped_fn = backprop.make_vjp(f, persistent=False)
    result, vjp = wrapped_fn(constant_op.constant(3.0))
    self.assertAllEqual(result, 9.0)
    self.assertAllEqual(vjp(2.0)[0], 12.0)

  def testPersistentMakeVJP(self):

    def f(x):
      return x * x

    wrapped_fn = backprop.make_vjp(f, persistent=True)
    _, vjp = wrapped_fn(constant_op.constant(3.0))
    vjp_result1 = vjp(2.0)[0]
    vjp_result2 = vjp(2.0)[0]
    self.assertAllEqual(vjp_result1, vjp_result2, 12.0)

  @test_util.assert_no_new_tensors
  def testGradGrad(self):

    def sq(x):
      return x * x

    def grad(x):
      value = backprop.gradients_function(sq, [0])(x)[0]
      return value

    gradgrad = backprop.gradients_function(grad, [0])

    self.assertAllEqual(gradgrad(constant_op.constant(3.0))[0], 2.0)

  @test_util.assert_no_new_tensors
  def testGradGradExp(self):

    def grad(x):
      value = backprop.gradients_function(math_ops.exp, [0])(x)[0]
      return value

    gradgrad = backprop.gradients_function(grad, [0])

    self.assertAllEqual(gradgrad(constant_op.constant(0.0))[0], 1.0)

  @test_util.assert_no_new_tensors
  def testStopGradient(self):
    grad = backprop.gradients_function(
        lambda x: array_ops.stop_gradient(math_ops.argmax(x)))
    self.assertAllEqual(grad([0.0])[0], None)

  @test_util.assert_no_new_tensors
  def testArgmax(self):

    def argmax(x):
      i = math_ops.argmax(x)
      return array_ops.stop_gradient(i)

    grad = backprop.gradients_function(argmax)
    self.assertAllEqual(grad([0.0])[0], None)

  @test_util.run_gpu_only
  @test_util.assert_no_new_tensors
  def testGPU(self):

    def fn(x):
      with context.device('/gpu:0'):
        b = constant_op.constant(2.0)
        c = math_ops.add(x.gpu(), b)
        # TODO(apassos): remove cpu below by making TensorVSPace aware
        # of devices.
        return math_ops.add(c, constant_op.constant(3.0)).cpu()

    grad = backprop.gradients_function(fn, [0])(constant_op.constant(1.0))[0]
    self.assertAllEqual(grad, 1.0)

  @test_util.run_gpu_only
  @test_util.assert_no_new_tensors
  def testGPUImplicitGrad(self):
    with context.device('gpu:0'):
      v = resource_variable_ops.ResourceVariable(
          constant_op.constant(1.0), name='v')

    def f():
      with context.device('gpu:0'):
        return v.read_value()

    self.assertEqual(backprop.implicit_grad(f)()[0][0].cpu().numpy(), 1.0)

  @test_util.assert_no_new_tensors
  def testCPU(self):

    def fn(x):
      b = constant_op.constant(2.0)
      c = math_ops.add(x, b)
      return math_ops.add(c, constant_op.constant(3.0))

    grad = backprop.gradients_function(fn, [0])(constant_op.constant(1.0))[0]
    self.assertAllEqual(grad, 1.0)

  @test_util.run_gpu_only
  @test_util.assert_no_new_tensors
  def testTensorCopyGPU2CPU2GPU(self):

    def f(a, b):
      return a.cpu() + b.cpu()

    with context.device('/gpu:0'):
      a = constant_op.constant(1.0)
      b = constant_op.constant(2.0)

    grad = backprop.gradients_function(f, [0])(a, b)[0]
    self.assertAllEqual(grad, 1.0)

  @test_util.assert_no_new_tensors
  def testEmptyParams(self):

    def fn(a, b):
      return a * b

    x = constant_op.constant(1.0)
    y = constant_op.constant(2.0)
    dx, dy = backprop.gradients_function(fn)(x, y)
    self.assertAllEqual(dx, y.numpy())
    self.assertAllEqual(dy, x.numpy())

  @test_util.assert_no_new_tensors
  def testUnconnectedNone(self):
    v = resource_variable_ops.ResourceVariable(1.0, name='testUnconnectedNone')

    def f():
      v.read_value()
      return constant_op.constant(1.0)

    self.assertEqual(backprop.implicit_grad(f)()[0][0], None)

  @test_util.assert_no_new_tensors
  def testGradientTapeReEnterContext(self):
    g = backprop.GradientTape()
    with g:
      x = constant_op.constant(3.0)
      g.watch(x)
      y = 2 * x
    with g:
      z = 2 * y
    grad = g.gradient(target=z, sources=[x])
    self.assertEqual(self.evaluate(grad), [4.0])

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testGradientTapeRepeatedSource(self):
    with backprop.GradientTape(persistent=False) as g:
      x = constant_op.constant(3.0)
      g.watch(x)
      y = 2 * x
    grad = g.gradient(target=y, sources=[x, x])
    self.assertEqual(self.evaluate(grad), [2.0, 2.0])

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testPersistentGradientTapeRepeatedSource(self):
    with backprop.GradientTape(persistent=True) as g:
      x = constant_op.constant(3.0)
      y = constant_op.constant(5.0)
      g.watch(x)
      g.watch(y)
      z = x * x + x * y
    grad = g.gradient(target=z, sources=[x, x])
    self.assertEqual(self.evaluate(grad), [11.0, 11.0])
    grad = g.gradient(target=z, sources=[y, x])
    self.assertEqual(self.evaluate(grad), [3.0, 11.0])

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testGradientTapeStructure(self):
    with backprop.GradientTape(persistent=True) as g:
      # Using different constant values because constant tensors are
      # cached, leading to a different gradient then what one might expect.
      x1 = constant_op.constant(3.0)
      x2 = constant_op.constant(3.1)
      x3 = constant_op.constant(3.2)
      g.watch(x1)
      g.watch(x2)
      g.watch(x3)
      y = x1 + 2 * x2 + 3 * x3
    self.assertEqual(self.evaluate(g.gradient(y, x1)), [1.0])
    self.assertEqual(self.evaluate(g.gradient(y, (x1,))), (1.0,))
    self.assertEqual(self.evaluate(g.gradient(y, (x1, x2))), (1.0, 2.0))
    self.assertEqual(
        self.evaluate(g.gradient(y, [(x1, x2), (x2, x3)])), [(1.0, 2.0),
                                                             (2.0, 3.0)])
    self.assertEqual(
        self.evaluate(g.gradient(y, (x1, x2, [x1, x3]))),
        (1.0, 2.0, [1.0, 3.0]))
    self.assertEqual(
        self.evaluate(g.gradient(y, [x1, {
            'x2': x2,
            'x3': x3
        }])), [1.0, {
            'x2': 2.0,
            'x3': 3.0
        }])

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testGradientTape(self):
    with backprop.GradientTape() as g:
      x = constant_op.constant(3.0)
      g.watch(x)
      y = x * x
      with backprop.GradientTape() as gg:
        gg.watch(y)
        z = 2 * y
      inner_grad = gg.gradient(z, [y])[0]
      self.assertEqual(self.evaluate(inner_grad), 2.0)
      y += inner_grad
    grad = g.gradient(y, [x])[0]
    self.assertEqual(self.evaluate(grad), 6.0)

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testGadientTapeCalledOnConstantTarget(self):
    with backprop.GradientTape() as g:
      x = variables.Variable([3.0])
      y = variables.Variable([2.0])
    grad = g.gradient(x, y)
    self.assertAllEqual(grad, None)

  @test_util.run_in_graph_and_eager_modes
  @test_util.run_v1_only('b/120545219')
  def testGradientTapeWithCond(self):
    x = constant_op.constant(3.0)

    def true_fn():
      return x

    def false_fn():
      return x * x

    with backprop.GradientTape() as g:
      g.watch(x)
      y = tf_cond.cond(x < x, true_fn, false_fn)

    if not context.executing_eagerly():
      with self.assertRaisesRegex(NotImplementedError, 'tf.gradients'):
        dy = g.gradient(y, [x])[0]
    else:
      dy = g.gradient(y, [x])[0]
      self.assertEqual(self.evaluate(dy), 6.0)

  @test_util.run_in_graph_and_eager_modes
  @test_util.run_v1_only('b/120545219')
  def testGradientTapeWithWhileLoop(self):
    i = constant_op.constant(1)
    x = constant_op.constant(2.)

    def cond(i, _):
      return i < 3

    def body(i, x):
      return i + 1, x * 2

    with backprop.GradientTape() as g:
      g.watch([x])
      _, y = while_loop.while_loop(cond, body, [i, x])

    if not context.executing_eagerly():
      with self.assertRaisesRegex(NotImplementedError, 'tf.gradients'):
        dy = g.gradient(y, [x])[0]
    else:
      dy = g.gradient(y, [x])[0]
      self.assertEqual(self.evaluate(dy), 4.0)

  @test_util.assert_no_new_tensors
  def testGradientTapeGradientCalledMultipleTimes(self):
    with backprop.GradientTape() as g:
      x = constant_op.constant(3.0)
      g.watch(x)
      y = x * x
      z = y * y
    g.gradient(z, [x])
    with self.assertRaisesRegex(
        RuntimeError, 'A non-persistent GradientTape can only'):
      g.gradient(y, [x])

  @test_util.assert_no_new_tensors
  def testGradientTapeJacobianCalledMultipleTimes(self):
    with backprop.GradientTape() as g:
      x = constant_op.constant(3.0)
      g.watch(x)
      y = x * x
      z = y * y
    g.jacobian(z, [x])
    with self.assertRaisesRegex(
        RuntimeError, 'A non-persistent GradientTape can only'):
      g.jacobian(y, [x])

  @test_util.assert_no_new_tensors
  def testJacobianInsideGradientTapeScope(self):
    with backprop.GradientTape() as g:
      x = constant_op.constant(3.0)
      g.watch(x)
      y = x * x
      z = y * y
      self.assertAllClose(4. * 3. ** 3., g.jacobian(z, x))

  @test_util.assert_no_new_tensors
  def testBatchJacobianInsideGradientTapeScope(self):
    with backprop.GradientTape(persistent=True) as g:
      x = constant_op.constant([[3.0]])
      g.watch(x)
      y = x * x
      z = y * y
      self.assertAllClose([[[4. * 3. ** 3.]]], g.batch_jacobian(z, x))

  def testBatchJacobianParallelIterations(self):
    @def_function.function
    def f(persistent):
      with backprop.GradientTape(persistent=persistent) as t:
        x = constant_op.constant([[3.0]])
        t.watch(x)
        y = x * x
        z = array_ops.tile(y * y, [1, 16])
      return t.batch_jacobian(z, x, parallel_iterations=8)
    with self.assertRaisesRegex(RuntimeError,
                                'persistent=True.*parallel_iterations'):
      f(persistent=False)
    self.assertAllClose([[[4. * 3. ** 3.]] * 16], f(persistent=True))

  @test_util.assert_no_new_tensors
  def testGradientTapeBatchJacobianCalledMultipleTimes(self):
    with backprop.GradientTape() as g:
      x = constant_op.constant([[3.0]])
      g.watch(x)
      y = x * x
      z = y * y
    g.batch_jacobian(z, x)
    with self.assertRaisesRegex(
        RuntimeError, 'A non-persistent GradientTape can only'):
      g.batch_jacobian(y, [x])

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  @test_util.run_v1_only('b/120545219')
  def testPersistentTape(self):
    with backprop.GradientTape(persistent=True) as g:
      x = constant_op.constant(3.0)
      g.watch(x)
      y = x * x
      z = y * y
    dz_dx = g.gradient(z, [x])[0]
    self.assertEqual(self.evaluate(dz_dx), 4 * 3 * 3 * 3)
    dy_dx = g.gradient(y, [x])[0]
    self.assertEqual(self.evaluate(dy_dx), 2 * 3)
    del g

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testHigherOrderGradient(self):
    with backprop.GradientTape(persistent=True) as g:
      x = constant_op.constant(3.0)
      g.watch(x)
      y = x**3  # y       := x^3
      dy_dx = g.gradient(y, x)  # dy/dx   := 3x^2
      d2y_dx2 = g.gradient(dy_dx, x)  # d2y/dx2 := 6x
    d3y_dx3 = g.gradient(d2y_dx2, x)  # d3y/dx3 := 6
    x = 3
    self.assertAllClose(self.evaluate(y), x**3)
    self.assertEqual(self.evaluate(dy_dx), 3 * x**2)
    self.assertEqual(self.evaluate(d2y_dx2), 6 * x)
    self.assertEqual(self.evaluate(d3y_dx3), 6)
    del g

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testPersistentNestedTape(self):
    with backprop.GradientTape(persistent=True) as g:
      x = constant_op.constant(3.0)
      g.watch(x)
      y = x * x
      with backprop.GradientTape(persistent=True) as gg:
        gg.watch(y)
        z = 2 * y
      for _ in range(2):
        inner_grad = gg.gradient(z, [y])[0]
        self.assertEqual(self.evaluate(inner_grad), 2.0)
      y += inner_grad
      del gg
    grad = g.gradient(y, [x])[0]
    self.assertEqual(self.evaluate(grad), 6.0)
    grad = g.gradient(z, [x])[0]
    self.assertEqual(self.evaluate(grad), 12.0)
    del g

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testGradientTapeVariable(self):
    v = resource_variable_ops.ResourceVariable(1.0, name='v')
    self.evaluate(v.initializer)
    with backprop.GradientTape() as g:
      y = v * v
    grad = g.gradient(y, [v])[0]
    self.assertAllEqual(self.evaluate(grad), 2.0)

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testNestedGradients(self):
    x = constant_op.constant(3.0)
    with backprop.GradientTape() as g:
      g.watch(x)
      y = x * x
      z = y * y
    dz_dx, dz_dy = g.gradient(z, [x, y])
    self.assertEqual(self.evaluate(dz_dx), 108.0)
    self.assertEqual(self.evaluate(dz_dy), 18.0)

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testUnconnectedGradientsDefault(self):
    x = constant_op.constant(1.0)
    y = constant_op.constant(3.0)
    with backprop.GradientTape() as g:
      g.watch([x, y])
      z = y * 2
    dz_dx = g.gradient(z, x)
    self.assertEqual(dz_dx, None)

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testUnconnectedGradientsZeros(self):
    x = constant_op.constant(1.0, shape=[2, 2])
    y = constant_op.constant(3.0)
    with backprop.GradientTape() as g:
      g.watch([x, y])
      z = y * 2
    dz_dx = g.gradient(z, x, unconnected_gradients='zero')
    self.assertAllEqual([[0.0, 0.0], [0.0, 0.0]], self.evaluate(dz_dx))

  @test_util.assert_no_new_tensors
  @test_util.run_in_graph_and_eager_modes
  def testUnconnectedGradientsVariablesZeros(self):
    x = resource_variable_ops.ResourceVariable(
        constant_op.constant(1., shape=[2, 2]))
    self.evaluate(x.initializer)
    y = resource_variable_ops.ResourceVariable(constant_op.constant(3.))
    self.evaluate(y.initializer)
    with backprop.GradientTape() as g:
      g.watch([x, y])
      z = y * 2
    dz_dx = g.gradient(z, x, unconnected_gradients='zero')
    self.assertAllEqual([[0.0, 0.0], [0.0, 0.0]], self.evaluate(dz_dx))

  @test_util.run_in_graph_and_eager_modes
  def testUnknownUnconnectedGradientsValueGiven(self):
    x = constant_op.constant(1.0)
    y = constant_op.constant(1.0)
    with backprop.GradientTape() as g:
      g.watch([x, y])
      z = y * 2
    with self.assertRaisesRegex(
        ValueError, "Unknown value for unconnected_gradients: 'nonsense'"):
      g.gradient(z, x, unconnected_gradients='nonsense')

  @test_util.run_in_graph_and_eager_modes
  def testUnconnectedGradientsNestedDefunZeros(self):

    @def_function.function
    def f(x):
      return x * x

    @def_function.function
    def h(y):
      z = f(y)
      return array_ops.stop_gradient(z)

    x = constant_op.constant(1.0)
    with backprop.GradientTape() as g:
      g.watch(x)
      k = x + 2.
      y = h(k)

    dy_dx = g.gradient(y, x, unconnected_gradients='zero')
    self.assertEqual(0.0, self.evaluate(dy_dx))

  def testInvalidRecordOperationMessage(self):
    y = constant_op.constant(2.)
    x = constant_op.constant(1.)
    with backprop.GradientTape() as g:
      g.watch(x)
      record.record_operation('InvalidBackprop', [y], [x], lambda dy: [])
    with self.assertRaisesRegex(errors_impl.InternalError,
                                'InvalidBackprop.*too few gradients'):
      g.gradient(y, x)

  @test_util.assert_no_new_tensors
  def testEmptyParamsForValueAndGradFunction(self):

    def fn(a, b):
      return a * b

    val_and_grads_fn = backprop.val_and_grad_function(fn)

    x = 2.0
    y = 3.0
    val, (dx, dy) = val_and_grads_fn(x, y)
    self.assertAllClose(val, x * y)
    self.assertAllEqual(dx, y)
    self.assertAllEqual(dy, x)

  @test_util.assert_no_new_tensors
  def testNonEmptyParamsForValueAndGradFunction(self):

    def fn(a, b):
      return a * b

    val_and_grad_fn = backprop.val_and_grad_function(fn, params=[1])

    x = 2.0
    y = 3.0
    val, grads = val_and_grad_fn(x, y)
    self.assertAllClose(val, x * y)
    self.assertEqual(1, len(grads))
    self.assertAllEqual(grads[0], x)

  @test_util.run_gpu_only
  @test_util.assert_no_new_tensors
  def testTensorCopyCPU2GPU2CPU(self):
    # forward: a (cpu->gpu) -> add (gpu) -> c (gpu->cpu) -> add (cpu) -> e (cpu)
    # back: e (cpu) -> add (cpu) -> c (cpu->gpu) -> add (gpu) -> grad (gpu->cpu)
    def f(a, b):
      with context.device('/gpu:0'):
        c = math_ops.add(a.gpu(0), b.gpu(0))
      return math_ops.add(c.cpu(), constant_op.constant(3.0))

    with context.device('/cpu:0'):
      a = constant_op.constant(1.0)
      b = constant_op.constant(2.0)

    grad = backprop.gradients_function(f, [0])(a, b)[0]
    self.assertAllEqual(grad, 1.0)

  def testGetAttrType(self):
    typ = backprop.op_attr_type('Add', 'T')
    self.assertEqual(typ, int(pywrap_tfe.TF_ATTR_TYPE))

  def testGetAttrList(self):
    typ = backprop.op_attr_type('MaxPool', 'ksize')
    self.assertEqual(typ, [int(pywrap_tfe.TF_ATTR_INT)])

  def testMakeAttrType(self):
    self.assertEqual(dtypes.float32,
                     backprop.make_attr(int(pywrap_tfe.TF_ATTR_TYPE), 1))

  def testMakeAttrTypeList(self):
    self.assertEqual([dtypes.float32],
                     backprop.make_attr([int(pywrap_tfe.TF_ATTR_TYPE)], [1]))

  def testMakeAttrString(self):
    self.assertEqual(b'a',
                     backprop.make_attr(int(pywrap_tfe.TF_ATTR_STRING), 'a'))

  def testMakeAttrStringList(self):
    self.assertEqual(
        [b'a'], backprop.make_attr([int(pywrap_tfe.TF_ATTR_STRING)], ['a']))

  def testMulType(self):

    def mul(x):
      return math_ops._mul_dispatch(x, x)  # pylint: disable=protected-access

    self.assertAllEqual(backprop.gradients_function(mul)(3.0)[0].numpy(), 6.0)

  def testMakeAttrShape(self):
    for s in ([], None, [1, 2, 3], [None, None], [1, None, 3]):
      expected = tensor_shape.TensorShape(s).as_proto()
      actual = backprop.make_attr(int(pywrap_tfe.TF_ATTR_SHAPE), s)
      self.assertEqual(
          expected,
          actual,
          msg=('For shape %r, expected %r != %r actual' %
               (s, expected, actual)))

  def testMakeAttrShapeList(self):
    shape_list = [[], None, [1, 2, 3], [None, None], [1, None, 3]]
    self.assertEqual(
        [tensor_shape.TensorShape(s).as_proto() for s in shape_list],
        backprop.make_attr([int(pywrap_tfe.TF_ATTR_SHAPE)], shape_list))

  def testArgsGradientFunction(self):

    def f(*args):
      return args[0] * args[0]

    grad = backprop.gradients_function(f)
    self.assertAllEqual(grad(1.0)[0], 2.0)

  def testPartial(self):

    def f(x, y):
      return x * y

    part = functools.partial(f, constant_op.constant(2.0))
    self.assertAllEqual(
        backprop.gradients_function(part)(constant_op.constant(1.0))[0], 2.0)

  def testReturnSameThing(self):

    def f(x):
      return x, 2 * x

    self.assertAllEqual(backprop.gradients_function(f)(1.0)[0], 3.0)

  @test_util.assert_no_new_tensors
  def testExceptionSafety(self):

    def f(unused_x):
      raise ValueError()

    try:
      backprop.gradients_function(f)(1.0)
    except ValueError:
      pass

    def real_f(x):
      return x * x

    self.assertAllEqual(backprop.gradients_function(real_f)(1.0)[0], 2.0)

  @test_util.assert_no_new_tensors
  def testMultiValueConvertToTensor(self):
    x = resource_variable_ops.ResourceVariable(
        initial_value=array_ops.constant([1.0]), name='x')

    def fn():
      a = math_ops.add(x.value(), 1.0)
      # Make sure convert_to_tensor works correctly with list of TensorNodes.
      b = array_ops_stack.stack([a, a], axis=0)
      return math_ops.reduce_mean(b)

    grad = backprop.implicit_grad(fn)()[0][0]
    self.assertAllEqual([1.0], grad)

  def testOutput(self):

    def multiout(x):
      return x + 2, x * x

    x = constant_op.constant([0.0, 1.0, 2.0])

    grad = backprop.gradients_function(multiout)(x)[0]
    self.assertAllEqual([1.0, 3.0, 5.0], grad)

  def testMultiValuePreservesIfNotDiffedAgainst(self):

    def tfe_conv2d(timage, tkernel, conv2dstrides):
      return nn_ops.conv2d(timage, tkernel, conv2dstrides, 'SAME')

    i = constant_op.constant([[[[1.0]]]])
    k = constant_op.constant([[[[2.0]]]])
    s = [1, 1, 1, 1]

    grad = backprop.gradients_function(tfe_conv2d, params=(0,))(i, k, s)[0]
    self.assertAllEqual([[[[2.0]]]], grad)

  def testSameObjectForMultipleArguments(self):

    def f(x, y):
      return math_ops.multiply(x, y)

    g = backprop.gradients_function(f)

    def np_g(x, y):
      dx, dy = g(x, y)
      return [dx.numpy(), dy.numpy()]

    x = constant_op.constant(1.)
    self.assertAllEqual([1., 1.], np_g(x, x))
    x = 1.
    self.assertAllEqual([1., 1.], np_g(x, x))
    x = constant_op.constant([[1.]])
    self.assertAllEqual([[[1.]], [[1.]]], np_g(x, x))
    x = [[1.]]
    self.assertAllEqual([[[1.]], [[1.]]], np_g(x, x))

    v = resource_variable_ops.ResourceVariable(
        initial_value=1., name='testSameObjectForMultipleArguments.Variable')
    self.assertAllEqual([1., 1.], np_g(v, v))

  @test_util.assert_no_new_tensors
  def testImplicitGradientsCustomGradientAndCachedVariableValue(self):

    @custom_gradient.custom_gradient
    def my_square(x):
      result = math_ops.square(x)

      def grad(dr):
        return 2 * dr * x + 1

      return result, grad

    x = resource_variable_ops.ResourceVariable(
        initial_value=3., name='X.' + self.id())

    def f():
      return my_square(x)

    g = backprop.implicit_grad(f)

    grads_and_vars = g()
    self.assertEqual(1, len(grads_and_vars))
    grad, var = grads_and_vars[0]
    self.assertAllEqual(7, grad)
    self.assertAllEqual(x, var)

  def testJacobianCustomGradient(self):

    class MyCallable(object):

      def __init__(self):
        self.a = variables.Variable(1.)
        self.b = variables.Variable(2.)
        self.c = variables.Variable(3.)

      def __call__(self, x):
        return self.a * x * x + self.b * x + self.c

    @def_function.function
    def call(c, x):

      @custom_gradient.custom_gradient
      def _call():
        y = c(x)

        def grad(dy, variables=None):  # pylint: disable=redefined-outer-name
          with backprop.GradientTape(persistent=True) as g:
            g.watch(variables)
            y = c(x)
          grad_vars = [
              2 * math_ops.reduce_sum(dy * g.jacobian(y, v)) for v in variables
          ]
          del g
          return (), grad_vars

        return y, grad

      return _call()

    c = MyCallable()
    x = constant_op.constant([1., 2., 3.])
    with backprop.GradientTape(persistent=True) as g:
      g.watch([c.a, c.b, c.c])
      y = call(c, x)
    self.assertAllEqual(g.gradient(y, x), None)

  @test_util.assert_no_new_tensors
  def testCustomGradient(self):

    @custom_gradient.custom_gradient
    def my_mul(x, y):
      result = x * y

      def grad(dr):
        return [dr * y, dr * x]

      return result, grad

    lr = 0.25
    x = resource_variable_ops.ResourceVariable(2., name='x')

    def loss(x):
      return my_mul(2., x.read_value())

    loss_grads_fn = backprop.implicit_val_and_grad(loss)

    losses = []
    for _ in range(5):
      loss, grads_and_vars = loss_grads_fn(x)
      losses.append(loss.numpy())
      for (grad, var) in grads_and_vars:
        var.assign_sub(lr * grad)
    self.assertAllEqual(losses, [4.0, 3., 2., 1., 0.])

  @test_util.assert_no_new_tensors
  def testCustomGradientIdentity(self):

    @custom_gradient.custom_gradient
    def my_identity(x):

      def grad(dresult):
        return [2 * dresult]

      return x, grad

    self.assertAllEqual(backprop.gradients_function(my_identity)(1.0)[0], 2.0)

  def testDifferentiatingFunctionThatReturnsNone(self):

    def fn(x, y):
      result = x * y  # pylint: disable=unused-variable

    x = constant_op.constant(1)
    y = constant_op.constant(2)

    loss_grads_fn = backprop.implicit_val_and_grad(fn)
    with self.assertRaisesRegex(
        ValueError, 'Cannot differentiate a function that returns None; '
        'did you forget to return a value from fn?'):
      loss_grads_fn(x, y)

    val_and_grads_fn = backprop.val_and_grad_function(fn)
    with self.assertRaisesRegex(
        ValueError, 'Cannot differentiate a function that returns None; '
        'did you forget to return a value from fn?'):
      val_and_grads_fn(x, y)

  def testZerosCacheDoesntLeakAcrossGraphs(self):
    with ops.Graph().as_default():

      def get_grad():
        with ops.Graph().as_default(), self.cached_session():
          t = constant_op.constant(1, dtype=dtypes.float32, shape=(10, 4))
          x = constant_op.constant(2, dtype=dtypes.float32, shape=(10, 4))
          with backprop.GradientTape() as tape:
            tape.watch(x)
            x1, _ = array_ops.split(x, num_or_size_splits=2, axis=1)
            y1 = x1**2
            y = array_ops.concat([y1, t], axis=1)
          return self.evaluate(tape.gradient(y, x))

      grad1 = get_grad()
      grad2 = get_grad()

      self.assertAllEqual(grad1, grad2)

  @test_util.run_in_graph_and_eager_modes
  def testSelectivelyWatchVariables(self):
    x1 = resource_variable_ops.ResourceVariable(1.0)
    x2 = resource_variable_ops.ResourceVariable(1.0)
    with backprop.GradientTape(watch_accessed_variables=False) as tape:
      tape.watch(x2)
      y = x1**2
      z = x2**3
    self.assertTupleEqual(tape.watched_variables(), (x2,))
    dy, dz = tape.gradient([y, z], [x1, x2])
    self.evaluate([x1.initializer, x2.initializer])
    self.assertIsNone(dy)
    self.assertEqual(self.evaluate(dz), 3.0)

  @test_util.run_in_graph_and_eager_modes
  def testDifferentiatingScalarCache(self):
    # In the following test, if x2 = x1 (i.e the objects are the exact same),
    # then y is essentially, 2*x1, and dy/dx1 = 2.
    # When we had a pure scalar cache in eager, this would be the case. This
    # test prevents us from going back to that case.
    with backprop.GradientTape(persistent=False) as g:
      x1 = constant_op.constant(3.0)
      x2 = constant_op.constant(3.0)
      g.watch(x1)
      g.watch(x2)
      y = x1 + x2
    grad = g.gradient(target=y, sources=[x1])
    self.assertEqual(self.evaluate(grad), [1.0])

  def testVariablesAndConstantsProduceTheSameGradients(self):

    # In the following test, differentiating [y, z] against [a, b] gives:
    # (dy/da + dz/da, dy/db + dz/db).
    # If a and b are the same constant, dz/da will not be 0 (which it should
    # be).
    # This is solved by using variable since doing a read_value on a tensor will
    # produce a new tensor and corresponding TensorHandle, and not reuse the
    # same tensor (which would happen if we are using a cache and reusing
    # EagerTensor objects).
    def get_grads(a, b):
      with backprop.GradientTape() as tape:
        tape.watch([a, b])
        y = a**3
        z = b**2
      return tape.gradient([y, z], [a, b])

    gradients_constants = get_grads(
        constant_op.constant(2.0), constant_op.constant(2.0))
    gradients_variables = get_grads(
        resource_variable_ops.ResourceVariable(2.0),
        resource_variable_ops.ResourceVariable(2.0))
    self.assertAllEqual(gradients_constants, gradients_variables)

  def testUnknownShapes(self):
    with ops.Graph().as_default():
      with backprop.GradientTape() as tape:
        a = array_ops.placeholder(dtype=dtypes.float32, shape=None)
        tape.watch(a)
        b = a**3

      db_da = tape.gradient(b, a)

      with self.cached_session() as sess:
        self.assertEqual((8.0, 12.0), sess.run((b, db_da), feed_dict={a: 2.0}))

  @test_util.run_in_graph_and_eager_modes
  def testCustomGradientInEagerAndGraph(self):

    @custom_gradient.custom_gradient
    def f(x):
      y = x * x

      def grad(dy):
        return [4 * dy]

      return y, grad

    with backprop.GradientTape() as t:
      c = constant_op.constant(1.0)
      t.watch(c)
      g = f(c)
    self.assertAllEqual(self.evaluate(t.gradient(g, c)), 4.0)

  def testOverrideSecondOrderWithCustomGradient(self):

    @custom_gradient.custom_gradient
    def f(x):

      def first_order_grad(dz):

        @custom_gradient.custom_gradient
        def first_order_custom(unused_x):

          def h(ddz):
            return -2.1 * ddz

          return -1.1, h

        return dz * first_order_custom(x)

      return x + 10., first_order_grad

    c = constant_op.constant(1.)
    with backprop.GradientTape() as outer:
      outer.watch(c)
      with backprop.GradientTape() as inner:
        inner.watch(c)
        d = f(c)**4.
      dd = inner.gradient(d, c)
      self.assertAllClose(4. * f(c)**3. * -1.1, dd)
    self.assertAllClose(3. * 4. * f(c)**2. * -1.1 * -1.1 + 4. * f(c)**3. * -2.1,
                        outer.gradient(dd, c))

  @test_util.run_in_graph_and_eager_modes
  def testCustomGradientForwardprop(self):

    @custom_gradient.custom_gradient
    def f(x):
      z = 2. * tensor_util.constant_value(x)

      def g(dz):

        @custom_gradient.custom_gradient
        def first_order(unused_x, unused_dz):

          def second_order_and_transpose(unused_ddz):
            return 2.2, 3.1

          return 2.1, second_order_and_transpose

        return first_order(x, dz)

      return z, g

    with backprop.GradientTape(persistent=True) as t:
      with backprop.GradientTape() as tt:
        c = constant_op.constant(1.)
        t.watch(c)
        tt.watch(c)
        output_grad = array_ops.ones([])
        t.watch(output_grad)
        output = f(c)
        self.assertAllClose(2., output)
      gc = tt.gradient(output, c, output_gradients=output_grad)
      self.assertAllClose(2.1, gc)
    ggc = t.gradient(gc, c)
    self.assertAllClose(2.2, ggc)
    # Note that executed eagerly this kind of transpose is not efficient. But
    # from a tf.function we could prune out the first-order gradient
    # computation.
    transpose = t.gradient(gc, output_grad)
    self.assertAllClose(3.1, transpose)

  @test_util.run_in_graph_and_eager_modes
  def testWatchBadThing(self):
    g = backprop.GradientTape()
    with self.assertRaisesRegex(ValueError, 'ndarray'):
      g.watch(np.array(1.))

  def testWatchComposite(self):
    """Test that tape.watch expands composites and watches component Tensors."""
    with backprop.GradientTape() as t:
      values = constant_op.constant([1.0, 2.0], dtypes.float32)
      s = sparse_tensor.SparseTensor(
          indices=[[0, 0], [1, 2]], values=values, dense_shape=[3, 4])
      t.watch(s)
      z = sparse_ops.sparse_reduce_sum_v2(s)
    result = t.gradient(z, values)
    self.assertAllEqual(result, [1.0, 1.0])

  def testWatchedVariablesAfterNonPersistentGradientCall(self):
    with backprop.GradientTape(persistent=False) as tape:
      x = resource_variable_ops.ResourceVariable(1.0)
      tape.watch(x)
    tape.gradient(x, x)
    self.assertEqual((x,), tape.watched_variables())

  def testWatchedVariablesOnlyHasVariablesFromLastTape(self):
    with backprop.GradientTape(persistent=False) as tape:
      x = resource_variable_ops.ResourceVariable(1.0)
      tape.watch(x)
    with backprop.GradientTape(persistent=False) as tape:
      z = resource_variable_ops.ResourceVariable(2.0)
      tape.watch(z)
    tape.gradient(z, z)
    self.assertEqual((z,), tape.watched_variables())

  def testWatchedVariablesRespectReset(self):
    with backprop.GradientTape(persistent=False) as tape:
      x = resource_variable_ops.ResourceVariable(1.0)
      tape.watch(x)
      self.assertEqual((x,), tape.watched_variables())
      tape.reset()
      z = resource_variable_ops.ResourceVariable(2.0)
      tape.watch(z)
      self.assertEqual((z,), tape.watched_variables())
    tape.gradient(z, z)
    self.assertEqual((z,), tape.watched_variables())

  def testNameScope(self):

    def fn(x):
      with ops.name_scope('my_scope'):
        a = math_ops.cos(x)
        b = math_ops.cos(x)
        return math_ops.add(a, b)

    @def_function.function
    def grad_fn(x):
      return backprop.gradients_function(fn)(x)

    grad_ops = grad_fn.get_concrete_function(
        constant_op.constant(1.0)).graph.get_operations()
    num_sin_ops_found = 0
    for op in grad_ops:
      if op.type == 'Sin':
        num_sin_ops_found += 1
        self.assertIn('gradient_tape/my_scope/', op.name)
    self.assertEqual(num_sin_ops_found, 2)

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testRecomputeGradWithDifferentShape(self):
    if sys.version_info.major == 3 and sys.version_info.minor in (11, 12):
      # TODO(b/264947738)
      self.skipTest('Not working in Python 3.11+')

    @custom_gradient.recompute_grad
    def outer(x):
      return [x[0] + 1, x[1] + 1]

    x = [
        variables.Variable([1.0, 2.0], name='a'),
        variables.Variable(1.0, name='b')
    ]
    with backprop.GradientTape():
      y = outer(x)
      self.assertAllEqual(y[0], [2.0, 3.0])
      self.assertAllEqual(y[1], 2.0)

    @custom_gradient.recompute_grad
    def outer_dict(x):
      for key in x.keys():
        x[key] = x[key] + 1
      return x

    x = {x[0].ref(): x[0], x[1].ref(): x[1]}
    with backprop.GradientTape():
      y = outer_dict(x)
      y = list(y.values())
      self.assertAllEqual(y[0], [2.0, 3.0])
      self.assertAllEqual(y[1], 2.0)

  @parameterized.parameters([(True), (False)])
  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testRecomputeGradWithNestedFunctionAndWhileLoop(self, reduce_retracing):
    if sys.version_info.major == 3 and sys.version_info.minor in (11, 12):
      # TODO(b/264947738)
      self.skipTest('Not working in Python 3.11+')

    @custom_gradient.recompute_grad
    @def_function.function(reduce_retracing=reduce_retracing)
    def outer(x):

      @def_function.function(reduce_retracing=reduce_retracing)
      def middle(y):

        @def_function.function(reduce_retracing=reduce_retracing)
        def inner(z):
          return z + 1

        i = constant_op.constant(0.0)
        c = lambda y, i: i < 10.
        b = lambda y, i: (inner(y), i + 1.0)
        y, i = while_loop.while_loop(c, b, [y, i])

        return y

      return middle(x)

    with MemoryChecker() as memory_checker:
      for _ in range(5):
        x = variables.Variable(1.0, name='x')
        with backprop.GradientTape():
          y = outer(x)
          self.assertAllEqual(y, 11.0)

    memory_checker.report()
    memory_checker.assert_no_leak_if_all_possibly_except_one()

class JacobianTest(test.TestCase):

  def _jacobian(self, experimental_use_pfor):
    persistent = context.executing_eagerly and not experimental_use_pfor
    with backprop.GradientTape(persistent=persistent) as g:
      x = constant_op.constant([1., 2.])
      y = constant_op.constant([3., 4.])
      g.watch(x)
      g.watch(y)
      z = x * x * y
    jacobian = g.jacobian(
        z, [x, y], experimental_use_pfor=experimental_use_pfor)
    answer = [array_ops.diag(2 * x * y), array_ops.diag(x * x)]
    return jacobian, answer

  @test_util.run_v1_only('b/120545219')
  def testPfor(self):
    jacobian, answer = self._jacobian(experimental_use_pfor=True)
    for j, a in zip(jacobian, answer):
      self.assertAllEqual(a, j)

  @test_util.run_v1_only('b/120545219')
  def testWhileLoop(self):
    jacobian, answer = self._jacobian(experimental_use_pfor=False)
    for j, a in zip(jacobian, answer):
      self.assertAllEqual(a, j)

  @test_util.run_v1_only('b/120545219')
  def testPforDefun(self):

    @def_function.function
    def _f():
      return self._jacobian(experimental_use_pfor=True)

    jacobian, answer = _f()
    for j, a in zip(jacobian, answer):
      self.assertAllEqual(a, j)

  @test_util.run_v1_only('b/120545219')
  def testWhileLoopDefun(self):

    @def_function.function
    def _f():
      return self._jacobian(experimental_use_pfor=False)

    jacobian, answer = _f()
    for j, a in zip(jacobian, answer):
      self.assertAllEqual(a, j)

  @test_util.run_v1_only('b/120545219')
  def testPersistentTape(self):
    if not context.executing_eagerly():
      return
    with backprop.GradientTape() as g:
      x = constant_op.constant([1.0, 2.0])
      g.watch(x)
      y = x * x
    with self.assertRaisesRegex(RuntimeError, 'persistent'):
      g.jacobian(y, x, experimental_use_pfor=False)

  @test_util.run_v1_only('b/120545219')
  def test_parallel_iterations(self):
    with backprop.GradientTape(persistent=True) as g:
      x = constant_op.constant([[1., 2], [3, 4]])
      g.watch(x)
      y = math_ops.matmul(x, x)
    self.assertAllClose(
        g.jacobian(y, x, parallel_iterations=2),
        g.jacobian(y, x, parallel_iterations=3))

  @test_util.run_in_graph_and_eager_modes
  def test_nested_jacobian(self):
    if context.executing_eagerly():
      # TODO(agarwal): b/128842926
      self.skipTest('Conversion of function calls not implemented yet.')
    x = array_ops.ones((10, 2))
    with backprop.GradientTape(persistent=False) as g:
      g.watch(x)
      with backprop.GradientTape(persistent=False) as gg:
        gg.watch(x)
        y = math_ops.reduce_sum(math_ops.square(x))
      dy_x = gg.jacobian(y, x)
    dy_xx = g.batch_jacobian(dy_x, x)
    dy_xx_answer = [[[2., 0], [0, 2.]]] * 10
    self.assertAllClose(dy_xx_answer, self.evaluate(dy_xx))

  def test_nested_batch_jacobian_foldl(self):
    def _grad(f):
      def _grad_function(primal):
        with backprop.GradientTape() as tape:
          tape.watch(primal)
          primal_out = f(primal)
        return tape.batch_jacobian(primal_out, primal)
      return _grad_function

    def _func(x):
      return array_ops.reshape(
          functional_ops.foldl_v2(lambda a, b: math_ops.cos(a + b),
                                  array_ops.transpose(x)),
          [1, 1])

    f = _func
    x = constant_op.constant([[1., 2.]])
    for _ in range(2):
      theoretical, numerical = gradient_checker_v2.compute_gradient(f, [x])
      self.assertAllClose(theoretical, numerical, rtol=1e-3)
      f = _grad(f)
      expected_flat = array_ops.reshape(numerical, [-1])
      self.assertAllClose(expected_flat,
                          array_ops.reshape(f(x), [-1]),
                          rtol=1e-3)
      self.assertAllClose(expected_flat,
                          array_ops.reshape(def_function.function(f)(x), [-1]),
                          rtol=1e-3)

  def test_grad_jacobian_conv(self):
    def _inner(x):
      kernel = array_ops.ones([3, 3, 1, 9])
      with backprop.GradientTape() as tape:
        tape.watch(x)
        y = nn_ops.conv2d(x, kernel, strides=(1, 1), padding='SAME',
                          data_format='NHWC')
        reduced = math_ops.reduce_sum(y ** 2., axis=[2, 3])
      return math_ops.reduce_sum(tape.batch_jacobian(reduced, x))

    theoretical, numerical = gradient_checker_v2.compute_gradient(
        def_function.function(_inner), [array_ops.ones([10, 4, 4, 1])])
    self.assertAllClose(numerical, theoretical, rtol=1e-1)

    @def_function.function
    def _outer():
      with backprop.GradientTape() as tape:
        x = array_ops.ones([10, 4, 4, 1])
        tape.watch(x)
        y = _inner(x)
      return tape.gradient(y, x)

    self.assertAllClose(array_ops.reshape(numerical, [-1]),
                        array_ops.reshape(_outer(), [-1]), rtol=1e-1)

  @test_util.run_in_graph_and_eager_modes
  def test_indexed_slices(self):
    with backprop.GradientTape(persistent=True) as g:
      inp = random_ops.random_uniform([3, 2])
      g.watch(inp)
      output = nn.embedding_lookup(inp, [0, 2])
    self.assertAllClose(
        g.jacobian(output, inp, experimental_use_pfor=True),
        g.jacobian(output, inp, experimental_use_pfor=False))

  def test_foldl_partial_function(self):
    x = array_ops.zeros([3])
    with backprop.GradientTape(persistent=True) as tape:
      tape.watch(x)
      result = def_function.function(
          functools.partial(functional_ops.foldl_v2, lambda a, b: a + b))(
              x)
    self.assertAllClose([1., 1., 1.],
                        tape.jacobian(result, x, experimental_use_pfor=True))
    self.assertAllClose([1., 1., 1.],
                        tape.jacobian(result, x, experimental_use_pfor=False))

    # Non-persistent tapes take a different function gradient path, but also
    # work with pfor=True.
    x = array_ops.zeros([3])
    with backprop.GradientTape() as tape:
      tape.watch(x)
      result = def_function.function(
          functools.partial(functional_ops.foldl_v2, lambda a, b: a + b))(
              x)
    self.assertAllClose([1., 1., 1.],
                        tape.jacobian(result, x, experimental_use_pfor=True))

  def test_foldl_pure_function(self):

    @def_function.function
    def compute_jacobian(use_pfor):
      x = array_ops.zeros([3])
      with backprop.GradientTape(persistent=True) as tape:
        tape.watch(x)
        result = functools.partial(functional_ops.foldl_v2, lambda a, b: a + b)(
            x)
      return tape.jacobian(result, x, experimental_use_pfor=use_pfor)

    self.assertAllClose(compute_jacobian(use_pfor=True),
                        compute_jacobian(use_pfor=False))

  def test_cond_func_grad_jacobian(self):

    @def_function.function
    def f(x):
      y = tf_cond.cond(x > 0., lambda: x**3., lambda: x**2.)
      return y

    with backprop.GradientTape(persistent=True) as tape:
      x = constant_op.constant(1.)
      tape.watch(x)
      y = f(x)
      grad = tape.gradient(y, x)
    self.assertAllClose(3., grad)
    jacobian = tape.jacobian(grad, x, experimental_use_pfor=False)
    self.assertAllClose(6., jacobian)
    jacobian_pfor = tape.jacobian(grad, x, experimental_use_pfor=True)
    self.assertAllClose(6., jacobian_pfor)

  def test_empty_tensor_consistent_jacobian(self):
    variable = variables.Variable(1.0)
    inputs = (
        constant_op.constant(np.random.uniform(size=(0, 4))),
        constant_op.constant(np.random.uniform(size=(0, 3))),
    )
    with backprop.GradientTape(persistent=True) as tape:
      outputs = variable * math_ops.cast(
          array_ops.concat(inputs, axis=-1), dtypes.float32
      )

    jacobians_pfor = tape.jacobian(
        outputs,
        variable,
        experimental_use_pfor=True,
    )
    jacobians_loop = tape.jacobian(
        outputs,
        variable,
        experimental_use_pfor=False,
    )
    self.assertAllClose(jacobians_pfor, jacobians_loop)

@test_util.run_all_in_graph_and_eager_modes
class BatchJacobianTest(test.TestCase, parameterized.TestCase):

  def _batch_jacobian(self, experimental_use_pfor):
    persistent = context.executing_eagerly and not experimental_use_pfor
    with backprop.GradientTape(persistent=persistent) as g:
      x = constant_op.constant([[1., 2.], [3., 4.]])
      y = constant_op.constant([[3., 4.], [5., 6.]])
      g.watch(x)
      z = x * x * y
    batch_jacobian = g.batch_jacobian(
        z, x, experimental_use_pfor=experimental_use_pfor)
    answer = array_ops_stack.stack(
        [array_ops.diag(2 * x[0] * y[0]),
         array_ops.diag(2 * x[1] * y[1])])
    return batch_jacobian, answer

  def testPfor(self):
    batch_jacobian, answer = self._batch_jacobian(experimental_use_pfor=True)
    self.assertAllEqual(answer, batch_jacobian)

  def testWhileLoop(self):
    batch_jacobian, answer = self._batch_jacobian(experimental_use_pfor=False)
    self.assertAllEqual(answer, batch_jacobian)

  def testPforDefun(self):

    @def_function.function
    def _f():
      return self._batch_jacobian(experimental_use_pfor=True)

    batch_jacobian, answer = _f()
    self.assertAllEqual(answer, batch_jacobian)

  def testWhileLoopDefun(self):

    @def_function.function
    def _f():
      return self._batch_jacobian(experimental_use_pfor=False)

    batch_jacobian, answer = _f()
    self.assertAllEqual(answer, batch_jacobian)

  def testPersistentTape(self):
    if not context.executing_eagerly():
      return
    with backprop.GradientTape() as g:
      x = constant_op.constant([[1.0, 2.0]])
      g.watch(x)
      y = x * x
    with self.assertRaisesRegex(RuntimeError, 'persistent'):
      g.batch_jacobian(y, x, experimental_use_pfor=False)

  def testBadShape(self):
    x = random_ops.random_uniform([2, 3])
    with backprop.GradientTape() as g:
      y = array_ops.concat([x, x], axis=0)
    with self.assertRaisesRegex(ValueError, 'Need first dimension'):
      g.batch_jacobian(y, x)

  def testBadInputRank(self):
    x = random_ops.random_uniform([2])
    with backprop.GradientTape() as g:
      y = random_ops.random_uniform([2, 2])
    with self.assertRaisesRegex(ValueError, 'must have rank at least 2'):
      g.batch_jacobian(y, x)

  def testBadOutputRank(self):
    x = random_ops.random_uniform([2, 2])
    with backprop.GradientTape() as g:
      y = random_ops.random_uniform([2])
    with self.assertRaisesRegex(ValueError, 'must have rank at least 2'):
      g.batch_jacobian(y, x)

  def test_parallel_iterations(self):
    with backprop.GradientTape(persistent=True) as g:
      x = constant_op.constant([[1., 2], [3, 4]])
      g.watch(x)
      w = constant_op.constant([[1., 2, 3, 4], [5, 6, 7, 8]])
      y = math_ops.matmul(x, w)
    self.assertAllClose(
        g.batch_jacobian(y, x, parallel_iterations=2),
        g.batch_jacobian(y, x, parallel_iterations=3))

  @parameterized.parameters((True, True), (True, False), (False, True),
                            (False, False))
  def test_degenerate_shape(self, use_function, use_pfor):

    def f(x):
      with backprop.GradientTape(persistent=True) as tape:
        tape.watch(x)
        y = x**2
      return tape.batch_jacobian(y, x, experimental_use_pfor=use_pfor)

    if use_function:
      f = def_function.function(f)
    self.assertAllEqual([1, 0, 0], array_ops.shape(f(array_ops.zeros([1, 0]))))

  @parameterized.parameters((True,), (False))
  def test_zeros_type_correct(self, use_pfor):
    for dtype in [dtypes.float32, dtypes.float64]:
      @def_function.function
      def f(x):
        del x
        return constant_op.constant([[1.]], dtype=dtype)  # pylint: disable=cell-var-from-loop

      with backprop.GradientTape(persistent=True) as tape:
        x = constant_op.constant([[2.]], dtype=dtype)
        tape.watch(x)
        y = f(x)
      jac = tape.batch_jacobian(y, x, experimental_use_pfor=use_pfor)
      self.assertEqual(dtype, jac.dtype)
      self.assertAllClose([[[0.]]], jac)

      with backprop.GradientTape(persistent=True) as tape:
        x = constant_op.constant([[2.]], dtype=dtype)
        tape.watch(x)
        y = f(x)
      jac = tape.batch_jacobian(y, x, unconnected_gradients='zero',
                                experimental_use_pfor=use_pfor)
      self.assertEqual(dtype, jac.dtype)
      self.assertAllClose([[[0.]]], jac)

  def test_strided_slice(self):
    x = array_ops.ones([2, 4, 2])
    length = constant_op.constant([2, 3, 4, 4], dtype=dtypes.int64)
    with backprop.GradientTape() as tape:
      tape.watch(x)
      y = array_ops.repeat(x, [2], axis=1)
      y = y[:, :math_ops.reduce_max(length), :]
    tape.batch_jacobian(y, x)

class AggregateIndexedSlicesGradientsTest(test_util.TensorFlowTestCase):

  def _assert_indexed_slices_equal(self, left, right):
    self.assertAllEqual(
        self.evaluate(ops.convert_to_tensor(left)),
        self.evaluate(ops.convert_to_tensor(right)))

  def testNoGradients(self):
    self.assertIsNone(backprop_util.AggregateIndexedSlicesGradients([]))

  def testOneGradient(self):
    t = math_ops._as_indexed_slices(
        constant_op.constant([[1., 2.], [0, 0], [3., 4.]]))
    result = backprop_util.AggregateIndexedSlicesGradients([t])
    self._assert_indexed_slices_equal(t, result)

  def testMultipleGradients(self):
    t0 = math_ops._as_indexed_slices(
        constant_op.constant([[1., 2.], [0, 0], [3., 4.]]))
    t1 = math_ops._as_indexed_slices(
        constant_op.constant([[0., 0.], [5, 6], [7., 8.]]))
    total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]])
    result = backprop_util.AggregateIndexedSlicesGradients([t0, t1])
    self._assert_indexed_slices_equal(total, result)

  def testMultipleGradientsWithNones(self):
    t0 = math_ops._as_indexed_slices(
        constant_op.constant([[1., 2.], [0, 0], [3., 4.]]))
    t1 = math_ops._as_indexed_slices(
        constant_op.constant([[0., 0.], [5, 6], [7., 8.]]))
    t3 = None
    total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]])
    result = backprop_util.AggregateIndexedSlicesGradients([t0, t1, t3])
    self._assert_indexed_slices_equal(total, result)

  def testMixedTensorAndIndexedSlices(self):
    t0 = math_ops._as_indexed_slices(
        constant_op.constant([[1., 2.], [0, 0], [3., 4.]]))
    t1 = constant_op.constant([[0., 0.], [5, 6], [7., 8.]])
    total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]])
    result = backprop_util.AggregateIndexedSlicesGradients([t0, t1])
    self._assert_indexed_slices_equal(total, result)

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


"""
This file runs the ``Kalmnan-SEM`` function on the CiCMOD netCDF files.
The settings will be stored as "settings.yaml" in the ``results_path``.
Also multiple plots will be created and stored in the same directory.

An example would be :
python plot_CiCMOD.py --file_name climate_indices_FOCI.nc --experiment_name FOCI_default

NOTE:
    The kalman settings need to be manually changes in the file and can not be passed as arguments using ``argparse``.
"""

# %%
import argparse

from pathlib import Path

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
    "--file_name",
    help="The filename should be is usually ``climate_indices_{model}.nc``.",
)
parser.add_argument(
    "--dir_path",
    default=Path("data") / Path("earth_system_models") / Path("CiCMOD"),
    help="The source directory path where the netCDF file is stored.\n Default to ``.data/earth_system_models/CiCMOD``",
)
parser.add_argument(
    "--results_path",
    default=Path("results") / Path("CiCMOD"),
    help="The name to store the figures. \n Default to ``results/CiCMOD``",
)
parser.add_argument(
    "--experiment_name",
    default=Path("default"),
    help="The name of directory below ``results_path`` to store the figures in. \n Default to ``default``",
)

args = parser.parse_args()
print("Start imports.")

# from scipy import fftpack
from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
import yaml

from kalman_reconstruction import pipeline
from kalman_reconstruction.custom_plot import (
    adjust_lightness,
    ncols_nrows_from_N,
    plot_state_with_probability,
    set_custom_rcParams,
)
from kalman_reconstruction.statistics import compute_fft_spectrum, crosscorr, normalize
from tqdm import tqdm

set_custom_rcParams()
plt.rcParams["axes.grid"] = True

# %%
# Get arguments
file_name = args.file_name
dir_path = args.dir_path
results_path = args.results_path
experiment_name = args.experiment_name

# extract model name
model = file_name.split("_")[-1]
model = model.split(".")[0]

# Set all paths
REPO_PATH = Path(__file__).parent.parent.parent
data_path = REPO_PATH / Path(dir_path) / Path(file_name)
results_path = REPO_PATH / Path(results_path) / Path(experiment_name)
results_path.mkdir(parents=True, exist_ok=True)

SAVE_FIGURES = True

def save_fig(fig, relative_path, **kwargs):
    store_path = results_path / relative_path
    store_path.parent.mkdir(parents=True, exist_ok=True)
    if SAVE_FIGURES:
        fig.savefig(store_path, **kwargs)
    else:
        pass

# %%
# Set the Kalman and plot settings
time_slice = slice(0, 201)
rolling_window = 10 * 12
# Set Kalman-SEM settings
rng_seed = 834567
random_variance = 1
nb_iter_SEM = 25
observation_variables = ["AMO", "PDO_PC", "NAO_PC"]
state_variables = ["AMO", "PDO_PC", "NAO_PC", "latent"]
smoothed_variables = ["AMO"]

settings = dict(
    rolling_window=rolling_window,
    rng_seed=rng_seed,
    random_variance=random_variance,
    observation_variables=observation_variables,
    state_variables=state_variables,
    smoothed_variables=smoothed_variables,
    nb_iter_SEM=nb_iter_SEM,
    time_slice=dict(
        start=time_slice.start,
        stop=time_slice.stop,
    ),
    data_path=str(data_path),
)
with open(results_path / "settings.yaml", "w") as stream:
    stream.write(
        "#Settings used in the application of the Kalman_SEM on the CiCOD dataset.\n"
    )
    yaml.dump(data=settings, stream=stream, default_flow_style=False)

# %%
data_original = xr.load_dataset(data_path)

# %%
data = data_original.sel(time=time_slice).copy()
try:
    for var in smoothed_variables:
        data[var] = data[var].rolling(time=rolling_window).mean()
    data = data.isel(time=slice(rolling_window, -rolling_window))
except:
    pass

random_vars = ["latent"]
for random_var in random_vars:
    rng = np.random.default_rng(seed=rng_seed)
    pipeline.add_random_variable(
        ds=data,
        var_name=random_var,
        random_generator=rng,
        variance=random_variance,
        dim="time",
    )

# %%
kalman_results = pipeline.xarray_Kalman_SEM(
    ds=data,
    observation_variables=observation_variables,
    state_variables=state_variables,
    nb_iter_SEM=nb_iter_SEM,
)

# %%
kalman_states = pipeline.from_standard_dataset(kalman_results)

# %% Plot loglikelihood
fig, ax = plt.subplots(1, 1, figsize=(8, 4))
ax.plot(kalman_results.kalman_itteration, kalman_results.log_likelihod)
ax.set_xlabel("kalman iteration")
ax.set_ylabel("log likelihood")
fig.suptitle("CiCMOD | Loglikelihood Kalman SEM")
save_fig(fig, "CiCMOD_loglikelihood.png", dpi=400)

# %% Plot scatter truth against reconstruction
fig, axs = plt.subplots(
    nrows=len(state_variables),
    ncols=1,
    layout="constrained",
    figsize=(12, 7),
    sharex=True,
)
axs = axs.flatten()
for idx, var in enumerate(state_variables):
    handle1, handle2 = plot_state_with_probability(
        ax=axs[idx],
        x_value=kalman_results.time,
        state=kalman_results.states.sel(state_name=var),
        prob=kalman_results.covariance.sel(state_name=var, state_name_copy=var),
        line_kwargs=dict(label=f"{var}"),
        output=True,
    )

    # if "latent" not in var:
    # color = adjust_lightness(handle1[0].get_color(), )
    # axs[idx].plot(data.time, data[var], label = f"{var} truth", alpha = 0.7, linestyle = ":")
    axs[idx].set_title(var)
    axs[idx].set_ylabel("value")
    axs[idx].legend()
axs[idx].set_xlabel("time in years")
fig.suptitle("CiCMOD | Reconstruction against truth")
save_fig(fig, "CiCMOD_recons_truth.png", dpi=400)

# %% Plot scatter truth against reconstruction
# n_cols = len(state_variables)
# fig, axs = plt.subplots(layout = "constrained", figsize = (12,4), ncols = n_cols)
# axs = axs.flatten()

# for idx, var in enumerate(state_variables):
#     axs[idx].scatter(
#         kalman_states[var],
#         data[var],
#     )
#     axs[idx].set_xlabel("reconstruction")
#     axs[idx].set_ylabel("truth")
#     axs[idx].set_title(var)

# fig.suptitle("CiCMOD | Reconstruction against truth")
# save_fig(fig, "CiCMOD_recons_truth_scatter.png", dpi = 400)

# %% Plot scatter ``latent`` variable against reconstruction
data_vars = data.data_vars
row_col = ncols_nrows_from_N(len(data_vars))

fig, axs = plt.subplots(layout="constrained", figsize=(20, 20), **row_col)
axs = axs.flatten()

for idx, var in enumerate(data_vars):
    axs[idx].scatter(
        normalize(kalman_states.latent),
        normalize(data[var]),
        marker=".",
        alpha=0.75,
    )
    axs[idx].set_xlabel(var)
    axs[idx].set_ylabel("variable")
    axs[idx].set_title(var)

fig.suptitle("CiCMOD | Latent variable against climate indices")
save_fig(fig, "CiCMOD_latent_indices_scatter.png", dpi=400)

# %% [markdown]
### Compute lagged cross correlation and covariance

# %% Calculate lagged cross correlation for all variables in the CiCMOD dataset against ``latent``
lag_years = np.arange(-30, 30, 1)
data_vars = data.data_vars
data_ccor_list = []
kalman_ccor_list = []
for idx, var in tqdm(enumerate(data_vars)):
    for lag in lag_years:
        # because data is stored in monthly form, we need to multiply the shift by 12 to have the lag in years
        lag_months = lag * 12
        # # No need to calculate covarinace
        # # calculate the covariance
        # ccov = xr.cov(data[var], kalman_states.latent.shift(time=lag*12), dim = "time").values
        # da_ccov = xr.DataArray(
        #     data = ccov[np.newaxis],
        #     dims=["lag_years"],
        #     coords = dict(
        #         lag_years = (["lag_years"], [lag]),
        #     )
        # )
        # da_ccov = da_ccov.rename(var)
        # da_ccov_list.append(da_ccov)

        # calculate the lagged cross correlation
        data_ccor = crosscorr(
            ds1=data[var], ds2=kalman_states.latent, lag=lag_months, dim="time"
        )
        # create the DataArray
        data_ccor = xr.DataArray(
            data=data_ccor.values[np.newaxis],
            dims=["lag_years"],
            coords=dict(
                lag_years=(["lag_years"], [lag]),
            ),
        )
        data_ccor = data_ccor.rename(var)
        data_ccor_list.append(data_ccor)
        # Also for all variables try to make cross correaltino with the results from the kalman SEM
        try:
            kalman_ccor = crosscorr(
                ds1=kalman_states[var],
                ds2=kalman_states.latent,
                lag=lag_months,
                dim="time",
            )
            # Create the DataArray
            kalman_ccor = xr.DataArray(
                data=kalman_ccor.values[np.newaxis],
                dims=["lag_years"],
                coords=dict(
                    lag_years=(["lag_years"], [lag]),
                ),
            )
            kalman_ccor = kalman_ccor.rename(var)
            kalman_ccor_list.append(kalman_ccor)
        except:
            pass

# Merge all DataArrays to Datasets
data_ccor = xr.merge(data_ccor_list)  # lagged cross correlation with all CICMOD vars
kalman_ccor = xr.merge(
    kalman_ccor_list
)  # lagged cross correlation with all vars avaibales in the kalman results.

# %% STATE VARIABLES | Plot Lagged correlation of latent variable against state variables
data_vars = state_variables
row_col = ncols_nrows_from_N(len(data_vars))

fig, axs = plt.subplots(
    layout="constrained",
    figsize=(12, 4),
    sharex=True,
    sharey=True,
    ncols=len(data_vars),
)
axs = axs.flatten()

for idx, var in enumerate(data_vars):
    axs[idx].step(data_ccor.lag_years, np.abs(data_ccor[var]), label="turth")
    axs[idx].step(kalman_ccor.lag_years, np.abs(kalman_ccor[var]), label="recons")
    axs[idx].set_xlabel("lag in years")
    axs[idx].set_ylabel("correlation")
    axs[idx].set_title(var)
    axs[idx].legend()

axs[idx].set_ylim((-1, 1))

fig.suptitle(
    "CiCMOD | Absolute value of lagged correlation of latent variable against state variables"
)
save_fig(fig, "CiCMOD_latent_states_lagged_corr.png", dpi=400)

# %% ALL INDICES | lagged correlation of latent variable against state variables
data_vars = data.data_vars
row_col = ncols_nrows_from_N(len(data_vars))

fig, axs = plt.subplots(
    layout="constrained", figsize=(20, 20), sharex=True, sharey=True, **row_col
)
axs = axs.flatten()

for idx, var in enumerate(data_vars):
    axs[idx].step(data_ccor.lag_years, data_ccor[var], label="turth")
    try:
        axs[idx].step(kalman_ccor.lag_years, kalman_ccor[var], label="recons")
    except:
        pass
    axs[idx].set_xlabel("lag in years")
    axs[idx].set_ylabel("correlation")
    axs[idx].set_title(var)
    axs[idx].legend()

extend = np.max(np.abs(axs[idx].get_ylim()))
axs[idx].set_ylim((-extend, extend))

fig.suptitle("CiCMOD | Lagged correlation of latent against climate indices.")
save_fig(fig, "CiCMOD_latent_indices_lagged_corr.png", dpi=400)

# %% ALL INDICES | Absolute value of lagged correlation of latent variable against state variables
data_vars = data.data_vars
row_col = ncols_nrows_from_N(len(data_vars))

fig, axs = plt.subplots(
    layout="constrained", figsize=(20, 20), sharex=True, sharey=True, **row_col
)
axs = axs.flatten()

for idx, var in enumerate(data_vars):
    axs[idx].step(data_ccor.lag_years, np.abs(data_ccor[var]), label="turth")
    try:
        axs[idx].step(kalman_ccor.lag_years, np.abs(kalman_ccor[var]), label="recons")
    except:
        pass
    axs[idx].set_xlabel("lag in years")
    axs[idx].set_ylabel("correlation")
    axs[idx].set_title(var)
    axs[idx].legend()

axs[idx].set_ylim((0, 1))

fig.suptitle(
    "CiCMOD | Absolute value lagged correlation of latent against climate indices."
)
save_fig(fig, "CiCMOD_latent_indices_lagged_corr_abs.png", dpi=400)

# %% [markdown]
# ### Perform frequency analysis on input and output data

# %% Frequency analysis of all state variables.
data_vars = state_variables
fig, axs = plt.subplots(
    layout="constrained",
    figsize=(12, 4),
    sharex=True,
    sharey=True,
    ncols=len(data_vars),
)
axs = axs.flatten()

for idx, var in enumerate(data_vars):
    ax = axs[idx]
    xf, yf, yf_plot, f_min, f_max = compute_fft_spectrum(
        time=data.time, signal=data[var].values
    )
    ax.loglog(xf, yf_plot, label="truth", alpha=0.7)
    (xf, yf, yf_plot, f_min, f_max) = compute_fft_spectrum(
        time=kalman_states.time, signal=kalman_states[var].values
    )
    ax.loglog(xf, yf_plot, label="reconst.", alpha=0.7)

    f_min = 1 / 200  #  years^{-1}
    ax.set_xlim((f_min, f_max))
    xticks = ax.get_xticks().copy()  # 1/years
    new_ticks = np.round(1 / xticks, decimals=2)  # years
    ax.set_xticks(ticks=xticks, labels=new_ticks)
    ax.tick_params(axis="x", rotation=45)
    ax.set_xlim((f_min, f_max))

    ax.set_xlabel(r"Period in years")
    ax.set_ylabel("Power in ????")
    ax.set_title(var)
    ax.legend()

ax.set_ylim(bottom=10 ** (-5.5), top=10 ** (0.5))
fig.suptitle("CiCMOD | Frequency spectrum state variables")
save_fig(fig, "CiCMOD_fft_states.png", dpi=400)


from PyQt5.QtGui import QTextFormat

from ..pktk import *

EXTRASELECTION_FILTER_KEEP =    0x0001
EXTRASELECTION_FILTER_REMOVE =  0x0002
EXTRASELECTION_FILTER_FCTTRUE = 0x0003

def flatten(items):
    """Return a flatten list whatever the number of nested level the list have

    f=flatten([1,2,[3,4],[5,[6,7],[8,[9,10]]]])

    f: [1,2,3,4,5,6,7,8,9,10]
    """
    returned = []
    for item in items:
        if isinstance(item, (list, tuple)):
            returned.extend(flatten(item))
        else:
            returned.append(item)
    return returned

def rotate(items, shiftValue=1):
    """Rotate list

    - Positive `shiftValue` will rotate to right
    - Negative `shiftValue` will rotate to left

    l=[1,2,3,4]

    x=rotate(l, 1)
    x: [4,1,2,3]

    x=rotate(l, -1)
    x: [2,3,4,1]

    """
    shiftValue = shiftValue % len(items)
    if shiftValue == 0:
        # no rotation...
        return items
    # do rotation
    return items[-shiftValue:] + items[:-shiftValue]

def unique(items):
    """Return list of items with duplicate removed

    Initial order list is preserved
    Note: works with non-hashable object
    """
    # a faster method could be return list(set(items))
    # but it doesn't work with non-hashable object (like a QColor for example)
    returned = []
    for item in items:
        if item not in returned:
            returned.append(item)
    return returned

def sortExtraSelections(extraSelection):
    """Sort given `extraSelection` list from item's format.property

    1) items with format.property(QTextFormat.UserProperty)
       - sorted by extraSelection.format.property(QTextFormat.UserProperty) value
    2) items without format.property(QTextFormat.UserProperty)
       - no special order, just set 'after items with property QTextFormat.UserProperty'
    """
    def checkType(value):
        returned = value.format.property(QTextFormat.UserProperty)
        if returned:
            return returned
        return 0xFFFFFFFF
    if not isinstance(extraSelection, list):
        raise EInvalidType('Given `extraSelection` must be a <list>')
    extraSelection.sort(key=checkType)

def filterExtraSelections(extraSelection, filterValue, filterRule=EXTRASELECTION_FILTER_KEEP, filterProperty=QTextFormat.UserProperty, sortResult=False, stopOnFirst=False, removed=[]):
    """Sort given `extraSelection` list from item's property

    If `filterProperty` is given, define on which extraSelection.format.property(filterProperty) value is filtered

    Given `filterRule` define what to do:
    - EXTRASELECTION_FILTER_KEEP (default):
        only keep items for which extraSelection.format.property(filterProperty) equal `filterValue`
        (in this case, items without property are removed)
    - EXTRASELECTION_FILTER_REMOVE:
        only keep items for which extraSelection.format.property(filterProperty) do not equal `filterValue`
        (in this case, items without property are removed)
    - EXTRASELECTION_FILTER_FCTTRUE
        in this cas `filterValue` is a callable function for which:
        - extraSelection item is provided
        - returned result is a boolean:
            . True = keep item
            . False = remove item

    If `sortResult` is True, filtered list is sorted, otherwise returned list keep initial item orders

    If `stopOnFirst` is True (default is False), stop to filter on first occurence found
    This can be used when we know there's normally no more than one occurence of item to filter in list
    """
    if not isinstance(extraSelection, list):
        raise EInvalidType('Given `extraSelection` must be a <list>')

    removed.clear()
    if filterRule == EXTRASELECTION_FILTER_FCTTRUE:
        if not callable(filterValue):
            raise EInvalidType("Given `filterValue` must be a callable when `filterRule` is set to EXTRASELECTION_FILTER_FCTTRUE")

        index = len(extraSelection) - 1
        while index >= 0:
            if filterValue(extraSelection[index]) is False:
                removed.append(extraSelection.pop(index))
                if stopOnFirst:
                    return
            index -= 1
    elif filterRule == EXTRASELECTION_FILTER_KEEP:
        index = len(extraSelection) - 1
        while index >= 0:
            if extraSelection[index].format.property(filterProperty) != filterValue:
                removed.append(extraSelection.pop(index))
                if stopOnFirst:
                    return
            index -= 1
    elif filterRule == EXTRASELECTION_FILTER_REMOVE:
        index = len(extraSelection) - 1
        while index >= 0:
            if extraSelection[index].format.property(filterProperty) == filterValue:
                removed.append(extraSelection.pop(index))
                if stopOnFirst:
                    return
            index -= 1


# python program to parse json

import json
import os

os.chdir('c:/Users/mayur/OneDrive/Documents/Study Material/Python/Interview/challenge-2/')

f = open('response.json', 'r')
result = f.read()
f.close()

my_dict = json.loads(result)

# Take key input from user

key = input("Enter metadata key")
val = my_dict[key]

print(val)


#!/usr/bin/env python
# -*- coding: utf-8
# Compute distance between nerve rootlets and PMJ, nerve rootlets and discs along centerline
#
# For usage, type: python compute_neck_angle.py -h

# Authors: Sandrine Bédard

import argparse
import csv
import math
from webbrowser import get
import numpy as np
import nibabel as nib
import os

def get_parser():
    parser = argparse.ArgumentParser(
        description="Compute cobb angle beteen C2 and C6")
    parser.add_argument('-centerline', required=True, type=str,
                        help=".csv file of the centerline.")
    parser.add_argument('-disclabel', required=True, type=str,
                        help="Labels of the intervertebral discs.")
    parser.add_argument('-subject', required=True, type=str,
                        help="Subject ID")
    parser.add_argument('-ses', required=True, type=str,
                        help="ses")
    parser.add_argument('-o', required=False, type=str,
                        default='angle.csv',
                        help="Output csv filename.")

    return parser

def save_Nifti1(data, original_image, filename):
    empty_header = nib.Nifti1Header()
    image = nib.Nifti1Image(data, original_image.affine, empty_header)
    nib.save(image, filename)

def find_angle(v1, v2):

    unit_vector1 = v1 / np.linalg.norm(v1)
    unit_vector2 = v2 / np.linalg.norm(v2)

    dot_product = np.dot(unit_vector1, unit_vector2)

    angle = np.arccos(dot_product) #angle in radian

    angle_deg = angle*180/np.pi
    return round(angle_deg, 4)

def unit_vector(vector):
    """ Returns the unit vector of the vector.  """
    return vector / np.linalg.norm(vector)

def get_vector_from_point(p1, p2):
    v = p2 - p1
    return unit_vector(v)

def project_vector(u, n):
    # Norm of n vector
    n_norm = np.linalg.norm(n)
    proj_of_u_on_n = (np.dot(u, n)/n_norm**2)*n
    return u - proj_of_u_on_n

def main():
    parser = get_parser()
    args = parser.parse_args()

    # Create an array with centerline coordinates
    centerline = np.genfromtxt(args.centerline, delimiter=',')
    #centerline = nib.load(args.centerline).get_fdata()
    disc_label = nib.load(args.disclabel)
    idx_2 = np.where(disc_label.get_fdata() == 2.0)[-1]
    idx_6 = np.where(disc_label.get_fdata() == 6.0)[-1]
    idx_2_corr = np.abs(centerline[2] - idx_2).argmin()  # centerline doesn't necessarly start at the index 0 if the segmentation is incomplete
    idx_6_corr = np.abs(centerline[2] - idx_6).argmin()  # centerline doesn't necessarly start at the index 0 if the segmentation is incomplete
    len = 3

    p1 = centerline[:,idx_2_corr - len]
    p2 = centerline[:,idx_2_corr + len]
    p3 = centerline[:,idx_6_corr - len]
    p4 = centerline[:,idx_6_corr + len]

    v1 = get_vector_from_point(p1=p1, p2=p2)

    print('v1', v1)
    v2 = get_vector_from_point(p1=p3, p2=p4)
    print('v1', v2)
    # Project v1 and v2 on X (R-L)
    n = np.array([1,0,0])
    v1_proj = project_vector(v1, n)
    v2_proj = project_vector(v2, n)
    print('v1_proj', v1_proj, 'v2_proj', v2_proj)
    angle = find_angle(v1, v2)
    angle_proj = find_angle(v1_proj, v2_proj)

    print('angle', angle, 'degrees')
    print('angle proj', angle_proj, 'degrees')
    subject = args.subject
    ses = args.ses
    fname_out = args.o
    if not os.path.isfile(fname_out):
        with open(fname_out, 'w') as csvfile:
            header = ['Subject', 'ses', 'Angle (degrees)']
            writer = csv.DictWriter(csvfile, fieldnames=header)
            writer.writeheader()
    with open(fname_out, 'a') as csvfile:
        spamwriter = csv.writer(csvfile, delimiter=',')
        line = [subject, ses, angle_proj]
        spamwriter.writerow(line)

if __name__ == '__main__':
    main()


import requests
from setup import INFO
import re
from bs4 import BeautifulSoup as Soup

class Engine:
    def __init__(e, name, create_query):
        e.name = name
        e.create_query = create_query

    def get_search_url(e, query):
        INFO("original query: {}".format(query))
        # lower case the query
        query = query.lower()
        # replace spaces
        query = '%20'.join(query.split(' '))
        # join together
        search_url = e.create_query(query)
        INFO("parsed search_url is {}".format(search_url))
        return search_url

    def get_magnet(e, query):
        search_url = e.get_search_url(query)
        INFO("requesting from {}".format(e.name))
        response = requests.get(search_url)
        if response.status_code != 200:
            INFO("an error {} occured with {} request".format(response.status_code, search_url))
            return None
        else:
            INFO("request {} succeeded".format(search_url))
            soup = Soup(response.text, features="html.parser")
            magnet = soup.find(href = re.compile('magnet'))
            if magnet is None:
                INFO("No result found in {}".format(e.name))
            magnet_link = magnet['href']
            INFO("returning magnet {}".format(magnet_link))
            return magnet_link

TPB = Engine(
    "TPB",
    lambda q: "https://thepiratebay.org/search/{}/0/7/0".format(q)
    )
KICKASS = Engine(
    "KICKASS",
    lambda q: "https://katcr.co/usearch/torrents-search.php?q={}".format(q)
    )
EZTV = Engine(
    "EZTV",
    lambda q: "https://eztv.ag/search/{}".format(q)
    )
L337 = Engine(
    "L337",
    lambda q: "https://1337x.to/search/{}/1/".format('+'.join(q.split('%20')))
    )

ENGINES = [
    TPB,
    EZTV,
    L337
]