returnn.tf.util.data

Provides Data, Dim, SearchBeam.

See Tensor and Dim for some higher-level description.

returnn.tf.util.data.FeatureDim(description, dimension, **kwargs)[source]

DEPRECATED. Use Dim instead, and setting the kind is not needed anymore.

Parameters:

  • description (str)

  • dimension (int|None)

Return type:

Dim

returnn.tf.util.data.SpatialDim(description, dimension=None, **kwargs)[source]

DEPRECATED. Use Dim instead, and setting the kind is not needed anymore.

Parameters:

  • description (str)

  • dimension (int|None)

Return type:

Dim

class returnn.tf.util.data.BatchInfo(base, new_dim, new_dim_index=None)[source]

A batched tensor is a tensor with batch dimension, i.e. consisting of multiple samples/sequences which are supposed to be totally independent of each other.

The batch dim can consists out of one or more flattened “virtual” dims, which BatchInfo keeps track of. This class provides some additional information about the batch dimension. Only one instance per different type of batch-dim is expected (i.e. batch_info1 is batch_info2 <==> same batch info).

When we pass data from the dataset to the network (in all cases (training, inference …) via Runner in the TF engine), we get a batch dimension due to the minibatch construction. This is a global batch size (usually dynamic, because every minibatch/step can have a different amount of samples, although we can also support static sizes, which is needed e.g. for TPUs) represented by BatchInfo.GlobalBatchDim.

When we do beam search (see SearchBeam), we have multiple hypotheses per batch item, and thus a different batch dimension.

We can also pack arrays (multiple sequences) (also referred to as “flattened tensors”, “non-padded tensors”, “ragged tensors”). See e.g. FlattenBatchLayer or flatten_with_seq_len_mask(). Also see tf.RaggedTensor which also represents packed tensors (but only batch-major, although this is just a reinterpretation). We do not directly use tf.RaggedTensor in Data to have robust and reliable code (which expects tf.Tensor). However, we maybe can make use of some of the functions in tf.ragged.

Parameters:

  • base (BatchInfo|None) – If this is extended or based on another batch. Except of the batch dim of the dataset minibatch, we would always have a base.

  • new_dim (BatchInfo.VirtualDimBase|None)

  • new_dim_index (int|None)

In most cases, this constructor would probably not be used directly by the user.

class VirtualDimBase[source]

Represents one virtual dim, flattened into the batch dim.

short_repr()[source]
Return type:

str

class FixedDim(size: Tensor | int, dim_tag: Dim | None = None)[source]

Represents a dim with fixed size.

Parameters:

  • size

  • dim_tag

short_repr()[source]
Return type:

str

class GlobalBatchDim(size: Tensor | int, dim_tag: Dim | None = None)[source]

Represents the global batch dim by the network (minibatch construction from the dataset).

Parameters:

  • size

  • dim_tag

short_repr()[source]
Return type:

str

class BeamDim(beam)[source]

Represents a search beam.

Parameters:

beam (SearchBeam)

short_repr()[source]
Return type:

str

class PaddedDim(dim_tag)[source]

Represents a dim with variable size, which is flattened with padding (not packed) into the batch.

Parameters:

dim_tag (Dim)

short_repr()[source]
Return type:

str

class PackedDim(dim_tag, key_axes)[source]

Represents a dim with variable sizes, which is packed (un-padded) into the batch. Variable w.r.t. other dims (must be per batch entry).

Parameters:

  • dim_tag (Dim)

  • key_axes (list[BatchInfo.VirtualDimBase]) – most common case would be [GlobalBatchDim(…)], but [GlobalBatchDim(…),BeamDim(…)] is also common.

property sizes[source]
Returns:

shape [B_flat]

Return type:

tf.Tensor

short_repr()[source]
Return type:

str

classmethod make_global_batch_info(batch_dim)[source]
Parameters:

batch_dim (tf.Tensor|int)

Returns:

global batch info w.r.t. the network / graph

Return type:

BatchInfo

classmethod make_global_broadcast_batch_info()[source]
Returns:

BatchInfo with no virtual dims, s.t. the dimension is 1 (== prod([])) (broadcastable)

Return type:

BatchInfo

classmethod get_common_batch_info(batches)[source]
Parameters:

batches (list[BatchInfo|None])

Return type:

BatchInfo|None

short_repr()[source]
Return type:

str

property dim[source]
Return type:

tf.Tensor|int

property static_dim[source]
Return type:

int|None

property beam[source]
Return type:

SearchBeam|None

get_base_chain()[source]
Return type:

list[BatchInfo]

get_global_base()[source]
Return type:

BatchInfo

get_global_batch_dim()[source]
Return type:

BatchInfo.GlobalBatchDim

is_global_batch()[source]
Return type:

bool

is_broadcast()[source]
Return type:

bool

copy_extend_with_beam(beam)[source]
Parameters:

beam (SearchBeam)

Return type:

BatchInfo

copy_remove_beam()[source]
Return type:

BatchInfo

copy_remove_dim(remove_dim)[source]
Parameters:

remove_dim (VirtualDimBase)

Return type:

BatchInfo

copy_set_beam(beam)[source]
Parameters:

beam (SearchBeam|None)

Return type:

BatchInfo

copy_extend_with_packed_dim_tag(dim_tag, batch_major)[source]
Parameters:

  • dim_tag (Dim)

  • batch_major (bool) – if True, add new dim in front. otherwise, add new dim at the end

Return type:

BatchInfo

copy_extend_with_padded_dim_tag(dim_tag, batch_major=None, new_dim_idx=None)[source]
Parameters:

  • dim_tag (Dim)

  • batch_major (bool|None) – if True, add new dim in front. otherwise, add new dim at the end

  • new_dim_idx (int|None)

Return type:

BatchInfo

copy_extend_with_padded_or_fixed_dim_tag(dim_tag, batch_major=None, new_dim_idx=None)[source]
Parameters:

  • dim_tag (Dim)

  • batch_major (bool|None) – if True, add new dim in front. otherwise, add new dim at the end

  • new_dim_idx (int|None)

Return type:

BatchInfo

class returnn.tf.util.data.SearchBeam(beam_size, dependency=<class 'returnn.util.basic.NotSpecified'>, name=None, _next_frame=None)[source]

Represents info about the beam from some beam search (e.g. via beam_search()), e.g. such as the beam size, but also the dependencies. This is somewhat parallel to SearchChoices, but simpler, and independent from the layers/network (returnn.tf.layers.base.LayerBase).

Parameters:
copy_as_prev_frame()[source]
Return type:

SearchBeam

classmethod get_combined_beam(beam1, beam2=None, *beams)[source]

Combines beams. This will throw an exception if they cannot be combined. Note that in beam search (see SearchChoices), the logic to combine beams from different search choices happens in a generic way for all layers automatically via TFNetwork._create_layer_layer_desc(), so normally we already have the same beam. Unless we are at template construction.

Parameters:
Return type:

SearchBeam|None