code: def _create_test_agent(self):
      # This dummy network allows us to deterministically anticipate that the
      # state-action-quantile outputs will be equal to sum of the
      # corresponding quantile inputs.
      # State/Quantile shapes will be k x 1, (N x batch_size) x 1,
      # or (N' x batch_size) x 1.
  
      class MockImplicitQuantileNetwork(linen.Module):
          """Custom Jax model used in tests."""
  
          num_actions: int
          quantile_embedding_dim: int
          inputs_preprocessed: bool = False
  
          @linen.compact
          def __call__(self, x, num_quantiles, rng):
              del rng
              x = x.reshape((-1))  # flatten
              state_net_tiled = jnp.tile(x, [num_quantiles, 1])
              x *= state_net_tiled
              quantile_values = linen.Dense(
                  features=self.num_actions,
                  kernel_init=linen.initializers.ones,
                  bias_init=linen.initializers.zeros,
              )(x)
              quantiles = jnp.ones([num_quantiles, 1])
              return atari_lib.ImplicitQuantileNetworkType(
                  quantile_values, quantiles
              )
  
      agent = implicit_quantile_agent.JaxImplicitQuantileAgent(
          network=MockImplicitQuantileNetwork,
          num_actions=self._num_actions,
          kappa=1.0,
          num_tau_samples=2,
          num_tau_prime_samples=3,
          num_quantile_samples=4,
          epsilon_eval=0.0,
      )
      # This ensures non-random action choices (since epsilon_eval = 0.0) and
      # skips the train_step.
      agent.eval_mode = True
      return agent
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import absltest
  from dopamine.discrete_domains import atari_lib
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.implicit_quantile import implicit_quantile_agent
  from flax import linen
  import gin
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the ImplicitQuantileAgentTest obj's underlying member method code with example usage: obj._create_test_agent() and returns: agent, atari_lib