code: def target_quantile_values(
      network_def,
      online_params,
      target_params,
      next_states,
      rewards,
      terminals,
      num_tau_prime_samples,
      num_quantile_samples,
      cumulative_gamma,
      double_dqn,
      rng,
  ):
      """Build the target for return values at given quantiles.
  
      Args:
        network_def: Linen Module used for inference.
        online_params: Parameters used for the online network.
        target_params: Parameters 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.
        num_tau_prime_samples: int, number of tau' samples (static_argnum).
        num_quantile_samples: int, number of quantile samples (static_argnum).
        cumulative_gamma: float, cumulative gamma to use (static_argnum).
        double_dqn: bool, whether to use double DQN (static_argnum).
        rng: Jax random number generator.
  
      Returns:
        Jax random number generator.
        The target quantile values.
      """
      rewards = jnp.tile(rewards, [num_tau_prime_samples])
      is_terminal_multiplier = 1.0 - terminals.astype(jnp.float32)
      # Incorporate terminal state to discount factor.
      gamma_with_terminal = cumulative_gamma * is_terminal_multiplier
      gamma_with_terminal = jnp.tile(
          gamma_with_terminal, [num_tau_prime_samples]
      )
      rng, rng1, rng2 = jax.random.split(rng, num=3)
      # Compute Q-values which are used for action selection for the next states
      # in the replay buffer. Compute the argmax over the Q-values.
      if double_dqn:
          outputs_action = network_def.apply(
              online_params,
              next_states,
              num_quantiles=num_quantile_samples,
              rng=rng1,
          )
      else:
          outputs_action = network_def.apply(
              target_params,
              next_states,
              num_quantiles=num_quantile_samples,
              rng=rng1,
          )
      target_quantile_values_action = outputs_action.quantile_values
      target_q_values = jnp.squeeze(
          jnp.mean(target_quantile_values_action, axis=0)
      )
      # Shape: batch_size.
      next_qt_argmax = jnp.argmax(target_q_values)
      # Get the indices of the maximium Q-value across the action dimension.
      # Shape of next_qt_argmax: (num_tau_prime_samples x batch_size).
      next_state_target_outputs = network_def.apply(
          target_params,
          next_states,
          num_quantiles=num_tau_prime_samples,
          rng=rng2,
      )
      next_qt_argmax = jnp.tile(next_qt_argmax, [num_tau_prime_samples])
      target_quantile_vals = jax.vmap(lambda x, y: x[y])(
          next_state_target_outputs.quantile_values, next_qt_argmax
      )
      target_quantile_vals = (
          rewards + gamma_with_terminal * target_quantile_vals
      )
      # We return with an extra dimension, which is expected by train.
      return rng, jax.lax.stop_gradient(target_quantile_vals[:, None])
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  from dopamine.jax import networks
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.metrics import statistics_instance
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: target_quantile_values(network_def, online_params, target_params, next_states, rewards, terminals, num_tau_prime_samples, num_quantile_samples, cumulative_gamma, double_dqn, rng) and returns: rng,  jax