ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tfp.experimental.auto_composite_tensor

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) Tensors or non-Tensor metadata. list and tuple attributes are supported, but must either contain only Tensors (or lists, etc, thereof), or no Tensors. E.g.,

  • object.attribute = [1., 2., 'abc'] # valid
  • object.attribute = [tf.constant(1.), [tf.constant(2.)]] # valid
  • object.attribute = ['abc', tf.constant(1.)] # invalid

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 Tensors may be expressed as functors that subclass AutoCompositeTensor and used in place of the original callable arg:

@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 Tensors should either (1) subclass AutoCompositeTensor, with the Tensors 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:

@tfp.experimental.auto_composite_tensor
class MyClass(tfp.experimental.AutoCompositeTensor):
  ...

mc = MyClass()
type(mc)
# ==> MyClass

No CompositeTensor base class (ok, but changes expected types)

@tfp.experimental.auto_composite_tensor
class MyClass(object):
  ...

mc = MyClass()
type(mc)
# ==> MyClass_AutoCompositeTensor

CompositeTensor base class, requiring trivial _type_spec

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

@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.

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__.

composite_tensor_subclass A subclass of cls and TF CompositeTensor.