code: def test_user_defined_pytree(self):
      k = np.asarray([3])
      pytree_shift = bijector_test_util.PytreeShift(k)
      td = transformed_distribution.TransformedDistribution(
          normal.Normal(0.0, 1), bijector=pytree_shift
      )
      leaves, treedef = jax.tree_util.tree_flatten(td)
      node_data = treedef.node_data()
  
      # `td` and `td.bijector` are both Pytrees, but only `td` was registered as a
      # Pytree via AutoCompositeTensor.
      self.assertFalse(jax.tree_util.treedef_is_leaf(treedef))
      self.assertFalse(
          jax.tree_util.treedef_is_leaf(
              jax.tree_util.tree_structure(td.bijector)
          )
      )
      self.assertIsInstance(td, tf.__internal__.CompositeTensor)
      self.assertNotIsInstance(
          td.bijector, tf.__internal__.CompositeTensor
      )
  
      # `"bijector"` is in the tuple of arg names for the Pytree children and not
      # the auxiliary data.
      self.assertIn("bijector", node_data[1][0])
      # The shift parameter (and both Normal parameters) are leaves.
      self.assertLen(leaves, 3)
  
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 3 for calling the PytreeTest obj's underlying member method code with example usage: obj.test_user_defined_pytree() without return types