code: def trace_jacobian_exact(ode_fn, state_shape, dtype):
      """Generates a function that computes `ode_fn` and trace of the jacobian.
  
      Augments provided `ode_fn` with explicit computation of the trace of the
      jacobian. This approach scales quadratically with the number of dimensions.
      This method computes unreduced trace, as reduction is performed inside of the
      bijector class.
  
      Args:
        ode_fn: `Callable(time, state)` that computes time derivative.
        state_shape: `TensorShape` representing the shape of the state.
        dtype: ``tf.DType` object representing the dtype of `state` tensor.
  
      Returns:
        augmented_ode_fn: `Callable(time, (state, log_det_jac))` that computes
          augmented time derivative `(state_time_derivative, trace_estimation)`.
      """
      del state_shape, dtype  # Not used by trace_jacobian_exact
  
      def augmented_ode_fn(time, state_log_det_jac, **kwargs):
          """Computes both time derivative and trace of the jacobian."""
          state, _ = state_log_det_jac
          ode_fn_with_time = lambda x: ode_fn(time, x, **kwargs)
          batch_shape = [ps.size0(state)]
          state_time_derivative, diag_jac = diag_jacobian(
              xs=state, fn=ode_fn_with_time, sample_shape=batch_shape
          )
          # tfp_math.diag_jacobian returns lists
          if isinstance(state_time_derivative, list):
              state_time_derivative = state_time_derivative[0]
          if isinstance(diag_jac, list):
              diag_jac = diag_jac[0]
          trace_value = diag_jac
          return state_time_derivative, trace_value
  
      return augmented_ode_fn
  
file_path: probability/tensorflow_probability/python/bijectors/ffjord.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.internal import cache_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math.diag_jacobian import diag_jacobian
  from tensorflow_probability.python.math.ode import dormand_prince
label_desc: Use image node_id 2 for calling a global function with example usage: trace_jacobian_exact(ode_fn, state_shape, dtype) and returns: augmented_ode_fn, state_time_derivative,  trace_value