code: def testFlattenUnflatten(self, dist_name, data):
      if (dist_name in PYTREE_BLOCKLIST) != FLAGS.blocklists_only:
          self.skipTest("Distribution currently broken.")
  
      dist = data.draw(
          dhps.distributions(
              enable_vars=False,
              dist_name=dist_name,
              validate_args=False,
              eligibility_filter=lambda dname: dname
              not in PYTREE_BLOCKLIST,
          )
      )
      flat_dist, dist_tree = jax.tree_util.tree_flatten(dist)
      new_dist = jax.tree_util.tree_unflatten(dist_tree, flat_dist)
      for old_param, new_param in zip(
          flat_dist, jax.tree_util.tree_leaves(new_dist)
      ):
          self.assertEqual(type(old_param), type(new_param))
          if isinstance(old_param, np.ndarray):
              self.assertTupleEqual(old_param.shape, new_param.shape)
  
imports: import functools
  import os
  from absl import flags
  from absl.testing import parameterized
  import hypothesis
  from hypothesis import strategies
  import jax
  from jax import random
  import jax.numpy
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal.backend import jax
  from tensorflow_probability.substrates.jax.bijectors import bijector_test_util
  from tensorflow_probability.substrates.jax.distributions import hypothesis_testlib
  from tensorflow_probability.substrates.jax.distributions import normal
  from tensorflow_probability.substrates.jax.distributions import transformed_distribution
  from tensorflow_probability.substrates.jax.internal import hypothesis_testlib
  from tensorflow_probability.substrates.jax.internal import test_util
label_desc: Use image node_id 1 for calling the PytreeTest obj's underlying member method code with example usage: obj.testFlattenUnflatten(dist_name, data) without return types