code: def test_get_distribution_params(self, legacy_distribution_network):
      ones = tf.Variable(tf.ones(shape=[2], dtype=tf.float32))
      distribution = (
          tfp.distributions.Categorical(logits=ones),
          tfp.distributions.Normal(ones, ones),
      )
      params = ppo_utils.get_distribution_params(
          distribution, legacy_distribution_network
      )
  
      self.assertEqual(
          [set(["logits"]), set(["loc", "scale"])],
          [set(d.keys()) for d in params],
      )  # pytype: disable=attribute-error
      self.assertAllEqual(
          [[[2]], [[2], [2]]],
          [[d[k].shape.as_list() for k in d] for d in params],
      )
  
file_path: agents/tf_agents/agents/ppo/ppo_utils_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import tensorflow
  import tensorflow_probability
  from tf_agents.agents.ppo import ppo_utils
  from tf_agents.trajectories import time_step
label_desc: Use image node_id 3 for calling the PPOUtilsTest obj's underlying member method code with example usage: obj.test_get_distribution_params(legacy_distribution_network) without return types