code: def auto_composite_tensor(
      cls=None,
      omit_kwargs=(),
      non_identifying_kwargs=(),
      module_name=None,
  ):
      """Automagically generate `CompositeTensor` behavior for `cls`.
  
      `CompositeTensor` objects are able to pass in and out of `tf.function` and
      `tf.while_loop`, or serve as part of the signature of a TF saved model.
  
      The contract of `auto_composite_tensor` is that all __init__ args and kwargs
      must have corresponding public or private attributes (or properties). Each of
      these attributes is inspected (recursively) to determine whether it is (or
      contains) `Tensor`s or non-`Tensor` metadata. Nested (`list`, `tuple`, `dict`,
      etc) attributes are supported, but must either contain *only* `Tensor`s (or
      lists, etc, thereof), or *no* `Tensor`s. E.g.,
        - object.attribute = [1., 2., 'abc']                        # valid
        - object.attribute = [tf.constant(1.), [tf.constant(2.)]]   # valid
        - object.attribute = ['abc', tf.constant(1.)]               # invalid
  
      All `__init__` args that may be `ResourceVariable`s must also admit `Tensor`s
      (or else `_convert_variables_to_tensors` must be overridden).
  
      If the attribute is a callable, serialization of the `TypeSpec`, and therefore
      interoperability with `tf.saved_model`, is not currently supported. As a
      workaround, callables that do not contain or close over `Tensor`s may be
      expressed as functors that subclass `AutoCompositeTensor` and used in place of
      the original callable arg:
  
      ```python
      @auto_composite_tensor(module_name='my.module')
      class F(AutoCompositeTensor):
  
        def __call__(self, *args, **kwargs):
          return original_callable(*args, **kwargs)
      ```
  
      Callable objects that do contain or close over `Tensor`s should either
      (1) subclass `AutoCompositeTensor`, with the `Tensor`s passed to the
      constructor, (2) subclass `CompositeTensor` and implement their own
      `TypeSpec`, or (3) have a conversion function registered with
      `type_spec.register_type_spec_from_value_converter`.
  
      If the object has a `_composite_tensor_shape_parameters` field (presumed to
      have `tuple` of `str` value), the flattening code will use
      `tf.get_static_value` to attempt to preserve shapes as static metadata, for
      fields whose name matches a name specified in that field. Preserving static
      values can be important to correctly propagating shapes through a loop.
      Note that the Distribution and Bijector base classes provide a
      default implementation of `_composite_tensor_shape_parameters`, populated by
      `parameter_properties` annotations.
  
      If the decorated class `A` does not subclass `CompositeTensor`, a *new class*
      will be generated, which mixes in `A` and `CompositeTensor`.
  
      To avoid this extra class in the class hierarchy, we suggest inheriting from
      `auto_composite_tensor.AutoCompositeTensor`, which inherits from
      `CompositeTensor` and implants a trivial `_type_spec` @property. The
      `@auto_composite_tensor` decorator will then overwrite this trivial
      `_type_spec` @property. The trivial one is necessary because `_type_spec` is
      an abstract property of `CompositeTensor`, and a valid class instance must be
      created before the decorator can execute -- without the trivial `_type_spec`
      property present, `ABCMeta` will throw an error! The user may thus do any of
      the following:
  
      #### `AutoCompositeTensor` base class (recommended)
      ```python
      @tfp.experimental.auto_composite_tensor
      class MyClass(tfp.experimental.AutoCompositeTensor):
        ...
  
      mc = MyClass()
      type(mc)
      # ==> MyClass
      ```
  
      #### No `CompositeTensor` base class (ok, but changes expected types)
      ```python
      @tfp.experimental.auto_composite_tensor
      class MyClass(object):
        ...
  
      mc = MyClass()
      type(mc)
      # ==> MyClass_AutoCompositeTensor
      ```
  
      #### `CompositeTensor` base class, requiring trivial `_type_spec`
      ```python
      from tensorflow.python.framework import composite_tensor
      @tfp.experimental.auto_composite_tensor
      class MyClass(composite_tensor.CompositeTensor):
        @property
        def _type_spec(self):  # will be overwritten by @auto_composite_tensor
          pass
        ...
  
      mc = MyClass()
      type(mc)
      # ==> MyClass
      ```
  
      ## Full usage example
  
      ```python
      @tfp.experimental.auto_composite_tensor(omit_kwargs=('name',))
      class Adder(tfp.experimental.AutoCompositeTensor):
        def __init__(self, x, y, name=None):
          with tf.name_scope(name or 'Adder') as name:
            self._x = tf.convert_to_tensor(x)
            self._y = tf.convert_to_tensor(y)
            self._name = name
  
        def xpy(self):
          return self._x + self._y
  
      def body(obj):
        return Adder(obj.xpy(), 1.),
  
      result, = tf.while_loop(
          cond=lambda _: True,
          body=body,
          loop_vars=(Adder(1., 1.),),
          maximum_iterations=3)
  
      result.xpy()  # => 5.
      ```
  
      Args:
        cls: The class for which to create a CompositeTensor subclass.
        omit_kwargs: Optional sequence of kwarg names to be omitted from the spec.
        non_identifying_kwargs: Optional sequence of kwarg names to be omitted from
          equality/comparison checks and the `__hash__` method of the spec.
        module_name: The module name with which to register the `TypeSpec`. If
          `None`, defaults to `cls.__module__`.
  
      Returns:
        composite_tensor_subclass: A subclass of `cls` and TF CompositeTensor.
      """
      if cls is None:
          return functools.partial(
              auto_composite_tensor,
              omit_kwargs=omit_kwargs,
              non_identifying_kwargs=non_identifying_kwargs,
              module_name=module_name,
          )
  
      if JAX_MODE:
          from jax import (
              tree_util,
          )  # pylint: disable=g-import-not-at-top
  
          tree_util.register_pytree_node(
              cls,
              pytree_flatten,
              functools.partial(pytree_unflatten, cls),
          )
  
      if module_name is None:
          module_name = cls.__module__
  
      type_spec_class_name = f"{cls.__name__}_ACTTypeSpec"
      type_spec_name = f"{module_name}.{type_spec_class_name}"
  
      # If the declared class is already a CompositeTensor subclass, we can avoid
      # affecting the actual type of the returned class. Otherwise, we need to
      # explicitly mix in the CT type, and hence create and return a newly
      # synthesized type.
      if issubclass(cls, composite_tensor.CompositeTensor):
  
          @type_spec_register(type_spec_name)
          class _AlreadyCTTypeSpec(_AutoCompositeTensorTypeSpec):
              @property
              def value_type(self):
                  return cls
  
          _AlreadyCTTypeSpec.__name__ = type_spec_class_name
  
          def _type_spec(obj):
              return _AlreadyCTTypeSpec.from_instance(
                  obj, omit_kwargs, non_identifying_kwargs
              )
  
          # pylint: disable=protected-access
          cls._type_spec = property(_type_spec)
          cls._convert_variables_to_tensors = (
              convert_variables_to_tensors
          )
          # pylint: enable=protected-access
          return cls
  
      clsid = (
          cls.__module__,
          cls.__name__,
          omit_kwargs,
          non_identifying_kwargs,
      )
  
      # Check for subclass if retrieving from the _registry, in case the user
      # has redefined the class (e.g. in a REPL/notebook).
      if clsid in _registry and issubclass(_registry[clsid], cls):
          return _registry[clsid]
  
      class _GeneratedCTTypeSpec(_AutoCompositeTensorTypeSpec):
          @property
          def value_type(self):
              return _registry[clsid]
  
      _GeneratedCTTypeSpec.__name__ = type_spec_class_name
  
      class _AutoCompositeTensor(cls, composite_tensor.CompositeTensor):
          """A per-`cls` subclass of `CompositeTensor`."""
  
          @property
          def _type_spec(self):
              return _GeneratedCTTypeSpec.from_instance(
                  self, omit_kwargs, non_identifying_kwargs
              )
  
          def _convert_variables_to_tensors(self):
              return convert_variables_to_tensors(self)
  
      _AutoCompositeTensor.__name__ = cls.__name__
      _registry[clsid] = _AutoCompositeTensor
      type_spec_register(type_spec_name)(_GeneratedCTTypeSpec)
      return _AutoCompositeTensor
  
imports: import contextlib
  import functools
  import threading
  import numpy
  import tensorflow.compat.v2
  from tensorflow.python.framework import composite_tensor
  from tensorflow.python.framework import type_spec
  from tensorflow.python.framework import type_spec_registry
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.util import tf_inspect
  from tensorflow.python.util import variable_utils
label_desc: Use image node_id 8 for calling a global function with example usage: auto_composite_tensor(cls, omit_kwargs, non_identifying_kwargs, module_name) and returns: _AutoCompositeTensor, functools, cls, _registry, _AlreadyCTTypeSpec, _registry, _GeneratedCTTypeSpec, convert_variables_to_tensors, cls