docstr: Tests the mean and vol of the univariate GBM sampled paths.
docstr_tok: ["Tests","the","mean","and","vol","of","the","univariate","GBM","sampled","paths","."]
code: def test_univariate_time_varying_vol_batched(
      self, supply_draws, dtype
  ):
      """Tests the mean and vol of the univariate GBM sampled paths."""
      initial_state = 2.0
      min_tol = 5e-3
      mu = 0.05
      sigma_times = np.array(
          [[0.0, 5.0, 10.0], [0.0, 7.0, 10.0]], dtype=dtype
      )
      sigma_values = np.array(
          [[0.4, 0.6, 0.8, 0.4], [0.5, 0.1, 0.3, 0.1]], dtype=dtype
      )
      sigma = tff.math.piecewise.PiecewiseConstantFunc(
          jump_locations=sigma_times, values=sigma_values, dtype=dtype
      )
      times = np.array([[0.0, 1.3, 4.5, 7.5, 19.0]], dtype=dtype)
      (
          mean,
          var,
          se_mean,
          se_var,
      ) = calculate_sample_paths_mean_and_variance(
          self,
          mu,
          sigma,
          times,
          initial_state,
          supply_draws,
          NUM_SAMPLES,
          dtype,
      )
      expected_mean = np.array(
          [
              [
                  0.0,  # t = 0
                  1.3 * mu - 0.5 * 1.3 * 0.6**2,  # t = 1.3
                  4.5 * mu - 0.5 * 4.5 * 0.6**2,  # t = 4.5
                  7.5 * mu
                  - 0.5 * (5.0 * 0.6**2 + 2.5 * 0.8**2),  # t = 7.5
                  19.0 * mu
                  - 0.5
                  * (5.0 * 0.6**2 + 5.0 * 0.8**2 + 9.0 * 0.4**2),
              ],
              [
                  0.0,  # mu = 0 at t = 0
                  1.3 * mu - 0.5 * 1.3 * 0.1**2,  # t = 1.3
                  4.5 * mu - 0.5 * 4.5 * 0.1**2,  # t = 4.5
                  7.5 * mu
                  - 0.5 * (7.0 * 0.1**2 + 0.5 * 0.3**2),  # t = 7.5
                  19.0 * mu
                  - 0.5
                  * (7.0 * 0.1**2 + 3.0 * 0.3**2 + 9.0 * 0.1**2),
              ],
          ],
          dtype=dtype,
      ) + np.log(initial_state)
  
      expected_var = np.array(
          [
              [
                  0.0,  # t = 0
                  1.3 * 0.6**2,  # t = 1.3
                  4.5 * 0.6**2,  # t = 4.5
                  5.0 * 0.6**2 + 2.5 * 0.8**2,  # t = 7.5
                  5.0 * 0.6**2
                  + 5.0 * 0.8**2
                  + 9.0 * 0.4**2,  # t = 19.0
              ],
              [
                  0.0,  # t = 0
                  1.3 * 0.1**2,  # t = 1.3
                  4.5 * 0.1**2,  # t = 4.5
                  7.0 * 0.1**2 + 0.5 * 0.3**2,  # t = 7.5
                  7.0 * 0.1**2
                  + 3.0 * 0.3**2
                  + 9.0 * 0.1**2,  # t = 19.0
              ],
          ],
          dtype=dtype,
      )
  
      mean_tol = np.maximum(se_mean * NUM_STDERRS, min_tol)
      var_tol = np.maximum(se_var * NUM_STDERRS, min_tol)
  
      arrays_all_close(
          self,
          tf.squeeze(mean),
          expected_mean,
          mean_tol,
          msg="comparing means",
      )
      arrays_all_close(
          self,
          tf.squeeze(var),
          expected_var,
          var_tol,
          msg="comparing variances",
      )
  
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.models.geometric_brownian_motion import geometric_brownian_motion_test_utils
  from tf_quant_finance.models.geometric_brownian_motion import univariate_geometric_brownian_motion
label_desc: Use image node_id 15 for calling the GeometricBrownianMotionTest obj's underlying member method code with example usage: obj.test_univariate_time_varying_vol_batched(supply_draws, dtype) without return types