code: def test_shapes(self, space, gin_config):
      gin.parse_config(gin_config)
  
      batch_shape = (2, 3)
      distribution = distributions.create_distribution(space)
      inputs = np.random.random(batch_shape + (distribution.n_inputs,))
      point = distribution.sample(inputs)
      self.assertEqual(point.shape, batch_shape + space.shape)
      # Check if the datatypes are compatible, i.e. either both floating or both
      # integral.
      self.assertEqual(
          isinstance(point.dtype, float), isinstance(space.dtype, float)
      )
      log_prob = distribution.log_prob(inputs, point)
      self.assertEqual(log_prob.shape, batch_shape)
  
file_path: trax/trax/rl/distributions_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import gin
  import gym
  import numpy
  from trax.rl import distributions
label_desc: Use image node_id 2 for calling the DistributionsTest obj's underlying member method code with example usage: obj.test_shapes(space, gin_config) without return types