class_docstr: Jax Rainbow network for Minatar.
class_docstr_tok: ["Jax","Rainbow","network","for","Minatar","."]
code: def __call__(
      self, x: jnp.ndarray, support: jnp.ndarray
  ) -> jnp.ndarray:
      initializer = nn.initializers.xavier_uniform()
      x = x.astype(jnp.float32)
      x = nn.Conv(
          features=16,
          kernel_size=(3, 3),
          strides=(1, 1),
          kernel_init=initializer,
      )(x)
      x = nn.relu(x)
      x = x.reshape(-1)  # flatten
      x = nn.Dense(
          features=self.num_actions * self.num_atoms,
          kernel_init=initializer,
      )(x)
      logits = x.reshape((self.num_actions, self.num_atoms))
      probabilities = nn.softmax(logits)
      q_values = jnp.sum(support * probabilities, axis=1)
      return atari_lib.RainbowNetworkType(
          q_values, logits, probabilities
      )
  
imports: import time
  from dopamine.discrete_domains import atari_lib
  from flax import linen
  import gin
  import jax
  import jax.numpy
  import minatar
label_desc: Use image node_id 1 for calling the MinatarRainbowNetwork obj's underlying member method code with example usage: obj.__call__(x, support) and returns: atari_lib