tf.experimental.ExtensionType

Base class for TensorFlow ExtensionType classes.

tf.experimental.ExtensionType(
    *args, **kwargs
)

Tensorflow ExtensionType classes are specialized Python classes that can be used transparently with TensorFlow -- e.g., they can be used with ops such as tf.cond or tf.while_loop and used as inputs or outputs for tf.function and Keras layers.

New ExtensionType classes are defined by creating a subclass of tf.ExtensionType that contains type annotations for all instance variables. The following type annotations are supported:

Type Example
Python integers i: int
Python floats f: float
Python strings s: str
Python booleans b: bool
Python None n: None
Tensors t: tf.Tensor
Composite Tensors rt: tf.RaggdTensor
Extension Types m: MyMaskedTensor
Tensor shapes shape: tf.TensorShape
Tensor dtypes dtype: tf.DType
Type unions length: typing.Union[int, float]
Tuples params: typing.Tuple[int, float, int, int]
Tuples w/ Ellipsis lengths: typing.Tuple[int, ...]
Mappings tags: typing.Mapping[str, str]

Fields annotated with typing.Mapping will be stored using an immutable mapping type.

ExtensionType values are immutable -- i.e., once constructed, you can not modify or delete any of their instance members.

Examples

class MaskedTensor(ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor
 
class Toy(ExtensionType):
  name: str
  price: ops.Tensor
  features: typing.Mapping[str, tf.Tensor]
 
class ToyStore(ExtensionType):
  name: str
  toys: typing.Tuple[Toy, ...]

Methods

__eq__

View source

__eq__(
    other
)

Return self==value.

__ne__

View source

__ne__(
    other
)

Return self!=value.