code: def _maybe_call_volatility_fn_and_grads(
      volatility_fn,
      state,
      volatility_fn_results=None,
      grads_volatility_fn=None,
      sample_shape=None,
      parallel_iterations=10,
  ):
      """Helper which computes `volatility_fn` results and grads, if needed."""
      state_parts = (
          list(state) if mcmc_util.is_list_like(state) else [state]
      )
      needs_volatility_fn_gradients = grads_volatility_fn is None
  
      # Convert `volatility_fn_results` to a list
      if volatility_fn_results is None:
          volatility_fn_results = volatility_fn(*state_parts)
  
      volatility_fn_results = (
          list(volatility_fn_results)
          if mcmc_util.is_list_like(volatility_fn_results)
          else [volatility_fn_results]
      )
      if len(volatility_fn_results) == 1:
          volatility_fn_results *= len(state_parts)
      if len(state_parts) != len(volatility_fn_results):
          raise ValueError(
              "`volatility_fn` should return a tensor or a list "
              "of the same length as `current_state`."
          )
  
      # The shape of 'volatility_parts' needs to have the number of chains as a
      # leading dimension. For determinism we broadcast 'volatility_parts' to the
      # shape of `state_parts` since each dimension of `state_parts` could have a
      # different volatility value.
  
      volatility_fn_results = _maybe_broadcast_volatility(
          volatility_fn_results, state_parts
      )
      if grads_volatility_fn is None:
          [
              _,
              grads_volatility_fn,
          ] = diag_jacobian(
              xs=state_parts,
              ys=volatility_fn_results,
              sample_shape=sample_shape,
              parallel_iterations=parallel_iterations,
              fn=volatility_fn,
          )
  
      # Compute gradient of `volatility_parts**2`
      if needs_volatility_fn_gradients:
          grads_volatility_fn = [
              2.0 * g * volatility
              if g is not None
              else tf.zeros_like(fn_arg)
              for g, volatility, fn_arg in zip(
                  grads_volatility_fn,
                  volatility_fn_results,
                  state_parts,
              )
          ]
  
      return volatility_fn_results, grads_volatility_fn
  
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 4 for calling a global function with example usage: _maybe_call_volatility_fn_and_grads(volatility_fn, state, volatility_fn_results, grads_volatility_fn, sample_shape, parallel_iterations) and returns: volatility_fn_results,  grads_volatility_fn