class_docstr: Network for Atari.
code: def __call__(self, x: dm_env.TimeStep, state):
      x = jax.tree_util.tree_map(lambda t: t[None, ...], x)
      return self.unroll(x, state)
  
imports: import collections
  import dm_env
  import haiku
  import jax.nn
  import jax.numpy
label_desc: Use image node_id 3 for calling the AtariNet obj's underlying member method code with example usage: obj.__call__(x, state) and returns: self