code: def as_stateful_builder(generator):
      """Wraps a generator to build trainable objects parameterized by Variables."""
  
      def error_no_variables(*args, **kwargs):
          """Raises error under JAX and Numpy backends."""
          raise ValueError(
              "TensorFlow is required for `tf.Variable`s. Only "
              "stateless representations are supported under the JAX "
              "and Numpy backends. "
          )
  
      if JAX_MODE or NUMPY_MODE:
          return error_no_variables
  
      # TF-specific imports.
      from tensorflow_probability.python.experimental.util import (
          deferred_module,
      )  # pylint: disable=g-import-not-at-top
  
      def build_stateful_trainable(*args, seed=None, **kwargs):
          g = functools.partial(generator, *args, **kwargs)
          params = _initialize_parameters(g, seed=seed)
          params_as_variables = []
          for name, value in params._asdict().items():
              # Params may themselves be structures, in which case there's no 1:1
              # mapping between param names and variable names. Currently we just give
              # the same name to all variables in a param structure and let TF sort
              # things out.
              params_as_variables.append(
                  tf.nest.map_structure(
                      lambda t, n=name: t
                      if t is None
                      else tf.Variable(t, name=n),
                      value,
                      expand_composites=True,
                  )
              )
          return deferred_module.DeferredModule(
              functools.partial(_apply_parameters, g),
              *params_as_variables,
              also_track=tf.nest.flatten((args, kwargs)),
          )
  
      # Update docstring.
      if generator.__doc__:
          doc = inspect.cleandoc(generator.__doc__)
          # Add `seed` to end of args list.
          seed_str = "seed: PRNG seed; see `tfp.random.sanitize_seed` for details."
          doc = re.sub(r"\nYields:", "  " + seed_str + "\nYields:", doc)
          # Replace `Yields` section with `Returns` a trainable instance.
          doc = re.sub(
              r"Yields:\n(.+\n)+\n",
              _STATEFUL_RETURNS_DOCSTRING,
              doc + "\n\n",
          )
          build_stateful_trainable.__doc__ = doc
  
      return build_stateful_trainable
  
file_path: probability/tensorflow_probability/python/internal/trainable_state_util.py
imports: import collections
  import functools
  import inspect
  import re
  import types
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import structural_tuple
label_desc: Use image node_id 7 for calling a global function with example usage: as_stateful_builder(generator) and returns: build_stateful_trainable, error_no_variables, deferred_module