returnn.frontend.dims

Utilities for dimension tags, dimensions, axes.

returnn.frontend.dims.range_over_dim(dim: Dim, *, dtype: str | None = None, device: str | None = None) → Tensor[T]

Parameters:

• dim

• dtype

:param device, :return: tensor with shape [dim]

returnn.frontend.dims.range_over_dim_strided(dim: Dim, *, stride: int | Tensor, out_dim: Dim | None = None, dtype: str | None = None, device: str | None = None) → Tuple[Tensor[T], Dim]

Parameters:

• dim

• stride

• out_dim

• dtype

:param device, :return: tensor with shape [dim], out_dim

returnn.frontend.dims.range_over_merged_dims(dims: Sequence[Dim], *, dtype: str | None = None, device: str | None = None) → Tensor[T]

This is if you want to index into a merged dim. Related: rf.merge_dims().

Parameters:

• dims

• dtype

• device

Returns:

tensor with shape [dim_0, …, dim_n] -> sparse_dim = merged_dim, where merged_dim = dim_0 * … * dim_n

returnn.frontend.dims.linspace_over_dim(dim: Dim, start: float | Tensor = 0.0, end: float | Tensor = 1.0, *, dtype: str | None = None, device: str | None = None) → Tensor

Linearly spaced values over a dim.

Parameters:

• dim – dim to range over

• start – start value

• end – end value

• dtype – dtype of the output tensor

• device – device of the output tensor

Returns:

tensor with shape [dim] containing linearly spaced values between start and end

returnn.frontend.dims.replace_dim(source: Tensor, *, in_dim: Dim, out_dim: Dim | None = None) → Tuple[Tensor, Dim]

Also see: replace_dim_v2(), rf.merge_dims(), rf.split_dims().

Parameters:

• source

• in_dim – should be in source.dims, to be replaced. If you want to replace the source.sparse_dim, see set_sparse_dim().

• out_dim – If not given, will create a new dim with the same size as in_dim. Note: If the size of out_dim is different from in_dim, currently the dim tag is replaced and there is no error – this is not checked.

Returns:

source with in_dim replaced by out_dim, and out_dim.

returnn.frontend.dims.replace_dim_v2(source: Tensor, *, in_dim: Dim, out_dim: Dim, allow_expand: bool = True, allow_shrink: bool = True) → Tensor

Extends replace_dim() by also allowing to expand or shrink the dim (or rather, to not ignore this; when replace_dim() is used on a dim with different size,

it will ignore this and anyway accept the new dim tag (currently)).

Parameters:

• source

• in_dim – should be in source.dims, to be replaced. If you want to replace the source.sparse_dim, see set_sparse_dim().

• out_dim – should not be in source.dims, to be replaced. Note: In contrast to replace_dim(), you must provide this explicitly.

• allow_expand – if True, allow to expand the dim, i.e. if out_dim.size > in_dim.size.

• allow_shrink – if True, allow to shrink the dim, i.e. if out_dim.size < in_dim.size.

Returns:

source with in_dim replaced by out_dim.

returnn.frontend.dims.set_sparse_dim(source: Tensor, sparse_dim: Dim) → Tensor

Parameters:

• source

• sparse_dim

Returns:

source with sparse_dim set

returnn.frontend.dims.dim_match_priority_when_needed(dim: Dim, *other_dims: Dim) → Dim

Returns:

maybe copy of dim with higher match_priority if needed to distinguish from other_dims

Why or when is this needed?

For activation values, this should never be needed, and all dims should be unique.

In case of self-attention, the standard way is to create a separate distinct dim to perform the attention reduction over. See SelfAttention.

However, in case of weight matrices, it is not unusual to have the same dim for both the input and output, so a square weight matrix. When reduction is performed in matmul(), we want to match the input feature dim to the dim in the weight matrix with higher match priority.

So dim_match_priority_when_needed() would be applied on the input feature dim.

https://github.com/rwth-i6/returnn/pull/871 https://github.com/rwth-i6/returnn_common/issues/17#issuecomment-997463222

returnn.frontend.dims.num_elements_of_shape(dims: Dim | Sequence[Dim], *, use_mask: bool = True, device: str | None = None) → int | Tensor

Parameters:

• dims

• use_mask

• device – only for the case when we return a Tensor. by default, this is CPU (just as the size tensor).

Returns:

num elements of a tensor of shape dims, properly considering masking

returnn.frontend.dims.masked_fraction_of_shape(dims: Dim | Sequence[Dim], *, inverse: bool = False) → int | float | Tensor

Parameters:

• dims

• inverse – if True, return the inverse of the fraction

Returns:

:func:`num_elements_of_shape`(dims) / prod(dims) if not inverse else prod(dims) / num_elements

returnn.frontend.dims.last_frame_position_of_dim(dim: Dim, *, device: str | None = None, allow_scalar_on_cpu: bool = True) → int | Tensor

Parameters:

• dim

• device

• allow_scalar_on_cpu – if device is not given, we would use rf.get_default_device() for dynamic sizes. If this is True, we would allow to return a scalar dynamic size on CPU.

Returns:

last frame position of dim

returnn.frontend.dims.use_mask_default(*, default: bool | None = None, default_false_for_behavior_version_up_to: int | None = None) → bool | None

Check the global RETURNN config for the rf_use_mask on what default we should use for the use_mask argument in various functions (e.g. conv(), pool(), reduce(), matmul(), …).

See issue #1691.

Parameters:

• default – what to return if it is not defined in the config, and default_false_for_behavior_version_up_to does not apply.

• default_false_for_behavior_version_up_to – if it is not defined in the config, and if this is set, and the behavior version is less or equal, then return False by default, i.e. do not use the mask by default, if it is not defined in the config. This takes precedence over default.

Returns:

what to use for the use_mask argument by default