code: def compute_hamiltonian(
      target_log_prob, momentum_parts, shard_axis_names=None
  ):
      """Compute the Hamiltonian of the current system."""
      shard_axis_names = shard_axis_names or (
          [None] * len(momentum_parts)
      )
      independent_chain_ndims = ps.rank(target_log_prob)
  
      def compute_sum_sq(v, shard_axes):
          sum_sq = tf.reduce_sum(
              v**2.0,
              axis=ps.range(independent_chain_ndims, ps.rank(v)),
          )
          if shard_axes is not None:
              sum_sq = distribute_lib.psum(sum_sq, shard_axes)
          return sum_sq
  
      momentum_sq_parts = (
          tf.cast(  # pylint: disable=g-complex-comprehension
              compute_sum_sq(m, axes), dtype=target_log_prob.dtype
          )
          for m, axes in zip(momentum_parts, shard_axis_names)
      )
      # TODO(jvdillon): Verify no broadcasting happening.
      return target_log_prob - 0.5 * sum(momentum_sq_parts)
  
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 6 for calling a global function with example usage: compute_hamiltonian(target_log_prob, momentum_parts, shard_axis_names) and returns: unknown, sum_sq