code: def test_rewrites_yield_to_return_in_docstring(self):
      wrapped = trainable_state_util.as_stateful_builder(
          generator_with_docstring
      )
      self.assertContainsExactSubsequence(
          generator_with_docstring.__doc__, "Yields:"
      )
      self.assertNotIn("Yields:", wrapped.__doc__)
      self.assertContainsExactSubsequence(
          wrapped.__doc__, "Test generator with a docstring."
      )
      self.assertContainsExactSubsequence(
          wrapped.__doc__,
          trainable_state_util._STATEFUL_RETURNS_DOCSTRING,
      )
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import joint_map
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.internal import trainable_state_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.minimize import minimize
  from tensorflow_probability.python.math.minimize import minimize_stateless
label_desc: Use image node_id 3 for calling the TestWrapGeneratorAsStateful obj's underlying member method code with example usage: obj.test_rewrites_yield_to_return_in_docstring() without return types