code: def project_distribution(supports, weights, target_support):
      """Projects a batch of (support, weights) onto target_support.
  
      Based on equation (7) in (Bellemare et al., 2017):
        https://arxiv.org/abs/1707.06887
      In the rest of the comments we will refer to this equation simply as Eq7.
  
      Args:
        supports: Jax array of shape (num_dims) defining supports for
          the distribution.
        weights: Jax array of shape (num_dims) defining weights on the
          original support points. Although for the CategoricalDQN agent these
          weights are probabilities, it is not required that they are.
        target_support: Jax array of shape (num_dims) defining support of the
          projected distribution. The values must be monotonically increasing. Vmin
          and Vmax will be inferred from the first and last elements of this Jax
          array, respectively. The values in this Jax array must be equally spaced.
  
      Returns:
        A Jax array of shape (num_dims) with the projection of a batch
        of (support, weights) onto target_support.
  
      Raises:
        ValueError: If target_support has no dimensions, or if shapes of supports,
          weights, and target_support are incompatible.
      """
      v_min, v_max = target_support[0], target_support[-1]
      # `N` in Eq7.
      num_dims = target_support.shape[0]
      # delta_z = `\Delta z` in Eq7.
      delta_z = (v_max - v_min) / (num_dims - 1)
      # clipped_support = `[\hat{T}_{z_j}]^{V_max}_{V_min}` in Eq7.
      clipped_support = jnp.clip(supports, v_min, v_max)
      # numerator = `|clipped_support - z_i|` in Eq7.
      numerator = jnp.abs(clipped_support - target_support[:, None])
      quotient = 1 - (numerator / delta_z)
      # clipped_quotient = `[1 - numerator / (\Delta z)]_0^1` in Eq7.
      clipped_quotient = jnp.clip(quotient, 0, 1)
      # inner_prod = `\sum_{j=0}^{N-1} clipped_quotient * p_j(x', \pi(x'))` in Eq7.
      inner_prod = clipped_quotient * weights
      return jnp.squeeze(jnp.sum(inner_prod, -1))
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  from dopamine.jax import losses
  from dopamine.jax import networks
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.metrics import statistics_instance
  from dopamine.replay_memory import prioritized_replay_buffer
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: project_distribution(supports, weights, target_support) and returns: jnp