code: def AtariCnnBody(
      n_frames=4,
      hidden_sizes=(32, 64, 64),
      output_size=512,
      mode="train",
      kernel_initializer=None,
      padding="VALID",
  ):
      """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],
              (8, 8),
              (4, 4),
              padding=padding,
              kernel_initializer=kernel_initializer,
          ),
          tl.Relu(),
          tl.Conv(
              hidden_sizes[1],
              (4, 4),
              (2, 2),
              padding=padding,
              kernel_initializer=kernel_initializer,
          ),
          tl.Relu(),
          tl.Conv(
              hidden_sizes[2],
              (3, 3),
              (1, 1),
              padding=padding,
              kernel_initializer=kernel_initializer,
          ),
          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 4 for calling a global function with example usage: AtariCnnBody(n_frames, hidden_sizes, output_size, mode, kernel_initializer, padding) and returns: tl