returnn.tf.util.data

Provides Data, Dim, SearchBeam.

See Data and Dim for some higher-level description.

class returnn.tf.util.data.Dim(kind=None, description=None, dimension=None, vocab=None, dyn_size=None, dyn_size_ext=None, undefined=False, generic=False, special=False, match_priority=0, derived_from_tag=None, derived_from_op=None, batch=None, control_flow_ctx=None, src_data=None, src_axis=None)

This identifies one axis/dimension, like a time-dimension, etc. This was called DimensionTag earlier, and referred to as dimension tag.

This is used by Data. See Data.dim_tags(). This would be passed as dim_tags when creating a Data instance.

It is not to specify the specific axis in a specific Data/tensor, but to specify the content and dimension. I.e. if we have the same Dim for two Data instances, the dimensions should match. I.e.:

data1.get_dim_tag(i) == data2.get_dim_tag(j)

=> tf.shape(data1.placeholder)[i] == tf.shape(data2.placeholder)[j]

This also includes further information such as sequence lengths or a vocabulary.

We differentiate between the batch dim, spatial dim or feature dim, although that is just flag and in many contexts there is no real difference between a spatial dim and a feature dim (the batch dim is often handled differently).

See SpatialDim() and FeatureDim() as easy wrappers to create dim tags for the user.

Parameters:

• kind (Entity|None) –
• description (str|None) – the description should be unique
• dimension (int|None) –
• vocab (returnn.datasets.util.vocabulary.Vocabulary|None) –
• dyn_size (tf.Tensor|None) – e.g. seq_len, (batch,)
• dyn_size_ext (Data|None) – seq_len or extended
• undefined (bool) – When this is specified as None by the user via shape.
• generic (bool) – This can not be a dim tag of Data but it would match to other existing dim tags Data.get_axis_from_description() and Dim.is_equal().
• special (bool) – Like generic, this can not be a dim tag of Data. But this dim tag also does not match anything except itself. So it can be used to represent special placeholders with special meanings like single_step.
• derived_from_tag (Dim|None) – Whether this new tag is reduced, down/up sampled, padded etc from this given other tag. In situations where dim tags are being matched (Data.get_common_data), the behavior is to consider them as equal, and assume that the chain of operations (e.g. padding + valid conv) results in the same dim.
• derived_from_op (Dim.Op|None) –
• match_priority (int) – when there is ambiguity between multiple dim tags, this value defines the order in which the dimension are assigned to their matching counterparts. A dimension tag with a higher priority value is assigned first. E.g. for a square matrix used for a linear transformation, the reduce dim tag should have a higher priority.
• batch (BatchInfo|None) – for batch-dim, or dynamic dims per batch
• control_flow_ctx (ControlFlowContext|None) –
• src_data (Data|None) –
• src_axis (int|None) –

class Types

Defines possible values for kind.

Unspecified = None

Batch = <batch>

Spatial = <spatial>

Time = <spatial>

Feature = <feature>

Types = (<batch>, <spatial>, <feature>)

short_repr()

Returns:some short repr Return type:str

copy(same_as_self=True, description=None, kind=None, match_priority=None)

Parameters:

• same_as_self (bool) –
• description (str|None) – new description
• kind (Entity|None) – if set, overwrites self.kind
• match_priority (int|None) –

Returns:

copy, maybe as new kind. setting same_as to self

Return type:

Dim

get_for_batch_ctx(batch, ctx, allow_none=False)

Parameters:

• batch (BatchInfo) –
• ctx (ControlFlowContext|None) –
• allow_none (bool) –

Return type:

Dim|None

set_dyn_size_ext_for_batch_ctx(batch, ctx, dyn_size_ext)

Parameters:

• batch (BatchInfo) –
• ctx (ControlFlowContext|None) –
• dyn_size_ext (Data) –

get_dyn_size_ext_for_batch_ctx(batch, ctx)

Parameters:

• batch (BatchInfo|None) –
• ctx (ControlFlowContext|None) –

Return type:

Data|None

dyn_size

Returns:dyn size / seq len (usually of shape [B]), or None If the dyn size can potentially be of a different shape, directly access dyn_size_ext. Return type:tf.Tensor|None

is_batch_dim()

Returns:whether this dim tag is of kind batch Return type:bool

is_feature_dim()

Returns:whether this dim tag is of kind feature Return type:bool

is_spatial_dim()

Returns:whether this dim tag is of kind spatial Return type:bool

is_dim_known()

Returns:whether we know the dimension; basically whether this is defined (although not self.undefined is defined slightly differently) Return type:bool

is_dynamic()

Returns:whether the dim is not static. usually means that it has seq lengths Return type:bool

can_be_used_as_dim()

Returns:whether this can be used as a dim in Data, i.e. it is not generic or special Return type:bool

is_same_size_tensor(x)

Parameters:x (tf.Tensor) – Returns:whether this dim tag for this specific batch (incl beam) is the same as the given size Return type:bool

set_tag_on_size_tensor(x, batch=None, same_as_before=False)

This function is used to couple a tf.Tensor instance representing the dyn size with the dim tag.

This is usually a newly created dim tag, which is yet unset.

It is also used to couple an existing dim tag with other dyn sizes which just differ by an expansion of the batch (e.g. search beam).

See also get_tag_from_size_tensor().

Parameters:

• x (tf.Tensor) –
• batch (BatchInfo|None) –
• same_as_before (bool) – implies it was set before, and the new size is the same. e.g. it could be some identity with added checks, or other change.

Returns:

self or new dim tag

Return type:

Dim

classmethod get_tag_from_size_tensor(x)

Parameters:x (tf.Tensor) – size tensor. has been set before via set_tag_on_size_tensor() Return type:Dim|None

complete_dyn_size()

In case we can calculate the dyn size, do that now.

is_equal(other, ignore_feature_dim=False, allow_same_feature_dim=False, allow_same_spatial_dim=None, treat_feature_as_spatial=False, broadcast_matches=False, unknown_spatial_matches=False, undefined_matches=False, derived_matches=False)

Compares self to other for equality.

Note that the default behavior is very restrictive. Use functions such as get_all_dimension_tags() or get_existing_tag_from_collection() to explicitly specify the behavior for the comparison.

Also note that the definition is slightly ad-hoc for some cases, and might potentially change in the future.

https://github.com/rwth-i6/returnn/issues/634

Parameters:

• other (Dim) –
• ignore_feature_dim (bool) –
• allow_same_feature_dim (bool) –
• allow_same_spatial_dim (bool|None) –
• treat_feature_as_spatial (bool) –
• broadcast_matches (bool) –
• unknown_spatial_matches (bool) –
• undefined_matches (bool) –
• derived_matches (bool) –

Return type:

bool

get_same_base()

Return type:Dim

get_same_derived_base()

Return type:Dim

undefined

Return type:bool

declare_same_as(other)

Parameters:other (Dim) –

classmethod get_existing_tag_from_collection(other, tags, is_equal_opts=None)

Parameters:

• other (Dim) –
• tags (list[Dim]|tuple[Dim]|set[Dim]) –
• is_equal_opts (dict[str]|None) – passed to Dim.is_equal

Return type:

Dim|None

classmethod get_all_dimension_tags(data_list, is_equal_opts=None, unique_separate_axes=True)

Parameters:

• data_list (list[Data]) –
• is_equal_opts (dict[str]|None) – passed to Dim.is_equal
• unique_separate_axes (bool) – e.g. data_list=[Data with shape (B,5,5,10)] results in 4 dim tags, not 3.

Returns:

list of dimension tags, dict for data -> list of dimension tags (for each axis)

Return type:

(list[Dim], dict[Data, list[Dim]])

classmethod get_uniq_collection(tags, is_equal_opts=None)

Parameters:

• tags (list[Dim]|tuple[Dim]|set[Dim]) –
• is_equal_opts (dict[str]|None) – passed to Dim.is_equal

Return type:

list[Dim]

get_dim_value()

Infers the dim this axis should have if unbroadcasted. If self.src_data has a placeholder, will use the shape from there. Otherwise, uses self.dimension (if static) or self.dyn_size (if dynamic). :rtype: int|tf.Tensor

axis_split_info()

Returns:axis split info. see get_param_axes_split_info() and usage (e.g. pretraining) Return type:list[int|None]

vocab

Return type:returnn.datasets.util.vocabulary.Vocabulary|None

sub_right(other)

Parameters:other (Dim|int) – Returns:self - other Return type:Dim

sub_left(other)

Parameters:other (Dim|int) – Returns:(-other) + self Return type:Dim

div_left(other)

Parameters:other (Dim|int) – Return type:Dim

div_right(other)

Parameters:other (Dim|int) – Return type:Dim

ceildiv_left(other)

Parameters:other (Dim|int) – Return type:Dim

ceildiv_right(other)

Parameters:other (Dim|int) – Return type:Dim

class Op(kind, inputs, attribs=None)

Op on Dim which results in a derived Dim.

Parameters:

• kind (str) – “add”, “sub”, “mul”, “ceildiv”
• inputs (list[Dim]) –
• attribs (dict[str]|None) –

returnn.tf.util.data.DimensionTag

alias of returnn.tf.util.data.Dim

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

Parameters:

• description (str) –
• dimension (int|None) –

Return type:

Dim

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

Parameters:

• description (str) –
• dimension (int|None) –

Return type:

Dim

class returnn.tf.util.data.ImplicitSparseDim(tag)

Represents an implicit dim via Data.sparse_dim.

Parameters:tag (Dim) –

class returnn.tf.util.data.ImplicitDynSizeDim(tag)

Represents an implicit dim via dynamic dim sizes. https://github.com/rwth-i6/returnn/issues/706 (For example via CumConcatLayer.)

Parameters:tag (Dim) –

class returnn.tf.util.data.OptionalDim(tag)

Represents a dim which might exist or not.

Parameters:tag (Dim) –

exception returnn.tf.util.data.VerifyOutShapeException

Exception via Data.verify_out_shape().

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

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

Represents one virtual dim, flattened into the batch dim.

short_repr()

Return type:str

class FixedDim(size, dim_tag=None)

Represents a dim with fixed size.

Parameters:

• size (tf.Tensor|int) –
• dim_tag (Dim|None) –

short_repr()

Return type:str

class GlobalBatchDim(size, dim_tag=None)

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

Parameters:

• size (tf.Tensor|int) –
• dim_tag (Dim|None) –

short_repr()

Return type:str

class BeamDim(beam)

Represents a search beam.

Parameters:beam (SearchBeam) –

short_repr()

Return type:str

class PaddedDim(dim_tag)

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

Parameters:dim_tag (Dim) –

short_repr()

Return type:str

class PackedDim(dim_tag, key_axes)

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.

sizes

Returns:shape [B_flat] Return type:tf.Tensor

short_repr()

Return type:str

classmethod make_global_batch_info(batch_dim)

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()

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

classmethod get_common_batch_info(batches)

Parameters:batches (list[BatchInfo|None]) – Return type:BatchInfo|None

short_repr()

Return type:str

dim

Return type:tf.Tensor|int

static_dim

Return type:int|None

beam

Return type:SearchBeam|None

get_base_chain()

Return type:list[BatchInfo]

get_global_base()

Return type:BatchInfo

get_global_batch_dim()

Return type:BatchInfo.GlobalBatchDim

is_global_batch()

Return type:bool

is_broadcast()

Return type:bool

copy_extend_with_beam(beam)

Parameters:beam (SearchBeam) – Return type:BatchInfo

copy_remove_beam()

Return type:BatchInfo

copy_remove_dim(remove_dim)

Parameters:remove_dim (VirtualDimBase) – Return type:BatchInfo

copy_set_beam(beam)

Parameters:beam (SearchBeam|None) – Return type:BatchInfo

copy_extend_with_packed_dim_tag(dim_tag, batch_major)

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)

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)

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)

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:

• beam_size (int) –
• dependency (SearchBeam|NotSpecified|None) –
• name (str|None) –
• _next_frame (SearchBeam|None) –

copy_as_prev_frame()

Return type:SearchBeam

classmethod get_combined_beam(beam1, beam2=None, *beams)

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:

• beam1 (SearchBeam|None) –
• beam2 (SearchBeam|None) –
• beams (SearchBeam|None) –

Return type:

SearchBeam|None

class returnn.tf.util.data.Data(name, shape=None, dtype=None, placeholder=None, sparse=None, sparse_dim=<class 'returnn.util.basic.NotSpecified'>, dim=<class 'returnn.util.basic.NotSpecified'>, size_placeholder=None, batch_dim_axis=<class 'returnn.util.basic.NotSpecified'>, time_dim_axis=<class 'returnn.util.basic.NotSpecified'>, feature_dim_axis=<class 'returnn.util.basic.NotSpecified'>, available_for_inference=True, auto_create_placeholders=False, vocab=None, dim_tags=None, same_dim_tags_as=None, batch=None, beam=None, control_flow_ctx=None)

This class is to describe a tensor, i.e. its shape and properties like whether we should consider it sparse data (i.e. it represents indices). Each dimension is described by Dim.

This is used in TFNetwork to describe the dataset external data (ExternData) as well as in every layer’s output and in many other parts of the code.

See Data and Dim.

Parameters:

• name (str) –
• shape (tuple[int|None]|list[int|None]) – including time-dim (can be None). excluding batch-dim. e.g. (time,feat)=(None,128)
• dtype (str) – e.g. “float32” or “int64”
• placeholder (tf.Tensor|None) – with added batch-dim
• sparse (bool|None) – whether to treat the value as an index. do not confuse with tf.SparseTensor
• sparse_dim (Dim|int|None|NotSpecified) –
• dim (int|None|NotSpecified) – feature dimension, shape[-1] if not sparse, otherwise like num_classes
• batch_dim_axis (int|None|NotSpecified) – where we add the batch-dim. e.g. shape=(time,…), 0 -> (batch,time,…), 1 -> (time,batch,…). Default is 0. This is normally always set, and a lot of code expects this. However, you can set it to None if this Data does not have a batch-dim.
• time_dim_axis (int|None|NotSpecified) – where we have the time dim axis, after we added the batch-dim. this is often 1. however, can be None if there is no time-dim.
• feature_dim_axis (int|None|NotSpecified) – feature dim axis. by default it’s the last one
• size_placeholder (dict[int,tf.Tensor]|None) – for every None in shape, this will describe the size. The size is always a tensor of shape (batch,), i.e. the size can be different for each sequence in a batch.
• available_for_inference (bool) – e.g. the extern data “classes” is usually not available for inference
• auto_create_placeholders (bool) – This will create a tf.placeholder.
• vocab (str|dict[str]|returnn.datasets.util.vocabulary.Vocabulary|None) –
• dim_tags (tuple[Dim]|list[Dim]|dict[int,Dim]|None) – If tuple/list, this specifies the whole (batch) shape. If dict, explicitly specified dimension tags per axis (axis counted with batch-dim)
• same_dim_tags_as (dict[int|str,Dim]|None) – will mark our dimension tags to be the same
• batch (BatchInfo|None) –
• beam (SearchBeam|None) – the batch-dim could be extended by a beam-size, such that it represents the merged dims [batch, beam_size].
• control_flow_ctx (ControlFlowContext|None) –

size_dtype = 'int32'

classmethod from_tensor(x)

Parameters:x (tf.Tensor) – Return type:Data

classmethod template_from_constant(x, name, dtype=None, shape=None, with_batch_dim=False, sparse_dim=None)

Parameters:

• x (int|float|bool|numpy.ndarray) – not actually assigned to the returned Data, just for the shape and dtype
• name (str) –
• dtype (str|None) –
• shape (list[Dim|int]|tuple[Dim|int]|None) – for verification, and defining dim tags. might also already include the batch-dim. (Then with_batch_dim is ignored.)
• with_batch_dim (bool) –
• sparse_dim (Dim|None) –

Returns:

data template

Return type:

Data

sanity_check(ignore_placeholder=False, assume_complete=True)

Performs some sanity checks on self, and raises exceptions if something is not sane.

Parameters:

• ignore_placeholder (bool) –
• assume_complete (bool) –

get_runtime_sanity_check_op()

Returns:op which does a couple of runtime sanity checks on the placeholder Return type:tf.Operation

verify_out_shape(out_shape)

Verifies that out_shape matches our shape, i.e. specifically the dim tags.

https://github.com/rwth-i6/returnn/issues/706

Throws an exception if this is not the case.

Parameters:out_shape (set[Dim|_MarkedDim]|tuple|list) – It must be a set, with the only exception when it is empty (then it doesn’t matter). See dim_tags_set().

get_placeholder_kwargs(with_batch=True)

Parameters:with_batch (bool) – Returns:kwargs for tf.compat.v1.placeholder Return type:dict[str]

get_axes_with_size()

Returns:list of axes which can vary in size for each entry of the batch-dim, e.g. the time-dim-axis. The axis index is counted without the batch-dim. Return type:list[int]

get_kwargs(include_special_axes=True)

Parameters:include_special_axes (bool) – whether to include time and feature special axis marker Returns:relevant attrib items for copying Return type:dict[str]

get_description(with_name=True, with_placeholder=False, catch_exceptions=False)

Parameters:

• with_name (bool) –
• with_placeholder (bool) –
• catch_exceptions (bool) –

Returns:

description of self. also used for __repr__

Return type:

str

get_batch_axes_short_description(special_axes=True)

Parameters:special_axes (bool) – special markers for old-style time_dim_axis and feature_dim_axis Return type:list[str]

get_compare_key()

Returns:some key which can be used for compare functions, i.e. such that cmp(get_compare_key(self), get_compare_key(other)) == cmp(self, other), i.e. we define some order by that. Note that this order is not totally fixed, and might change. Return type:object

copy(name=None)

Parameters:name (str) – if given, will overwrite this name Returns:copy of myself, using self.get_kwargs(), and with placeholder and size_placeholder Return type:Data

copy_as_batch_major()

Returns:copy of myself with batch_dim_axis == 0 Return type:Data

copy_as_time_major()

Returns:copy of myself with time_dim_axis == 0 Return type:Data

copy_with_batch_dim_axis(batch_dim_axis)

Parameters:batch_dim_axis (int) – Returns:copy of myself with specific batch_dim_axis Return type:Data

copy_with_time_dim_axis(time_dim_axis)

Parameters:time_dim_axis (int) – Returns:copy of myself with specific time_dim_axis Return type:Data

copy_transpose(perm)

Parameters:perm (list[int]) – permutation of the axes, counted with batch-dim. Maps the new axes to the old axes Returns:copy of myself with permuted axes Return type:Data

copy_move_axis(old_axis, new_axis)

Parameters:

• old_axis (int) – counted with batch-dim
• new_axis (int) – counted with batch-dim

Returns:

copy of myself with moved axis (see move_axis())

Return type:

Data

copy_swap_axes(axis1, axis2)

Like Data.copy_move_axis(), but keeps all other axes unchanged. :param int axis1: counted with batch-dim :param int axis2: counted with batch-dim :return: copy of myself with moved axis (see swapaxes()) :rtype: Data

copy_as_bt_or_tb_major()

Return type:Data Returns:copy of myself in batch-time-major or time-batch-major

copy_with_feature_dim_axis(feature_dim_axis)

Parameters:feature_dim_axis (int) – can also be negative Returns:copy of myself with specific feature dim axis Return type:Data

copy_as_batch_feature_major()

Returns:copy of self with batch_dim_axis == 0 and feature_dim_axis == 1 Return type:Data

copy_as_time_batch_major()

Returns:copy of self with batch_dim_axis == 1 and time_dim_axis == 0 Return type:Data

copy_as_batch_spatial_major()

Returns:copy with batch_dim_axis == 0, then all dynamic axes, then any other spatial axes, last feature axis Return type:Data

copy_with_feature_last()

Returns:copy of self with feature_dim_axis being the very last axis Return type:Data

copy_add_batch_dim(batch_dim_axis, batch=None, dim_tag=None)

Parameters:

• batch_dim_axis (int) –
• batch (BatchInfo|None) –
• dim_tag (Dim|None) –

Returns:

copy of myself with added batch-dim

Return type:

Data

copy_add_spatial_dim(spatial_dim_axis=None, dim=1, auto_time_dim_axis=True)

Parameters:

• spatial_dim_axis (int|None) – counted with batch-dim. if there is no time-dim, this will be it.
• dim (int|None) –
• auto_time_dim_axis (bool) –

Returns:

copy of myself with added spatial-dim

Return type:

Data

copy_add_feature_dim(axis=None)

Parameters:axis (int|None) – Returns:self with a new feature dim axis with dim 1. If there is an existing feature dim, the new feature dim will be added right after. If we are sparse, we don’t add a feature dim, but it becomes a spatial dim instead. Return type:Data

get_default_new_axis_for_dim_tag(dim_tag)

Parameters:dim_tag (Dim) – Return type:int

copy_add_dim_by_tag(dim_tag, unbroadcast=False, axis=None)

Parameters:

• dim_tag (Dim) –
• unbroadcast (bool) – If True unbroadcast the newly added axis. Will infer the unbroadcast shape via Dim.get_dim_value()
• axis (int|None) –

Return type:

Data

copy_split_feature_dim(new_feature_dim)

Split it into (new_spatial_dim, new_feat_dim), in that order. This will increase the feature_dim_axis by one.

Parameters:new_feature_dim (int) – will be the new dim Return type:Data

copy_extend_batch(batch)

Similar as copy_compatible_to with unbroadcast=True, we would possibly extend/expand our batch dim. See BatchInfo. This assumes that we already have a batch dim (otherwise see copy_add_batch_dim()).

This excludes any beam expansion, which is handled explicitly elsewhere (e.g. see copy_extend_with_beam()).

Parameters:batch (BatchInfo) – Return type:Data

copy_compatible_to(data, add_dims=True, unbroadcast=False, except_feature=False, except_axis=None, check_sparse=True, check_dtype=True)

Parameters:

• data (Data) – other data which the returned tensor should be compatible to It would add any missing axes with a dim 1 axis for automatic broadcasting (with add_dims=True). It currently does not check whether existing dims match.
• add_dims (bool) – whether to add (broadcast, or unbroadcasted) dims. throws error if missing dim
• unbroadcast (bool) – if True, all added broadcast axes (axes with dim 1) will be tiled such that they match
• except_feature (bool) – if unbroadcast, do not unbroadcast the feature dim
• except_axis (Dim|int|None) – if unbroadcast, do not unbroadcast this axis
• check_sparse (bool) –
• check_dtype (bool) –

Returns:

Data, might add broadcast dimensions

Return type:

Data

copy_time_flattened()

Returns:copy of myself where the time-axis is flattened away into the batch-dim-axis. See get_placeholder_time_flattened() and :func:`flatten_with_seq_len_mask for more details. Return type:Data

copy_extend_with_beam(beam)

Parameters:beam (SearchBeam|None) – Returns:copy of myself where the batch-dim is extended/multiplied by beam_size, using tile_transposed Return type:Data

copy_merge_into_batch(axes)

Parameters:axes (list[int]) – All axes to be merged into the batch axis. Must include the batch_dim_axis. The order is kept. Returns:copy of myself where the the given axes are merged into the batch dim Return type:Data

copy_squeeze_axes(axes)

Parameters:axes (list[int]) – counted with batch dim Returns:copy of myself, with squeezed axes Return type:Data

copy_template(name=None, dtype=None)

Parameters:

• name (str|None) –
• dtype (str|None) –

Returns:

copy of myself, using self.get_kwargs(), without placeholder

Return type:

Data

copy_template_dense(name=None, dtype=None)

Parameters:

• name (str|None) –
• dtype (str|None) –

Returns:

copy of myself, using self.get_kwargs(), without placeholder

Return type:

Data

copy_template_excluding_axis(exclude_axis, name=None)

Parameters:

• exclude_axis (int) – axis to be removed.
• name (str|None) – if set, this will be the new name.

Returns:

copy of myself excluding exclude_axis axis, without placeholder.

Return type:

Data

copy_template_excluding_spatial_dim(spatial_axis_num, name=None)

Parameters:

• spatial_axis_num (int) – index in self.get_spatial_batch_axes()
• name (str|None) – if set, this will be the new name

Returns:

copy of myself excluding the time-dimension without placeholder

Return type:

Data

copy_template_excluding_time_dim(name=None)

Parameters:name (str|None) – if set, this will be the new name Returns:copy of myself excluding the time-dimension without placeholder Return type:Data

copy_template_adding_time_dim(name=None, time_dim_axis=0)

Adds a time-dim-axis. If a time-dim-axis already exists, it will anyway create this new one.

Parameters:

• name (str|None) – if set, this will be the new name
• time_dim_axis (int) – the new time-dim-axis index

Returns:

copy of myself adding the time-dimension without placeholder

Return type:

Data

copy_template_replace_dim_tag(axis, new_dim_tag, name=None)

Parameters:

• axis (int) –
• new_dim_tag (Dim) –
• name (str|None) – new name

Return type:

Data

copy_template_replace_dim(axis, new_dim, new_size=None)

Parameters:

• axis (int) –
• new_dim (int|None) –
• new_size (tf.Tensor|None) –

Return type:

Data

copy_template_new_dim_tags(new_dim_tags, name=None, keep_special_axes=False)

Parameters:

• new_dim_tags (list[Dim]|tuple[Dim]) –
• name (str|None) –
• keep_special_axes (bool) –

Return type:

Data

copy_template_set_ctx(ctx)

Parameters:ctx (ControlFlowContext) – Returns:new Data instance Return type:Data

copy_template_unpack_batch()

If the batch dim contains a BatchInfo.PackedDim, unpack it and restore the data from before the packing.

Return type:Data

matches_var_dim_pattern(other)

Parameters:other (Data) – Returns:whether the variable-dims pattern matches, i.e. same variable dims (get_variable_dim_pattern), same time dim, excluding batch-dim. i.e. the size_placeholder should be compatible. Return type:bool

dim_tags

Return type:tuple[Dim]

shape

Returns:shape without batch-dim. e.g. (time,feat) = (None,128) Return type:tuple[int|None]

batch_shape

Returns:shape with added batch-dim. e.g. (batch,time,feat) = (None,None,128) Return type:tuple[int|None]

get_batch_shape(batch_dim)

Parameters:batch_dim (int|tf.Tensor|None) – Returns:shape with added batch-dim. e.g. (batch,time,feat) = (None,None,128) Return type:tuple[int|None]

get_dynamic_batch_shape()

Return type:list[int|tf.Tensor]

have_varying_shape_in_ctx()

Returns:whether the (dynamic) shape can change in this control flow context. E.g. when self.control_flow_context is a loop, and we have one dynamic dim where dyn_size_ext has the same control_flow_context (such that dyn_size_ext has e.g. shape [B,T] outside the loop). This can be relevant for accumulating values of self.placeholder e.g. via tf.TensorArray. Return type:bool

size_placeholder

For compatibility, return a proxy object which behaves like the original dict.

Return type:_SizePlaceholderProxy

shape_dense

Returns:shape with feature dim axis Return type:tuple[int|None]

batch_shape_dense

Return type:tuple[int|None]

dim_tags_sparse

Returns:dim tags without feature dim axis Return type:tuple[Dim]

dim_tags_set_implicit_only_wrapped

Returns:Dim tags implicit by sparse dim, or dynamic sizes, and not present as explicit dims. Also see dim_tags_set(). Return type:set[_ImplicitDim]

dim_tags_set_implicit_only

Returns:Dim tags implicit by sparse dim, or dynamic sizes, and not present as explicit dims. Also see dim_tags_set(). Return type:set[Dim]

dim_tags_set_implicit

This is mostly intended to be used for verification, such as out_shape in a layer.

https://github.com/rwth-i6/returnn/issues/706

We return a set because when dim tags (dimensions, and the shape) are checked, we never want that the order plays any role.

https://github.com/rwth-i6/returnn/wiki/RETURNN-principles

Further, dimension tags should ideally be unique.

https://github.com/rwth-i6/returnn/issues/632

(This is not enforced currently, but we should not treat this specially now.)

Returns:set of dim tags Return type:set[Dim]

ndim

Return type:int Returns:ndim counted without batch-dim

ndim_dense

Return type:int Returns:ndim counted without batch-dim, added by 1 if we are sparse

batch_ndim

Return type:int Returns:ndim counted with batch-dim

batch_ndim_dense

Return type:int Returns:ndim counted with batch-dim, added by 1 if we are sparse

is_time_major

Returns:whether this is in time-major format, i.e. (time,batch,…) Return type:bool

is_batch_major

Returns:whether this is in batch-major format, i.e. (batch,…) Return type:bool

is_batch_feature_major

Returns:whether this is in batch-feature-major format, i.e. (batch,feature,…) (NC…) Return type:bool

batch_dim_axis

Returns:batch dim axis, counted with batch-dim Return type:int|None

feature_dim_axis

Returns:feature dim axis, counted with batch-dim Return type:int|None

feature_dim_axis_or_unspecified

Returns:feature dim axis, counted with batch-dim. could also be unspecified Return type:int|None|NotSpecified

time_dim_axis_excluding_batch

Return type:int|None

placeholder

Return type:tf.Tensor|None

batch

Return type:BatchInfo|None

beam

Return type:SearchBeam|None

dim

Return type:int|None

feature_dim_or_sparse_dim

Returns:if we have a feature dim, return its dim tag. if we are sparse, return the sparse_dim. otherwise None Return type:Dim|None

sparse

Return type:bool Returns:whether the values represent class indices. see sparse_dim

vocab

Return type:returnn.datasets.util.vocabulary.Vocabulary|None

time_dimension()

Returns:shape(placeholder)[time_dim_axis], int scalar Return type:tf.Tensor

get_dim(axis)

Parameters:axis (int) – counted with batch-dim Returns:shape[axis] Return type:tf.Tensor|int

get_placeholder_as_time_major()

Return type:tf.Tensor

get_placeholder_as_batch_major()

Return type:tf.Tensor

get_placeholder_with_specific_batch_dim_axis(batch_dim_axis)

Parameters:batch_dim_axis (int) – Return type:tf.Tensor

get_placeholder_with_runtime_sanity_checks()

Returns:identity(self.placeholder) with added checks Return type:tf.Tensor

get_placeholder_time_flattened()

Returns:via flatten_with_seq_len_mask() Return type:tf.Tensor

get_placeholder_flattened(keepdims=False)

Parameters:keepdims (bool) – if set, it will add broadcast dimensions after the flattening behind the first axis Return type:tf.Tensor Returns:placeholder where all dynamic axes are flattened into a single axis. e.g. for the usual case (batch, time, dim), it becomes (batch’|time’, dim), or (batch, time, height, dim) will also become (batch’|time’, dim). with keep_dims, (batch, time, height, dim) will become (batch’|time’, 1, 1, dim).

get_axes(exclude_time=False, exclude_batch=False, exclude_feature=False)

Parameters:

• exclude_time (bool) – will filter out the time-axis
• exclude_batch (bool) – will filter out the batch-axis
• exclude_feature (bool) – will filter out the feature-axis

Returns:

list of axes, like range(len(self.shape)), calculated with batch dim.

Return type:

list[int]

get_axes_from_description(axes, allow_int=<class 'returnn.util.basic.NotSpecified'>)

Parameters:

• axes (int|list[int]|str|typing.Sequence[str|Dim]|Dim|None) – one axis or multiple axis, or none. This is counted with batch-dim, which by default is axis 0 (see enforce_batch_dim_axis). It also accepts the special tokens “B”|”batch”, “spatial”, “spatial_except_time”, or “F”|”feature”, and more (see the code).
• allow_int (bool|NotSpecified) – whether to allow an int directly. in almost all cases, it is better to use a symbolic name to specify an axis, as different layers could reorder them, and maybe also change their behavior in the future.

Returns:

list of axes, counted with batch-dim

Return type:

list[int]

get_dim_tag_from_description(axis)

Parameters:axis (str|Dim) – Returns:our matching dim tag. this assumes it exists. Return type:Dim

get_axis_from_description(axis, allow_int=<class 'returnn.util.basic.NotSpecified'>)

Parameters:

• axis (int|str|Dim) –
• allow_int (bool|NotSpecified) –

Returns:

axis, counted with batch-dim

Return type:

int

get_description_from_axis(axis)

Parameters:axis (int) – Returns:some canonical description, such that self.get_axis_from_description(res) == axis. This is quite heuristically for now. We use both strings as also Dim when appropriate. The behavior could potentially change in the future, also the condition will always hold. Return type:str|Dim

has_axis(axis)

Parameters:axis (str|Dim) – Returns:whether the axis exists Return type:bool

get_axes_by_tag_name(name, spatial_only=False)

Parameters:

• name (str) – the tag name, or part of it (must be unique, and must exist)
• spatial_only (bool) –

Return type:

list[int]

get_axis_by_tag_name(name, spatial_only=False)

Parameters:

• name (str) – the tag name, or part of it (must be unique, and must exist)
• spatial_only (bool) –

Return type:

int

get_batch_axis_excluding_batch(axis)

Parameters:axis (int) – counted with batch-dim Returns:axis counted without batch-dim Return type:int|None

have_dim_tag(tag, include_implicit=True, unique=False)

Parameters:

• tag (Dim) –
• include_implicit (bool) –
• unique (bool) –

Return type:

bool

get_batch_axis(axis)

Parameters:axis (int) – counted without batch-dim Returns:axis counted with batch-dim Return type:int

have_batch_axis()

Return type:bool

have_time_axis()

Return type:bool

have_feature_axis()

Return type:bool

is_time_axis_dynamic()

Returns:whether there are different seq-lens for the time, or all the same (static) Return type:bool

is_axis_dynamic(axis)

Parameters:axis (int) – counted with batch-dim axis Returns:dynamic, i.e. we have it in size_placeholder. Note that this does not perfectly match with get_dynamic_axes(), but more with is_time_axis_dynamic(), although probably in most (all?) cases it should match. If True, you can get the size via get_dynamic_size(). Return type:bool

has_dynamic_size(axis)

Parameters:axis (int) – counted with batch-dim axis. implies that you can call get_dynamic_size(). Return type:bool

get_dynamic_size(axis)

Parameters:axis (int) – counted with batch-dim axis. get_dynamic_size() should be True. Returns:shape (B,) Return type:tf.Tensor

set_dynamic_size(axis, sizes)

Parameters:

• axis (int) – counted with batch-dim
• sizes (tf.Tensor) – shape [B]

get_dynamic_axes()

Returns:list of axes, counted with batch-dim axis (but we exclude the batch dim axis itself) Return type:list[int]

get_static_axes()

Returns:list of axes, counted with batch-dim axis (but we exclude the batch dim axis itself) Return type:list[int]

mark_same_time(tags, must_match=False)

If the given dimension tag matches any of our axes, we set our time axis to the selected one.

Parameters:

• tags (set[Dim]|Dim) –
• must_match (bool) – if True, throw an exception if not found

Returns:

whether we have found the same

Return type:

bool

is_same_time_dim(other)

Checks whether we have a matching/compatible time dim.

Parameters:other (Data) – Return type:bool

get_sequence_lengths()

Returns:seq lens tensor of shape (batch,) of dtype int32. also see get_dynamic_size() Return type:tf.Tensor

get_sequence_mask()

Returns:seq mask of shape (batch,time) if we are batch-major, else (time,batch) if we are time-major Return type:tf.Tensor

get_sequence_mask_broadcast(axis=None)

Parameters:axis (int|None) – Returns:seq mask of shape ((batch,time) or (time,batch)) + (1,)s for remaining dims if BT or TB major, and axis is T or None. In general compatible to placeholder, i.e. same ndim, with broadcast dims. We assert here that the axis is dynamic (is_axis_dynamic()), i.e. we have the size. Return type:tf.Tensor

copy_masked(mask_value)

Parameters:mask_value (float|int) – Return type:Data

get_batch_dim()

Return type:tf.Tensor|int

get_batch_dim_tag()

Return type:Dim

get_static_batch_dim()

Return type:int|None

get_spatial_batch_axes()

Return type:list[int] Returns:list of axes which are not batch axes and not feature or which are time axis or dynamic. counted with batch-dim.

get_spatial_axes()

Return type:list[int] Returns:list of axes which are not feature and batch axes, counted without batch-dim.

get_feature_batch_axes()

Return type:list[int] Returns:list of axes which are feature axes, counted with batch-dim. currently there is only one or zero such axis.

get_feature_axes()

Return type:list[int] Returns:list of axes which are feature axes, counted without batch-dim.

SpecialAxesNames = ('time_dim_axis', 'feature_dim_axis')

get_special_axes_dict(counted_with_batch_dim=True, only_available=False)

Parameters:

• counted_with_batch_dim (bool) –
• only_available (bool) –

Returns:

dict axis-name -> axis

Return type:

dict[str,int]

get_bc_spatial_batch_shape()

Returns:shape which will broadcast along all spatial dimensions and time/batch dim Return type:tuple[int|None]

get_bc_shape(opts=None)

Parameters:opts (dict[Dim|str|list[Dim|str]|tuple[Dim|str],int|str|None]|None) – key specifies the axes. value 1 (‘x’) is broadcasting, -1 (None) is not broadcasting Axes should not be defined multiple times. The default behavior if an axis is not specified is like get_bc_spatial_batch_shape(), i.e. it will broadcast in batch and spatial dims only. Or if “*” is in the dict, this overwrites the default behavior for all axes. Returns:shape where 1 means broadcasting, None or >1 means not broadcasting. can be used for TFUtil.dropout() Return type:tuple[int|None]

get_scope_name()

Returns:via self.placeholder or any self.size_placeholder, or None Return type:str|None

get_full_name()

Returns:if we have a defined scope (via self.get_scope_name()), then scope_name + “/” + self.name, otherwise just self.name Return type:str

get_dim_tag(axis)

Parameters:axis (int) – counted with batch-dim Return type:Dim

get_time_dim_tag()

Return type:Dim

get_size_dim_tag(number)

Parameters:number (int) – index in sorted(size_placeholder.keys()) Return type:Dim

get_batch_shape_dim_tags()

Returns:list of dimension tags, for each axis (counted with batch dim, i.e. len is batch_ndim) Return type:tuple[Dim]

classmethod get_common_data(sources, ignore_feature_dim=False, allow_broadcast_all_sources=<class 'returnn.util.basic.NotSpecified'>, name=None)

Parameters:

• sources (list[Data]) –
• ignore_feature_dim (bool) – when set, the feature dim does not have to match in the sources
• allow_broadcast_all_sources (bool|NotSpecified) –
• name (str|None) –

Returns:

some generic data where the sources should be compatible to (with copy_compatible_to), i.e. it contains the union of all axes from all sources (least common multiple). This is always a template, and a new copy.

Return type:

Data|None

find_matching_dims(dim_tag, is_equal_opts)

Finds the dimensions of this Data that match another Dim

Parameters:

• dim_tag (Dim) –
• is_equal_opts (dict[str,bool]|None) – passed to Dim.is_equal

Return type:

list[int] a list of matching axes, counted with batch dim. Sorted in ascending order

find_matching_dim_map(other, other_axes, is_equal_opts=None)

Looks up all other_axes of another Data in this Data. Does not allow duplicates.

Parameters:

• other (Data) –
• other_axes (list[int]) – a list of axes of other, counted with batch dim
• is_equal_opts (dict[str,bool]|None) – passed to Dim.is_equal

Returns:

a dict mapping other axes to own axes, all counted with batch dim

Return type:

dict[int,int]

class returnn.tf.util.data.ControlFlowContext(kind, outer_ctx=None)

This is a simple wrapper around the TF ControlFlowContext, i.e. tf.while_loop or tf.cond.

We have this wrapper to refer to a context which might not exist yet (e.g. at template construction time). Also, we might want to store additional information, such the spatial dim tag of the loop.

Parameters:

• kind (str) – from ControlFlowContext.Types
• outer_ctx (ControlFlowContext) –

class Types

Possible types of context.

Loop = 'loop'

Cond = 'cond'

repr_inner()

Return type:str

classmethod abs_ctx_stack(ctx)

Parameters:ctx (ControlFlowContext|None) – Return type:list[ControlFlowContext]

classmethod abs_ctx_stack_with_root(ctx)

Parameters:ctx (ControlFlowContext|None) – Return type:list[ControlFlowContext|None] Returns:chain of ctx, last is self, first is None

classmethod is_parent_or_same(parent, child)

Parameters:

• parent (ControlFlowContext|None) –
• child (ControlFlowContext|None) –

Return type:

bool

classmethod collect_parent_dims(ctx)

Parameters:ctx (ControlFlowContext|None) – Return type:list[Dim]

is_loop()

Return type:bool

is_cond()

Return type:bool

outer_ctx

Return type:ControlFlowContext|None

tf_control_flow_ctx

Return type:tensorflow.python.ops.control_flow_ops.ControlFlowContext|None

loop_spatial_dim

Return type:Dim|None