class_docstr: Impala Network which also outputs penultimate representation layers.
class_docstr_tok: ["Impala","Network","which","also","outputs","penultimate","representation","layers","."]
code: def __call__(self, x, stop_grad_representation=False):
      # This method sets up the MLP for inference, using the specified number of
      # layers and units.
      logging.info(
          (
              "Creating Impala network with %d nn_scale, %s stack_sizes, %d "
              "num_blocks, normalize_representation %s"
          ),
          self.nn_scale,
          str(self.stack_sizes),
          self.num_blocks,
          self.normalize_representation,
      )
      initializer = nn.initializers.xavier_uniform()
      if not self.inputs_preprocessed:
          x = preprocess_atari_inputs(x)
  
      conv_out = self.encoder(x)
      conv_out = conv_out.reshape(-1)
  
      conv_out = nn.Dense(features=512, kernel_init=initializer)(
          conv_out
      )
      representation = nn.relu(conv_out)
      if self.normalize_representation:
          representation /= jnp.linalg.norm(representation, ord=2)
      if stop_grad_representation:
          representation = jax.lax.stop_gradient(representation)
      q_values = nn.Dense(
          features=self.num_actions, kernel_init=initializer
      )(representation)
      return NetworkType(q_values, representation)
  
imports: import collections
  from typing import Tuple
  from absl import logging
  from dopamine.discrete_domains import atari_lib
  from flax import linen
  import gin
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the ImpalaNetworkWithRepresentations obj's underlying member method code with example usage: obj.__call__(x, stop_grad_representation) and returns: NetworkType