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_drift_batched(
      self, supply_draws, dtype
  ):
      """Tests the mean and vol of the univariate GBM sampled paths."""
      initial_state = 2.0
      min_tol = _tolerance_by_dtype(dtype)
  
      # Batch dimensions = [2].
      mu_times = np.array(
          [[0.0, 5.0, 10.0], [0.0, 7.0, 10.0]], dtype=dtype
      )
      mu_values = np.array(
          [[0.0, 0.0, 0.05, 0.05], [0.01, 0.01, 0.07, 0.07]],
          dtype=dtype,
      )
  
      mu = tff.math.piecewise.PiecewiseConstantFunc(
          jump_locations=mu_times, values=mu_values, dtype=dtype
      )
      sigma = 0.0
      times = np.array(
          [[0.0, 1.0, 5.0, 7.0, 10.0], [0.0, 1.5, 3.2, 4.8, 25.3]],
          dtype=dtype,
      )
      mean, var, se_mean, _ = calculate_sample_paths_mean_and_variance(
          self,
          mu,
          sigma,
          times,
          initial_state,
          supply_draws,
          NUM_SAMPLES,
          dtype,
      )
      expected_mean = np.array(
          [
              [
                  0.0,  # mu = 0 at t = 0
                  0.0,  # mu = 0 for t <= 1.0
                  0.0,  # mu = 0 for t < 5.0
                  2.0 * 0.05,  # mu = 0.05 for 5.0 < t <= 7.0
                  5.0 * 0.05,  # mu = 0.05 for 5.0 < t <= 10.0
              ],
              [
                  0.0,  # mu = 0.01 at t = 0
                  1.5 * 0.01,  # mu = 0.01 for t <= 1.5
                  3.2 * 0.01,  # mu = 0.01 for t < 3.2
                  4.8 * 0.01,  # mu = 0.01 t <= 4.8
                  7.0 * 0.01
                  + 18.3 * 0.07  # mu = 0.01 for t <= 7.0 and.
                  # 0.07 after
              ],
          ],
          dtype=dtype,
      ) + np.log(initial_state)
      expected_var = np.zeros((2, 5), dtype=dtype)  # As sigma is zero.
  
      # Set acceptable tolerances based on the predicted variance and a minimum
      # tolerance based on the precision.
      mean_tol = np.maximum(se_mean * NUM_STDERRS, min_tol)
      var_tol = np.ones_like(expected_var) * min_tol
  
      arrays_all_close(
          self,
          tf.squeeze(mean),
          expected_mean,
          atol=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 13 for calling the GeometricBrownianMotionTest obj's underlying member method code with example usage: obj.test_univariate_time_varying_drift_batched(supply_draws, dtype) without return types