code: def AtariCnn(
      n_frames=4, hidden_sizes=(32, 32), output_size=128, mode="train"
  ):
      """An Atari CNN."""
      del mode
  
      # TODO(jonni): Include link to paper?
      # Input shape: (B, T, H, W, C)
      # Output shape: (B, T, output_size)
      return tl.Serial(
          _BytesToFloats(),
          _FrameStack(n_frames=n_frames),  # (B, T, H, W, 4C)
          tl.Conv(hidden_sizes[0], (5, 5), (2, 2), "SAME"),
          tl.Relu(),
          tl.Conv(hidden_sizes[1], (5, 5), (2, 2), "SAME"),
          tl.Relu(),
          tl.Flatten(n_axes_to_keep=2),  # B, T and rest.
          tl.Dense(output_size),
          tl.Relu(),
      )
  
imports: from trax import layers
label_desc: Use image node_id 3 for calling a global function with example usage: AtariCnn(n_frames, hidden_sizes, output_size, mode) and returns: tl