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, eval_mode=False, key=None):
      # Generate a random number generation key if not provided
      # if key is None:
      #   key = jax.random.PRNGKey(int(time.time() * 1e6))
  
      if not self.inputs_preprocessed:
          x = x.astype(jnp.float32) / 255.0
  
      hidden_sizes = [32, 64, 64]
      kernel_sizes = [8, 4, 3]
      stride_sizes = [4, 2, 1]
      for hidden_size, kernel_size, stride_size in zip(
          hidden_sizes, kernel_sizes, stride_sizes
      ):
          layer = nn.Conv(
              features=hidden_size,
              kernel_size=(kernel_size, kernel_size),
              strides=(stride_size, stride_size),
              kernel_init=nn.initializers.xavier_uniform(),
          )
          x = layer(x)
          x = nn.relu(x)
          x = self._record_activations(x, layer)
      x = x.reshape((-1))  # flatten
  
      layer = nn.Dense(
          features=512, kernel_init=nn.initializers.xavier_uniform()
      )
      x = layer(x)
      x = nn.relu(x)
      x = self._record_activations(x, layer)
  
      if self.dueling:
          layer = nn.Dense(
              features=self.num_actions * self.num_atoms,
              kernel_init=nn.initializers.xavier_uniform(),
          )
          adv = layer(x)
          adv = self._record_activations(adv, layer)
          layer = nn.Dense(
              features=self.num_atoms,
              kernel_init=nn.initializers.xavier_uniform(),
          )
          value = layer(x)
          value = self._record_activations(value, layer)
          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:
          layer = nn.Dense(
              features=self.num_actions * self.num_atoms,
              kernel_init=nn.initializers.xavier_uniform(),
          )
          x = layer(x)
          logits = x.reshape((self.num_actions, self.num_atoms))
          logits = self._record_activations(logits, layer)
  
      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 math
  from dopamine.discrete_domains import atari_lib
  from flax import linen
  import gin
  import jax.numpy
label_desc: Use image node_id 2 for calling the FullRainbowNetwork obj's underlying member method code with example usage: obj.__call__(x, support, eval_mode, key) and returns: atari_lib, atari_lib