class_docstr: Jax Rainbow network for Full Rainbow.
  
  Attributes:
    num_actions: int, number of actions the agent can take at any state.
    num_atoms: int, the number of buckets of the value function distribution.
    noisy: bool, Whether to use noisy networks.
    dueling: bool, Whether to use dueling network architecture.
    distributional: bool, whether to use distributional RL.
class_docstr_tok: ["Jax","Rainbow","network","for","Full","Rainbow",".","Attributes",":","num_actions",":","int",",","number","of","actions","the","agent","can","take","at","any","state",".","num_atoms",":","int",",","the","number","of","buckets","of","the","value","function","distribution",".","noisy",":","bool",",","Whether","to","use","noisy","networks",".","dueling",":","bool",",","Whether","to","use","dueling","network","architecture",".","distributional",":","bool",",","whether","to","use","distributional","RL","."]
code: def __call__(
      self,
      x,
      support,
      key=None,
      return_to_condition=None,
      eval_mode=False,
  ):
      del key, eval_mode
      if not self.inputs_preprocessed:
          x = preprocess_atari_inputs(x)
      if return_to_condition is not None:
          x = transform_and_concat_return(x, return_to_condition)
  
      initializer = nn.initializers.xavier_uniform()
      x = self.encoder(x)
      x = x.reshape((-1))  # flatten
      # Single hidden layer of size feature_dim
      x = nn.Dense(
          features=self.feature_dim * self.nn_scale,
          kernel_init=initializer,
      )(x)
      x = nn.relu(x)
  
      if self.dueling:
          adv = nn.Dense(features=self.num_actions * self.num_atoms)(x)
          value = nn.Dense(features=self.num_atoms)(x)
          adv = adv.reshape((self.num_actions, self.num_atoms))
          value = value.reshape((1, self.num_atoms))
          logits = value + (
              adv - (jnp.mean(adv, axis=0, keepdims=True))
          )
      else:
          x = nn.Dense(features=self.num_actions * self.num_atoms)(x)
          logits = x.reshape((self.num_actions, self.num_atoms))
  
      if self.distributional:
          probabilities = nn.softmax(logits)
          q_values = jnp.sum(support * probabilities, axis=1)
          return atari_lib.RainbowNetworkType(
              q_values, logits, probabilities
          )
      q_values = jnp.sum(logits, axis=1)  # Sum over all the num_atoms
      return atari_lib.DQNNetworkType(q_values)
  
imports: import collections
  from typing import Tuple
  from absl import logging
  from dopamine.discrete_domains import atari_lib
  from flax import linen
  import gin
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the ParameterizedRainbowNetwork obj's underlying member method code with example usage: obj.__call__(x, support, key, return_to_condition, eval_mode) and returns: atari_lib, atari_lib