code: def _get_drift(
      step_size_parts,
      volatility_parts,
      grads_volatility,
      grads_target_log_prob,
      name=None,
  ):
      """Compute diffusion drift at the current location `current_state`.
  
      The drift of the diffusion at is computed as
  
      ```none
      0.5 * `step_size` * volatility_parts * `target_log_prob_fn(current_state)`
      + `step_size` * `grads_volatility`
      ```
  
      where `volatility_parts` = `volatility_fn(current_state)**2` and
      `grads_volatility` is a gradient of `volatility_parts` at the `current_state`.
  
      Args:
        step_size_parts: Python `list` of `Tensor`s representing the step size for
          Euler-Maruyama method. Must broadcast with the shape of
          `volatility_parts`.  Larger step sizes lead to faster progress, but
          too-large step sizes make rejection exponentially more likely. When
          possible, it's often helpful to match per-variable step sizes to the
          standard deviations of the target distribution in each variable.
        volatility_parts: Python `list` of `Tensor`s representing the value of
          `volatility_fn(*state_parts)`.
        grads_volatility: Python list of `Tensor`s representing the value of the
          gradient of `volatility_parts**2` wrt the state of the chain.
        grads_target_log_prob: Python list of `Tensor`s representing
          gradient of `target_log_prob_fn(*state_parts`) wrt `state_parts`. Must
          have same shape as `volatility_parts`.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., 'mala_get_drift').
  
      Returns:
        drift_parts: Tensor or Python list of `Tensor`s representing the
          state(s) of the Markov chain(s) at each result step. Has same shape as
          input `current_state_parts`.
      """
  
      with tf.name_scope(name or "mala_get_drift"):
          drift_parts = []
  
          for (
              step_size,
              volatility,
              grad_volatility,
              grad_target_log_prob,
          ) in zip(
              step_size_parts,
              volatility_parts,
              grads_volatility,
              grads_target_log_prob,
          ):
              volatility_squared = tf.square(volatility)
              drift = (
                  0.5
                  * step_size
                  * (
                      volatility_squared * grad_target_log_prob
                      + grad_volatility
                  )
              )
              drift_parts.append(drift)
  
          return drift_parts
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.math.diag_jacobian import diag_jacobian
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import metropolis_hastings
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 2 for calling a global function with example usage: _get_drift(step_size_parts, volatility_parts, grads_volatility, grads_target_log_prob, name) and returns: drift_parts