docstr: Tests CIR drift and volatility functions.
docstr_tok: ["Tests","CIR","drift","and","volatility","functions","."]
code: def test_drift_and_volatility_batch_shape(self):
      """Tests CIR drift and volatility functions."""
      dim = 1
      # model_shape = [3, 2]
      theta = [[[0.04], [0.9]], [[0.03], [0.8]], [[0.05], [0.7]]]
      mean_reversion = [
          [[0.06], [0.01]],
          [[0.05], [0.02]],
          [[0.07], [0.03]],
      ]
      sigma = [[[0.1], [0.11]], [[0.2], [0.12]], [[0.3], [0.13]]]
      # model_shape + whatever_shape + [dim]
      state_shape = (3, 2, 10, 8, 7, dim)
      state = np.full(shape=state_shape, fill_value=2.0)
  
      process = tff.models.cir.CirModel(
          theta=theta,
          mean_reversion=mean_reversion,
          sigma=sigma,
          dtype=np.float64,
      )
      drift_fn = process.drift_fn()
      volatility_fn = process.volatility_fn()
  
      with self.subTest("Drift"):
          drift = drift_fn(0.2, state)
          self.assertEqual(state_shape, drift.shape)
  
          # Shape [3, 2, 1, 1, 1, 1]
          theta_expand = np.array(
              [
                  [[[[[0.04]]]], [[[[0.9]]]]],
                  [[[[[0.03]]]], [[[[0.8]]]]],
                  [[[[[0.05]]]], [[[[0.7]]]]],
              ]
          )
          mean_reversion_expand = np.array(
              [
                  [[[[[0.06]]]], [[[[0.01]]]]],
                  [[[[[0.05]]]], [[[[0.02]]]]],
                  [[[[[0.07]]]], [[[[0.03]]]]],
              ]
          )
          # Shape [3, 2, 10, 8, 7, 1]
          expected_drift = theta_expand - mean_reversion_expand * state
          self.assertAllClose(
              expected_drift, drift, atol=1e-7, rtol=1e-7
          )
      with self.subTest("Volatility"):
          vol = volatility_fn(0.2, state)
          self.assertEqual(state_shape + (dim,), vol.shape)
  
          sigma_expand = np.array(
              [
                  [[[[[0.1]]]], [[[[0.11]]]]],
                  [[[[[0.2]]]], [[[[0.12]]]]],
                  [[[[[0.3]]]], [[[[0.13]]]]],
              ]
          )
          expected_vol = np.expand_dims(
              sigma_expand * np.sqrt(state), axis=-1
          )
          self.assertAllClose(expected_vol, vol, atol=1e-7, rtol=1e-7)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
  from tf_quant_finance.math import random_ops
label_desc: Use image node_id 2 for calling the CirTest obj's underlying member method code with example usage: obj.test_drift_and_volatility_batch_shape() without return types