code: def target_distribution(
      target_network,
      next_states,
      rewards,
      terminals,
      support,
      cumulative_gamma,
  ):
      """Builds the C51 target distribution as per Bellemare et al. (2017).
  
      First, we compute the support of the Bellman target, r + gamma Z'. Where Z'
      is the support of the next state distribution:
  
        * Evenly spaced in [-vmax, vmax] if the current state is nonterminal;
        * 0 otherwise (duplicated num_atoms times).
  
      Second, we compute the next-state probabilities, corresponding to the action
      with highest expected value.
  
      Finally we project the Bellman target (support + probabilities) onto the
      original support.
  
      Args:
        target_network: Jax Module used for the target network.
        next_states: numpy array of batched next states.
        rewards: numpy array of batched rewards.
        terminals: numpy array of batched terminals.
        support: support for the distribution.
        cumulative_gamma: float, cumulative gamma to use.
  
      Returns:
        The target distribution from the replay.
      """
      is_terminal_multiplier = 1.0 - terminals.astype(jnp.float32)
      # Incorporate terminal state to discount factor.
      gamma_with_terminal = cumulative_gamma * is_terminal_multiplier
      target_support = rewards + gamma_with_terminal * support
      next_state_target_outputs = target_network(next_states)
      q_values = jnp.squeeze(next_state_target_outputs.q_values)
      next_qt_argmax = jnp.argmax(q_values)
      probabilities = jnp.squeeze(
          next_state_target_outputs.probabilities
      )
      next_probabilities = probabilities[next_qt_argmax]
      return jax.lax.stop_gradient(
          project_distribution(
              target_support, next_probabilities, support
          )
      )
  
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 2 for calling a global function with example usage: target_distribution(target_network, next_states, rewards, terminals, support, cumulative_gamma) and returns: jax