code: def test_discrete_valued_head(
      self,
      input_shape,
      vmin_shape,
      vmax_shape,
      num_atoms,
      expected_logits_shape,
  ):
      vmin = np.zeros(vmin_shape, float)
      vmax = np.ones(vmax_shape, float)
      head = distributional.DiscreteValuedHead(
          vmin=vmin, vmax=vmax, num_atoms=num_atoms
      )
      input_array = np.zeros(input_shape, dtype=float)
      output_distribution = head(input_array)
      self.assertEqual(
          output_distribution.logits_parameter().shape,
          expected_logits_shape,
      )
  
      values = output_distribution._values
  
      # Can't do assert_allclose(values[..., 0], vmin), because the args may
      # have broadcast-compatible but unequal shapes. Do the following instead:
      npt.assert_allclose(
          values[..., 0] - vmin, np.zeros_like(values[..., 0])
      )
      npt.assert_allclose(
          values[..., -1] - vmax, np.zeros_like(values[..., -1])
      )
  
      # Check that values are monotonically increasing.
      intervals = values[..., 1:] - values[..., :-1]
      npt.assert_array_less(np.zeros_like(intervals), intervals)
  
      # Check that the values are equally spaced.
      npt.assert_allclose(
          intervals[..., 1:] - intervals[..., :1],
          np.zeros_like(intervals[..., 1:]),
          atol=1e-7,
      )
  
file_path: acme/acme/tf/networks/distributional_test.py
imports: from acme.tf.networks import distributional
  import numpy
  from numpy import testing
  from absl.testing import absltest
  from absl.testing import parameterized
label_desc: Use image node_id 1 for calling the DistributionalTest obj's underlying member method code with example usage: obj.test_discrete_valued_head(input_shape, vmin_shape, vmax_shape, num_atoms, expected_logits_shape) without return types