code: def testReplayQuantileValueShape(self):
      agent = self._create_test_agent()
      batch_size = 32
      batch_states = onp.ones(
          (batch_size,) + self.observation_shape + (self.stack_size,)
      )
      for params in [agent.online_params, agent.target_network_params]:
          agent._rng, rng_input = jax.random.split(agent._rng)
          model_output = jax.vmap(
              lambda n, p, x, y, z: n.apply(
                  p, x=x, num_quantiles=y, rng=z
              ),
              in_axes=(None, None, 0, None, None),
          )(
              agent.network_def,
              params,
              batch_states,
              agent.num_tau_samples,
              rng_input,
          )
          quantile_values = model_output.quantile_values
          quantile_values = jnp.squeeze(quantile_values)
          self.assertEqual(quantile_values.shape[0], batch_size)
          self.assertEqual(
              quantile_values.shape[1], agent.num_tau_samples
          )
          self.assertEqual(quantile_values.shape[2], agent.num_actions)
  
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 6 for calling the ImplicitQuantileAgentTest obj's underlying member method code with example usage: obj.testReplayQuantileValueShape() without return types