tf.experimental.ExtensionTypeBatchEncoder

View source on GitHub

Class used to encode and decode extension type values for batching.

In order to be batched and unbatched by APIs such as tf.data.Dataset, tf.keras, and tf.map_fn, extension type values must be encoded as a list of tf.Tensors, where stacking, unstacking, or concatenating these encoded tensors and then decoding the result must be equivalent to stacking, unstacking, or concatenating the original values. ExtensionTypeBatchEncoders are responsible for implementing this encoding.

The default ExtensionTypeBatchEncoder that is used by BatchableExtensionType assumes that extension type values can be stacked, unstacked, or concatenated by simply stacking, unstacking, or concatenating every nested Tensor, ExtensionType, CompositeTensor, and TensorShape field.

Extension types where this is not the case will need to override __batch_encoder__ with a custom encoder that overrides the batch, unbatch, encode, and decode methods. E.g.:

class CustomBatchEncoder(ExtensionTypeBatchEncoder):
  pass # Override batch(), unbatch(), encode(), and decode().

class CustomType(BatchableExtensionType):
  x: tf.Tensor
  y: tf.Tensor
  shape: tf.TensorShape
  __batch_encoder__ = CustomBatchEncoder()

For example, tf.RaggedTensor and tf.SparseTensor both use custom batch encodings which define ops to "box" and "unbox" individual values into tf.variant tensors.

Methods

batch

View source

batch(
    spec, batch_size
)

Returns the TypeSpec representing a batch of values described by spec.

The default definition returns a TypeSpec that is equal to spec, except that an outer axis with size batch_size is added to every nested TypeSpec and TensorShape field. Subclasses may override this default definition, when necessary.

decode

View source

decode(
    spec, encoded_value
)

Decodes value from a batchable tensor encoding.

See encode for a description of the default encoding. Subclasses may override this default definition, when necessary.

encode

View source

encode(
    spec, value, minimum_rank=0
)

Encodes value as a nest of batchable Tensors or CompositeTensors.

The default definition returns a flat tuple of all the Tensors, CompositeTensors, and ExtensionTypes from a depth-first traversal of value's fields. Subclasses may override this default definition, when necessary.

encoding_specs

View source

encoding_specs(
    spec
)

Returns a list of TensorSpec(s) describing the encoding for spec.

See encode for a description of the default encoding. Subclasses may override this default definition, when necessary.

unbatch

View source

unbatch(
    spec
)

Returns the TypeSpec for a single unbatched element in spec.

The default definition returns a TypeSpec that is equal to spec, except that the outermost axis is removed from every nested TypeSpec, and TensorShape field. Subclasses may override this default definition, when necessary.