docstr: Test batching.
docstr_tok: ["Test","batching","."]
code: def test_batch(self):
      """Test batching."""
      dtype = np.float64
      times = tf.constant([[2.0, 2.5, 3], [2.0, 2.5, 3]], dtype=dtype)
      strikes = tf.constant([[16, 22, 35], [16, 22, 35]], dtype=dtype)
  
      times_data = tf.constant(
          [[1.5, 2.5, 3.5, 4.5, 5.5], [1.2, 2.2, 3.5, 4.5, 5.5]],
          dtype=dtype,
      )
      # Corresponding squared volatility values
      sigma_square_data = tf.constant(
          [
              [0.15, 0.25, 0.35, 0.4, 0.45, 0.4],
              [0.2, 0.35, 0.55, 0.45, 0.4, 0.6],
              [0.3, 0.45, 0.25, 0.4, 0.5, 0.65],
              [0.25, 0.25, 0.45, 0.25, 0.5, 0.55],
              [0.35, 0.35, 0.25, 0.4, 0.55, 0.65],
          ],
          dtype=dtype,
      )
      # Interpolation is done for the total variance
      total_variance = (
          tf.expand_dims(times_data, -1) * sigma_square_data
      )
      # Corresponding strike values. Notice we need to broadcast to the shape of
      # `sigma_square_data`
      strike_data = tf.broadcast_to(
          tf.constant([15, 25, 35, 40, 50, 55], dtype=dtype), [5, 6]
      )
      sigma_square_data = tf.broadcast_to(sigma_square_data, [2, 5, 6])
      strike_data = tf.broadcast_to(strike_data, [2, 5, 6])
      # Interpolate total variance on the grid [times, strikes]
      interpolator = interpolation_2d.Interpolation2D(
          times_data, strike_data, total_variance, dtype=dtype
      )
      interpolated_values = interpolator.interpolate(times, strikes)
  
      with self.subTest("CorrectDtype"):
          self.assertEqual(
              interpolated_values.dtype.as_numpy_dtype, dtype
          )
      with self.subTest("CorrectShape"):
          self.assertAllClose(
              interpolated_values.shape.as_list(), [2, 3]
          )
      expected_vols = np.array(
          [
              [0.38239871, 0.71669375, 1.125],
              [0.40785739, 0.85479298, 1.00384615],
          ]
      )
      with self.subTest("CorrectResult"):
          self.assertAllClose(
              interpolated_values, expected_vols, rtol=1e-04, atol=1e-04
          )
  
file_path: tf-quant-finance/tf_quant_finance/math/interpolation/interpolation_2d/interpolation_2d_test.py
imports: import numpy
  import tensorflow.compat.v1
  import tf_quant_finance
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 2 for calling the Interpolation2DTest obj's underlying member method code with example usage: obj.test_batch() without return types