docstr: Tests the mean of a univariate GBM sampled paths has the right shape.
docstr_tok: ["Tests","the","mean","of","a","univariate","GBM","sampled","paths","has","the","right","shape","."]
code: def test_univariate_time_varying_vol_batched_time_broadcast(
      self, supply_draws, dtype
  ):
      """Tests the mean of a univariate GBM sampled paths has the right shape."""
      initial_state = 2.0
      mu = 0.05
      expected_shape = (2, 3, 4, 5)
      sigma_times = np.array(
          [0.0, 5.0, 10.0, 15.0, 20.0], dtype=dtype
      ) * np.ones(expected_shape, dtype=dtype)
      sigma_values = np.array(
          [0.2, 0.2, 0.4, 0.4, 0.3, 0.1], dtype=dtype
      ) * np.ones((2, 3, 4, 6), dtype=dtype)
      sigma = tff.math.piecewise.PiecewiseConstantFunc(
          jump_locations=sigma_times, values=sigma_values, dtype=dtype
      )
      times = np.array(
          [0.0, 1.0, 5.0, 7.0, 12.0], dtype=dtype
      ) * np.ones((4, 5), dtype=dtype)
      mean, var, _, _ = calculate_sample_paths_mean_and_variance(
          self,
          mu,
          sigma,
          times,
          initial_state,
          supply_draws,
          NUM_SAMPLES,
          dtype,
      )
  
      with self.subTest("Means"):
          self.assertAllClose(mean.shape, expected_shape, atol=1e-3)
      with self.subTest("Variances"):
          self.assertAllClose(var.shape, expected_shape, atol=1e-3)
  
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 17 for calling the GeometricBrownianMotionTest obj's underlying member method code with example usage: obj.test_univariate_time_varying_vol_batched_time_broadcast(supply_draws, dtype) without return types