code: def test_construct_drift_callable(self):
      dtype = tf.float64
      a, b = 0.1, -0.8
  
      def test_drift_fn(t):
          return tf.expand_dims(t * a + b, axis=-1)
  
      def test_total_drift_fn(t1, t2):
          res = (t2**2 - t1**2) * a / 2 + (t2 - t1) * b
          return tf.expand_dims(res, axis=-1)
  
      drift_fn, total_drift_fn = bm_utils.construct_drift_data(
          test_drift_fn, test_total_drift_fn, 1, dtype
      )
      times = tf.constant([0.0, 1.0, 2.0], dtype=dtype)
      drift_vals = self.evaluate(drift_fn(times))
      np.testing.assert_array_equal(drift_vals.shape, [3, 1])
      np.testing.assert_allclose(drift_vals, [[-0.8], [-0.7], [-0.6]])
  
      t1 = tf.constant([1.0, 2.0, 3.0], dtype=dtype)
      t2 = tf.constant([1.5, 3.0, 5.0], dtype=dtype)
  
      total_vals = self.evaluate(total_drift_fn(t1, t2))
      np.testing.assert_array_equal(total_vals.shape, [3, 1])
      np.testing.assert_allclose(
          total_vals, [[-0.3375], [-0.55], [-0.8]], atol=1e-7
      )
      # Tests that total drift is None if drift is a callable and no total_drift
      # is supplied
  
      _, total_drift = bm_utils.construct_drift_data(
          test_drift_fn, None, 1, dtype
      )
      self.assertIsNone(total_drift)
  
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 5 for calling the BrownianMotionUtilsTest obj's underlying member method code with example usage: obj.test_construct_drift_callable() and returns: tf, tf