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)[source]

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[source]

Defines possible values for kind.

Unspecified = None[source]
Batch = <batch>[source]
Spatial = <spatial>[source]
Time = <spatial>[source]
Feature = <feature>[source]
Types = (<batch>, <spatial>, <feature>)[source]
short_repr()[source]
Returns:some short repr
Return type:str
copy(same_as_self=True, description=None, kind=None, match_priority=None)[source]
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)[source]
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)[source]
Parameters:
  • batch (BatchInfo) –
  • ctx (ControlFlowContext|None) –
  • dyn_size_ext (Data) –
get_dyn_size_ext_for_batch_ctx(batch, ctx)[source]
Parameters:
  • batch (BatchInfo|None) –
  • ctx (ControlFlowContext|None) –
Return type:

Data|None

dyn_size[source]
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()[source]
Returns:whether this dim tag is of kind batch
Return type:bool
is_feature_dim()[source]
Returns:whether this dim tag is of kind feature
Return type:bool
is_spatial_dim()[source]
Returns:whether this dim tag is of kind spatial
Return type:bool
is_dim_known()[source]
Returns:whether we know the dimension; basically whether this is defined (although not self.undefined is defined slightly differently)
Return type:bool
is_dynamic()[source]
Returns:whether the dim is not static. usually means that it has seq lengths
Return type:bool
can_be_used_as_dim()[source]
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)[source]
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)[source]

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)[source]
Parameters:x (tf.Tensor) – size tensor. has been set before via set_tag_on_size_tensor()
Return type:Dim|None
complete_dyn_size()[source]

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)[source]

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.

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()[source]
Return type:Dim
get_same_derived_base()[source]
Return type:Dim
undefined[source]
Return type:bool
declare_same_as(other)[source]
Parameters:other (Dim) –
classmethod get_existing_tag_from_collection(other, tags, is_equal_opts=None)[source]
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)[source]
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)[source]
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()[source]

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()[source]
Returns:axis split info. see get_param_axes_split_info() and usage (e.g. pretraining)
Return type:list[int|None]
vocab[source]
Return type:returnn.datasets.util.vocabulary.Vocabulary|None
sub_right(other)[source]
Parameters:other (Dim|int) –
Returns:self - other
Return type:Dim
sub_left(other)[source]
Parameters:other (Dim|int) –
Returns:(-other) + self
Return type:Dim
div_left(other)[source]
Parameters:other (Dim|int) –
Return type:Dim
div_right(other)[source]
Parameters:other (Dim|int) –
Return type:Dim
ceildiv_left(other)[source]
Parameters:other (Dim|int) –
Return type:Dim
ceildiv_right(other)[source]
Parameters:other (Dim|int) –
Return type:Dim
class Op(kind, inputs, attribs=None)[source]

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[source]

alias of returnn.tf.util.data.Dim

returnn.tf.util.data.FeatureDim(description, dimension, **kwargs)[source]
Parameters:
  • description (str) –
  • dimension (int|None) –
Return type:

Dim

returnn.tf.util.data.SpatialDim(description, dimension=None, **kwargs)[source]
Parameters:
  • description (str) –
  • dimension (int|None) –
Return type:

Dim

class returnn.tf.util.data.ImplicitSparseDim(tag)[source]

Represents an implicit dim via Data.sparse_dim.

Parameters:tag (Dim) –
class returnn.tf.util.data.ImplicitDynSizeDim(tag)[source]

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)[source]

Represents a dim which might exist or not.

Parameters:tag (Dim) –
exception returnn.tf.util.data.VerifyOutShapeException[source]

Exception via Data.verify_out_shape().

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, dim_tag=None)[source]

Represents a dim with fixed size.

Parameters:
  • size (tf.Tensor|int) –
  • dim_tag (Dim|None) –
short_repr()[source]
Return type:str
class GlobalBatchDim(size, dim_tag=None)[source]

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

Parameters:
  • size (tf.Tensor|int) –
  • dim_tag (Dim|None) –
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.
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
dim[source]
Return type:tf.Tensor|int
static_dim[source]
Return type:int|None
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:
  • beam_size (int) –
  • dependency (SearchBeam|NotSpecified|None) –
  • name (str|None) –
  • _next_frame (SearchBeam|None) –
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:
  • 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)[source]

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'[source]
classmethod from_tensor(x)[source]
Parameters:x (tf.Tensor) –
Return type:Data
classmethod template_from_constant(x, name, dtype=None, shape=None, with_batch_dim=False, sparse_dim=None)[source]
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)[source]

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()[source]
Returns:op which does a couple of runtime sanity checks on the placeholder
Return type:tf.Operation
verify_out_shape(out_shape)[source]
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)[source]
Parameters:with_batch (bool) –
Returns:kwargs for tf.compat.v1.placeholder
Return type:dict[str]
get_axes_with_size()[source]
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)[source]
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)[source]
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)[source]
Parameters:special_axes (bool) – special markers for old-style time_dim_axis and feature_dim_axis
Return type:list[str]
get_compare_key()[source]
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)[source]
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()[source]
Returns:copy of myself with batch_dim_axis == 0
Return type:Data
copy_as_time_major()[source]
Returns:copy of myself with time_dim_axis == 0
Return type:Data
copy_with_batch_dim_axis(batch_dim_axis)[source]
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)[source]
Parameters:time_dim_axis (int) –
Returns:copy of myself with specific time_dim_axis
Return type:Data
copy_transpose(perm)[source]
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)[source]
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)[source]

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()[source]
Return type:Data
Returns:copy of myself in batch-time-major or time-batch-major
copy_with_feature_dim_axis(feature_dim_axis)[source]
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()[source]
Returns:copy of self with batch_dim_axis == 0 and feature_dim_axis == 1
Return type:Data
copy_as_time_batch_major()[source]
Returns:copy of self with batch_dim_axis == 1 and time_dim_axis == 0
Return type:Data
copy_as_batch_spatial_major()[source]
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()[source]
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)[source]
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)[source]
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)[source]
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)[source]
Parameters:dim_tag (Dim) –
Return type:int
copy_add_dim_by_tag(dim_tag, unbroadcast=False, axis=None)[source]
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)[source]

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)[source]

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)[source]
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()[source]
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)[source]
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)[source]
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)[source]
Parameters:axes (list[int]) – counted with batch dim
Returns:copy of myself, with squeezed axes
Return type:Data
copy_template(name=None, dtype=None)[source]
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)[source]
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)[source]
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)[source]
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)[source]
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)[source]

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)[source]
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)[source]
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)[source]
Parameters:
  • new_dim_tags (list[Dim]|tuple[Dim]) –
  • name (str|None) –
  • keep_special_axes (bool) –
Return type:

Data

copy_template_set_ctx(ctx)[source]
Parameters:ctx (ControlFlowContext) –
Returns:new Data instance
Return type:Data
copy_template_unpack_batch()[source]

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)[source]
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[source]
Return type:tuple[Dim]
shape[source]
Returns:shape without batch-dim. e.g. (time,feat) = (None,128)
Return type:tuple[int|None]
batch_shape[source]
Returns:shape with added batch-dim. e.g. (batch,time,feat) = (None,None,128)
Return type:tuple[int|None]
get_batch_shape(batch_dim)[source]
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()[source]
Return type:list[int|tf.Tensor]
have_varying_shape_in_ctx()[source]
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[source]

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

Return type:_SizePlaceholderProxy
shape_dense[source]
Returns:shape with feature dim axis
Return type:tuple[int|None]
batch_shape_dense[source]
Return type:tuple[int|None]
dim_tags_sparse[source]
Returns:dim tags without feature dim axis
Return type:tuple[Dim]
dim_tags_set_implicit_only_wrapped[source]
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[source]
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[source]
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.

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[source]
Return type:int
Returns:ndim counted without batch-dim
ndim_dense[source]
Return type:int
Returns:ndim counted without batch-dim, added by 1 if we are sparse
batch_ndim[source]
Return type:int
Returns:ndim counted with batch-dim
batch_ndim_dense[source]
Return type:int
Returns:ndim counted with batch-dim, added by 1 if we are sparse
is_time_major[source]
Returns:whether this is in time-major format, i.e. (time,batch,…)
Return type:bool
is_batch_major[source]
Returns:whether this is in batch-major format, i.e. (batch,…)
Return type:bool
is_batch_feature_major[source]
Returns:whether this is in batch-feature-major format, i.e. (batch,feature,…) (NC…)
Return type:bool
batch_dim_axis[source]
Returns:batch dim axis, counted with batch-dim
Return type:int|None
feature_dim_axis[source]
Returns:feature dim axis, counted with batch-dim
Return type:int|None
feature_dim_axis_or_unspecified[source]
Returns:feature dim axis, counted with batch-dim. could also be unspecified
Return type:int|None|NotSpecified
time_dim_axis_excluding_batch[source]
Return type:int|None
placeholder[source]
Return type:tf.Tensor|None
batch[source]
Return type:BatchInfo|None
beam[source]
Return type:SearchBeam|None
dim[source]
Return type:int|None
feature_dim_or_sparse_dim[source]
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[source]
Return type:bool
Returns:whether the values represent class indices. see sparse_dim
vocab[source]
Return type:returnn.datasets.util.vocabulary.Vocabulary|None
time_dimension()[source]
Returns:shape(placeholder)[time_dim_axis], int scalar
Return type:tf.Tensor
get_dim(axis)[source]
Parameters:axis (int) – counted with batch-dim
Returns:shape[axis]
Return type:tf.Tensor|int
get_placeholder_as_time_major()[source]
Return type:tf.Tensor
get_placeholder_as_batch_major()[source]
Return type:tf.Tensor
get_placeholder_with_specific_batch_dim_axis(batch_dim_axis)[source]
Parameters:batch_dim_axis (int) –
Return type:tf.Tensor
get_placeholder_with_runtime_sanity_checks()[source]
Returns:identity(self.placeholder) with added checks
Return type:tf.Tensor
get_placeholder_time_flattened()[source]
Returns:via flatten_with_seq_len_mask()
Return type:tf.Tensor
get_placeholder_flattened(keepdims=False)[source]
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)[source]
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'>)[source]
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)[source]
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'>)[source]
Parameters:
  • axis (int|str|Dim) –
  • allow_int (bool|NotSpecified) –
Returns:

axis, counted with batch-dim

Return type:

int

get_description_from_axis(axis)[source]
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)[source]
Parameters:axis (str|Dim) –
Returns:whether the axis exists
Return type:bool
get_axes_by_tag_name(name, spatial_only=False)[source]
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)[source]
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)[source]
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)[source]
Parameters:
  • tag (Dim) –
  • include_implicit (bool) –
  • unique (bool) –
Return type:

bool

get_batch_axis(axis)[source]
Parameters:axis (int) – counted without batch-dim
Returns:axis counted with batch-dim
Return type:int
have_batch_axis()[source]
Return type:bool
have_time_axis()[source]
Return type:bool
have_feature_axis()[source]
Return type:bool
is_time_axis_dynamic()[source]
Returns:whether there are different seq-lens for the time, or all the same (static)
Return type:bool
is_axis_dynamic(axis)[source]
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)[source]
Parameters:axis (int) – counted with batch-dim axis. implies that you can call get_dynamic_size().
Return type:bool
get_dynamic_size(axis)[source]
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)[source]
Parameters:
  • axis (int) – counted with batch-dim
  • sizes (tf.Tensor) – shape [B]
get_dynamic_axes()[source]
Returns:list of axes, counted with batch-dim axis (but we exclude the batch dim axis itself)
Return type:list[int]
get_static_axes()[source]
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)[source]

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)[source]

Checks whether we have a matching/compatible time dim.

Parameters:other (Data) –
Return type:bool
get_sequence_lengths()[source]
Returns:seq lens tensor of shape (batch,) of dtype int32. also see get_dynamic_size()
Return type:tf.Tensor
get_sequence_mask()[source]
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)[source]
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)[source]
Parameters:mask_value (float|int) –
Return type:Data
get_batch_dim()[source]
Return type:tf.Tensor|int
get_batch_dim_tag()[source]
Return type:Dim
get_static_batch_dim()[source]
Return type:int|None
get_spatial_batch_axes()[source]
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()[source]
Return type:list[int]
Returns:list of axes which are not feature and batch axes, counted without batch-dim.
get_feature_batch_axes()[source]
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()[source]
Return type:list[int]
Returns:list of axes which are feature axes, counted without batch-dim.
SpecialAxesNames = ('time_dim_axis', 'feature_dim_axis')[source]
get_special_axes_dict(counted_with_batch_dim=True, only_available=False)[source]
Parameters:
  • counted_with_batch_dim (bool) –
  • only_available (bool) –
Returns:

dict axis-name -> axis

Return type:

dict[str,int]

get_bc_spatial_batch_shape()[source]
Returns:shape which will broadcast along all spatial dimensions and time/batch dim
Return type:tuple[int|None]
get_bc_shape(opts=None)[source]
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()[source]
Returns:via self.placeholder or any self.size_placeholder, or None
Return type:str|None
get_full_name()[source]
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)[source]
Parameters:axis (int) – counted with batch-dim
Return type:Dim
get_time_dim_tag()[source]
Return type:Dim
get_size_dim_tag(number)[source]
Parameters:number (int) – index in sorted(size_placeholder.keys())
Return type:Dim
get_batch_shape_dim_tags()[source]
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)[source]
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)[source]

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)[source]

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)[source]

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:
class Types[source]

Possible types of context.

Loop = 'loop'[source]
Cond = 'cond'[source]
repr_inner()[source]
Return type:str
classmethod abs_ctx_stack(ctx)[source]
Parameters:ctx (ControlFlowContext|None) –
Return type:list[ControlFlowContext]
classmethod abs_ctx_stack_with_root(ctx)[source]
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)[source]
Parameters:
  • parent (ControlFlowContext|None) –
  • child (ControlFlowContext|None) –
Return type:

bool

classmethod collect_parent_dims(ctx)[source]
Parameters:ctx (ControlFlowContext|None) –
Return type:list[Dim]
is_loop()[source]
Return type:bool
is_cond()[source]
Return type:bool
outer_ctx[source]
Return type:ControlFlowContext|None
tf_control_flow_ctx[source]
Return type:tensorflow.python.ops.control_flow_ops.ControlFlowContext|None
loop_spatial_dim[source]
Return type:Dim|None