class_docstr: Jax Rainbow network for Full Rainbow.
  
  Attributes:
      num_actions: The number of actions the agent can take at any state.
      num_atoms: The number of buckets of the value function distribution.
      noisy: Whether to use noisy networks.
      dueling: Whether to use dueling network architecture.
      distributional: Whether to use distributional RL.
code: def __call__(
      self,
      x,
      support,
      actions=None,
      do_rollout=False,
      eval_mode=False,
      key=None,
  ):
      if not self.inputs_preprocessed:
          x = preprocess_atari_inputs(x)
  
      # Generate a random number generation key if not provided
      if key is None:
          key = random.PRNGKey(int(time.time() * 1e6))
  
      latent = self.encode(x)
      x = self.apply_projection(
          latent.reshape(-1), key, eval_mode
      )  # Single hidden layer of size 512
      x = nn.relu(x)
  
      if self.dueling:
          key, rng1, rng2 = random.split(key, 3)
          _, adv_net = feature_layer(
              self.noisy, self.num_actions * self.num_atoms
          )
          _, val_net = feature_layer(self.noisy, self.num_atoms)
          adv = adv_net(x, rng1, eval_mode)
          value = val_net(x, rng2, eval_mode)
          adv = adv.reshape((self.num_actions, self.num_atoms))
          value = value.reshape((1, self.num_atoms))
          logits = value + (adv - (jnp.mean(adv, -2, keepdims=True)))
      else:
          key, rng1 = random.split(key, 2)
          _, adv_net = feature_layer(
              self.noisy, self.num_actions * self.num_atoms
          )
          x = adv_net(x, rng1, eval_mode)
          logits = x.reshape((self.num_actions, self.num_atoms))
  
      if do_rollout:
          latent = self.spr_rollout(latent, actions, key)
  
      if self.distributional:
          probabilities = jnp.squeeze(nn.softmax(logits))
          q_values = jnp.squeeze(
              jnp.sum(support * probabilities, axis=-1)
          )
          return SPROutputType(q_values, logits, probabilities, latent)
  
      q_values = jnp.squeeze(logits)
      return SPROutputType(q_values, None, None, latent)
  
imports: import collections
  import functools
  import time
  from typing import Any, Callable, Optional, Tuple
  from dopamine.jax.networks import preprocess_atari_inputs
  from flax import linen
  import gin
  import jax
  from jax import lax
  from jax import random
  from jax.nn import initializers
  import jax.numpy
  import numpy
label_desc: Use image node_id 6 for calling the SPRNetwork obj's underlying member method code with example usage: obj.__call__(x, support, actions, do_rollout, eval_mode, key) and returns: SPROutputType, SPROutputType