code: def trace_jacobian_hutchinson(
      ode_fn,
      state_shape,
      dtype,
      sample_fn=tf.random.normal,
      num_samples=1,
      seed=None,
  ):
      """Generates a function that computes `ode_fn` and expectation of the trace.
  
      Uses Hutchinson's trick to estimate the trace of the Jacobian using automatic
      differentiation. This is the approach used in the original FFJORD paper [1].
      This method computes unreduced trace, as reduction is performed inside of the
      bijector class.
  
      The trace estimate is obtained by computing
      ```None
      Tr[A] approx_equal sum_{i} r[i]^{T} @ A @ r[i]; r[i] - gaussian sample.
      ```
  
      For details on the original work see [2].
  
      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.
        sample_fn: `Callable(shape, dtype, seed)` that generates random samples with
          zero mean and covariance of an identity matrix.
          Default value: `tf.random.normal`
        num_samples: `Integer` number of random samples to use for trace estimation.
          Default value: '1'
        seed: PRNG seed compatible with `sample_fn`.
  
      Returns:
        augmented_ode_fn: `Callable(time, (state, log_det_jac))` that computes
          augmented time derivative `(state_time_derivative, trace_estimation)`.
  
      #### References
  
      [1]:  Grathwohl, W., Chen, R. T., Betterncourt, J., Sutskever, I.,
            & Duvenaud, D. (2018). Ffjord: Free-form continuous dynamics for
            scalable reversible generative models. arXiv preprint arXiv:1810.01367.
            http://arxiv.org.abs/1810.01367
  
      [2]:  Hutchinson, M. F. (1989). A stochastic estimator of the trace of the
            influence matrix for Laplacian smoothing splines. Communications in
            Statistics-Simulation and Computation, 18(3), 1059-1076.
      """
  
      random_samples = sample_fn(
          ps.concat([[num_samples], state_shape], axis=0),
          dtype=dtype,
          seed=seed,
      )
  
      def augmented_ode_fn(time, state_log_det_jac, **kwargs):
          """Computes both time derivative and trace of the jacobian."""
          state, _ = state_log_det_jac
  
          if JAX_MODE:
              import jax  # pylint:disable=g-import-not-at-top
  
              state_time_derivative, f_vjp = jax.vjp(
                  lambda s: ode_fn(time, s, **kwargs), state
              )
  
              def estimate_trace(random_sample):
                  (jvp,) = f_vjp(random_sample)
                  return random_sample * jvp
  
          else:
              with tf.GradientTape(
                  persistent=True, watch_accessed_variables=False
              ) as tape:
                  tape.watch(state)
                  state_time_derivative = ode_fn(time, state, **kwargs)
  
              def estimate_trace(random_sample):
                  """Computes stochastic trace estimate based on a single sample."""
                  #  We use first use gradient with `output_gradients` to compute the
                  #  jacobian-value-product and then take a dot product with the random
                  #  sample to obtain the trace estimate as formula above.
                  jvp = tape.gradient(
                      state_time_derivative, state, random_sample
                  )
                  return random_sample * jvp
  
          results = tf.vectorized_map(estimate_trace, random_samples)
          trace_estimates = tf.reduce_mean(results, axis=0)
          return state_time_derivative, trace_estimates
  
      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 1 for calling a global function with example usage: trace_jacobian_hutchinson(ode_fn, state_shape, dtype, sample_fn, num_samples, seed) and returns: augmented_ode_fn, state_time_derivative,  trace_estimates, unknown, unknown