code: def test_construct_drift_constant(self):
      dtypes = [tf.float64, tf.float32]
  
      def make_total_drift_fn(v, dtype):
          def fn(t1, t2):
              return bm_utils.outer_multiply(
                  t2 - t1, v * tf.ones([2], dtype=dtype)
              )
  
          return fn
  
      for dtype in dtypes:
          drift_const = tf.constant(2.0, dtype=dtype)
          drift_fn, total_drift_fn = bm_utils.construct_drift_data(
              drift_const, None, 2, dtype
          )
          times = tf.constant([0.3, 0.9, 1.5], dtype=dtype)
          drift_vals = self.evaluate(drift_fn(times))
          np.testing.assert_array_equal(drift_vals.shape, [3, 2])
          np.testing.assert_allclose(
              drift_vals, [[2.0, 2], [2, 2], [2, 2]]
          )
          total_vals = self.evaluate(total_drift_fn(times - 0.2, times))
          np.testing.assert_array_equal(total_vals.shape, [3, 2])
          np.testing.assert_allclose(
              total_vals,
              [[0.4, 0.4], [0.4, 0.4], [0.4, 0.4]],
              atol=1e-7,
          )
          # Tests if both are supplied. Note we deliberately supply inconsistent
          # data to ensure that the method doesn't do anything clever.
  
          (
              drift_fn_alt,
              total_drift_fn_alt,
          ) = bm_utils.construct_drift_data(
              drift_const, make_total_drift_fn(4.0, dtype), 2, dtype
          )
          drift_vals = self.evaluate(drift_fn_alt(times))
          np.testing.assert_array_equal(drift_vals.shape, [3, 2])
          np.testing.assert_allclose(
              drift_vals, [[2.0, 2], [2, 2], [2, 2]]
          )
          total_vals_alt = self.evaluate(
              total_drift_fn_alt(times - 0.2, times)
          )
          np.testing.assert_array_equal(total_vals.shape, [3, 2])
          np.testing.assert_allclose(
              total_vals_alt,
              [[0.8, 0.8], [0.8, 0.8], [0.8, 0.8]],
              atol=1e-5,
          )
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow.python.framework import test_util
  from tf_quant_finance.models.legacy import brownian_motion_utils
label_desc: Use image node_id 4 for calling the BrownianMotionUtilsTest obj's underlying member method code with example usage: obj.test_construct_drift_constant() and returns: fn, bm_utils