class_docstr: Jacobian test cases.
class_docstr_tok: ["Jacobian","test","cases","."]
docstr: Test function value_and_jacobian.
docstr_tok: ["Test","function","value_and_jacobian","."]
code: def test_value_and_jacobian(self, dtype):
      """Test function value_and_jacobian."""
      # Shape [2]
      x = tf.range(1, 3, dtype=dtype)
      func = functools.partial(quadratic, x=x)
      with self.subTest("Quadratic"):
          with self.subTest("SingleTensor"):
              func = functools.partial(quadratic, x=x)
              # Shape [3]
              ps = tf.constant([1.0, 2.0, -1.0], dtype=dtype)
              values, jacobian = self.evaluate(
                  tff.math.value_and_jacobian(func, ps)
              )
              with self.subTest("Values"):
                  # Shape [2]
                  expected_values = [2.0, 7.0]
                  self.assertEqual(values.shape, (2,))
                  np.testing.assert_allclose(values, expected_values)
              with self.subTest("Jacobian"):
                  # Shape [2, 3]
                  expected_jacobian = [[1.0, 1.0, 1.0], [4.0, 2.0, 1.0]]
                  self.assertEqual(
                      jacobian.shape,
                      (
                          2,
                          3,
                      ),
                  )
                  np.testing.assert_allclose(
                      jacobian, expected_jacobian
                  )
          with self.subTest("BatchedTensor"):
              func = functools.partial(quadratic, x=x)
              # Shape [2, 3]
              ps = tf.constant(
                  [[1.0, 2.0, -1.0], [2.0, 0.0, 0.0]], dtype=dtype
              )
              values, jacobian = self.evaluate(
                  tff.math.value_and_jacobian(func, ps)
              )
              with self.subTest("Values"):
                  # Shape [2, 2]
                  expected_values = [[2.0, 7.0], [2.0, 8.0]]
                  self.assertEqual(
                      values.shape,
                      (
                          2,
                          2,
                      ),
                  )
                  np.testing.assert_allclose(values, expected_values)
              with self.subTest("Jacobian"):
                  # Shape [2, 2, 3]
                  expected_jacobian = [
                      [[1.0, 1.0, 1.0], [4.0, 2.0, 1.0]],
                      [[1.0, 1.0, 1.0], [4.0, 2.0, 1.0]],
                  ]
                  self.assertEqual(jacobian.shape, (2, 2, 3))
                  np.testing.assert_allclose(
                      jacobian, expected_jacobian
                  )
  
      x = tf.range(1, 3, dtype=dtype)
      y = tf.range(-3, -1, dtype=dtype)
      func = functools.partial(multivariate_quadratic, x=x, y=y)
      with self.subTest("MultivariateQuadratic"):
          with self.subTest("SingleTensor"):
              # Shape [4]
              ps = tf.constant([1.0, 2.0, -1.0, 0.0], dtype=dtype)
              values, jacobian = self.evaluate(
                  tff.math.value_and_jacobian(func, ps)
              )
              with self.subTest("Values"):
                  # Shape [2]
                  expected_values = [-14.0, -8.0]
                  self.assertEqual(values.shape, (2,))
                  np.testing.assert_allclose(values, expected_values)
              with self.subTest("Jacobian"):
                  # Shape [2, 4]
                  expected_jacobian = [
                      [1.0, -3.0, 9.0, 1.0],
                      [4.0, -4.0, 4.0, 1.0],
                  ]
                  self.assertEqual(jacobian.shape, (2, 4))
                  np.testing.assert_allclose(
                      jacobian, expected_jacobian
                  )
          with self.subTest("BatchedTensor"):
              # Shape [2, 4]
              ps = tf.constant(
                  [[1.0, 2.0, -1.0, 0.0], [0.0, 1.0, -2.0, 1.0]],
                  dtype=dtype,
              )
              values, jacobian = self.evaluate(
                  tff.math.value_and_jacobian(func, ps)
              )
              with self.subTest("Values"):
                  # Shape [2, 2]
                  expected_values = [[-14.0, -8.0], [-20.0, -11.0]]
                  self.assertEqual(
                      values.shape,
                      (
                          2,
                          2,
                      ),
                  )
                  np.testing.assert_allclose(values, expected_values)
              with self.subTest("Jacobian"):
                  # Shape [2, 2, 4]
                  expected_jacobian = [
                      [[1.0, -3.0, 9.0, 1.0], [4.0, -4.0, 4.0, 1.0]],
                      [[1.0, -3.0, 9.0, 1.0], [4.0, -4.0, 4.0, 1.0]],
                  ]
                  self.assertEqual(jacobian.shape, (2, 2, 4))
                  np.testing.assert_allclose(
                      jacobian, expected_jacobian
                  )
  
file_path: tf-quant-finance/tf_quant_finance/math/jacobian_test.py
imports: import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 2 for calling the JacobianTest obj's underlying member method code with example usage: obj.test_value_and_jacobian(dtype) without return types