code: def testShapesAndValues(self):
      agent = self._create_test_agent()
      state = onp.ones((1, 28224))
      net_outputs = agent.network_def.apply(agent.online_params, state)
      self.assertEqual(
          net_outputs.logits.shape, (self.num_actions, self._num_atoms)
      )
      self.assertEqual(
          net_outputs.probabilities.shape, net_outputs.logits.shape
      )
      self.assertEqual(net_outputs.q_values.shape, (self.num_actions,))
      # Check probabilities are uniform for all except the first action.
      expected_probabilities = (
          onp.ones_like(net_outputs.probabilities)
          * 1.0
          / self._num_atoms
      )
      # The first action will have probability 1.0 at the highest quantile (it is
      # set this way by design in the mock network to guarantee that action 1 is
      # always selected).
      expected_probabilities[0] = onp.zeros(self._num_atoms)
      expected_probabilities[0, -1] = 1.0
      onp.testing.assert_allclose(
          net_outputs.probabilities, expected_probabilities
      )
  
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.quantile import quantile_agent
  from dopamine.utils import test_utils
  from flax import linen
  import gin
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling the JaxQuantileAgentTest obj's underlying member method code with example usage: obj.testShapesAndValues() without return types