code: def construct_drift_data(drift, total_drift_fn, dim, dtype):
      """Constructs drift functions."""
      # Six cases based on drift being (None, callable, constant) and total drift
      # being (None, callable).
      # 1. Neither drift nor total drift given -> return defaults.
      # 2. total drift is none and drift is a callable -> raise error.
      # 3. total drift is none and drift is a constant -> calculate.
      # 4. total drift is given and drift is none -> differentiate.
      # 5. total drift is given and drift is callable -> return them.
      # 6. total drift is given and drift is constant -> wrap and return.
  
      # Case 1, 2 and 3.
      if total_drift_fn is None:
          # Case 1
          if drift is None:
              return _default_drift_data(dim, dtype)
          if is_callable(drift):
              # Case 2: drift is a function and total_drift_fn needs to be computed.
              # We need numerical integration for this which will be added later.
              # For now, return None.
              # TODO(b/141091950): Add numerical integration.
              return drift, None
  
          # Case 3. Drift is a constant.
          def total_drift(t1, t2):
              dt = tf.convert_to_tensor(t2 - t1, dtype=dtype)
              return outer_multiply(
                  dt, tf.ones([dim], dtype=dt.dtype) * drift
              )
  
          return (
              _make_drift_fn_from_const(drift, dim, dtype),
              total_drift,
          )
  
      # Total drift is not None
      # Case 4.
      if drift is None:
  
          def drift_from_total_drift(t):
              start_time = tf.zeros_like(t)
              return gradient.fwd_gradient(
                  lambda x: total_drift_fn(start_time, x), t
              )
  
          return drift_from_total_drift, total_drift_fn
  
      # Case 5
      if is_callable(drift):
          return drift, total_drift_fn
  
      return (
          _make_drift_fn_from_const(drift, dim, dtype),
          total_drift_fn,
      )
  
file_path: tf-quant-finance/tf_quant_finance/models/legacy/brownian_motion_utils.py
imports: import tensorflow.compat.v2
  from tf_quant_finance.math import gradient
label_desc: Use image node_id 3 for calling a global function with example usage: construct_drift_data(drift, total_drift_fn, dim, dtype) and returns: _make_drift_fn_from_const,  total_drift_fn, _make_drift_fn_from_const,  total_drift, drift_from_total_drift,  total_drift_fn, drift,  total_drift_fn, _default_drift_data, drift,  None, outer_multiply, gradient