class_docstr: Simple convolutional stack commonly used for Atari.
code: def __call__(self, inputs: Images) -> jnp.ndarray:
      inputs_rank = jnp.ndim(inputs)
      batched_inputs = inputs_rank == 4
      if inputs_rank < 3 or inputs_rank > 4:
          raise ValueError(
              "Expected input BHWC or HWC. Got rank %d" % inputs_rank
          )
  
      outputs = self._network(inputs)
  
      if batched_inputs:
          return jnp.reshape(outputs, [outputs.shape[0], -1])  # [B, D]
      return jnp.reshape(outputs, [-1])
  
imports: from typing import Any, Optional, Sequence, Tuple
  from acme.jax.networks import base
  from acme.jax.networks import duelling
  from acme.jax.networks import embedding
  from acme.jax.networks import policy_value
  from acme.jax.networks import resnet
  from acme.wrappers import observation_action_reward
  import haiku
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the AtariTorso obj's underlying member method code with example usage: obj.__call__(inputs) and returns: jnp, jnp