class_docstr: A Jax implementation of the standard 3-layer CNN used in Atari.
  
  Attributes:
    padding: which padding style to use. Defaults to SAME, which yields larger
      final latents.
code: def __call__(self, x):
      # x = x[None, Ellipsis]
      hidden_sizes = self.stack_sizes
      kernel_sizes = [8, 4, 3]
      stride_sizes = [4, 2, 1]
      for layer in range(3):
          x = nn.Conv(
              features=hidden_sizes[layer],
              kernel_size=(kernel_sizes[layer], kernel_sizes[layer]),
              strides=(stride_sizes[layer], stride_sizes[layer]),
              kernel_init=nn.initializers.xavier_uniform(),
              padding=self.padding,
          )(x)
          x = nn.relu(x)  # flatten
      return x
  
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 1 for calling the RainbowCNN obj's underlying member method code with example usage: obj.__call__(x) and returns: x