Specify dim tags for layers that create new axes

[Zettelkasten]

E.g. CumConcatLayer from #589 allows (even enforces) you to specify dim_tag: DimensionTag for the output axis.

I propose something similar for all other layers that introduce new dims. Basically for all layers introducing a new dynamic seq length here that could be useful.

This is already possible to do using an additional ReinterpretDataLayer (or so, something that lets you specify new dim tags. I don't think this exists yet, but is easy to add). Then you would do:

'x_': {'class': 'split_dims',  ...},
'x': {'class': 'reinterpret_data', 'from': 'x_', 'set_dim_tags': {
  'stag:split_dims0': DimensionTag(...), 'stag:split_dims1': DimensionTag(...)}}

Here, split_dims0 and so on are the tag descriptions that SplitDimsLayer internally sets.
Not only is this solution pretty verbose, it also relies on these internal tag descriptions. As a user, you would need to look them up frequently, and we would need to make sure that they don't change internally.
I would like to avoid this, proposing to add a dim_tags: List[DimensionTag] option directly to SplitDimsLayer.
With other layers, its the same idea.
This also promotes the best-practice that you should always name your axes properly.

Maybe, for some layers, we would even enforce the use of this. Especially for layers that really introduce an entirely new dim (like SplitDimsLayer). But also for others. When enforced, this would need a new behavior version (#508).

(Related to that I would also allow to specify just a string with the description as dim_tag to make this easier to use. Calling the DimensionTag constructor is not that nice in Sisyphus configs.)

---

List of layers which need this

Also see: Operations on dimension tags

Layers (potentially) introducing a new dynamic seq length

List

• ConvLayer (done, out_spatial_dims, out_dim, in_spatial_dims, in_dim, ConvLayer and PoolLayer, in_dim, in_spatial_dims, out_dim, out_spatial_dims #789)
• PoolLayer (done, out_spatial_dims, out_dim, in_spatial_dims, in_dim, ConvLayer and PoolLayer, in_dim, in_spatial_dims, out_dim, out_spatial_dims #789)
• TransposedConvLayer (done, out_spatial_dims, out_dim, in_spatial_dims, in_dim, TransposedConvLayer in_dim, out_dim, in_spatial_dims, out_spatial_dims #791)
• WindowLayer (done, out_spatial_dim and window_dim, WindowLayer, window_dim and out_spatial_dim #776)
• PadLayer (done, out_dims and axes, via PadLayer, out_spatial_dims option #778 and PadLayer, rename out_spatial_dims to out_dims #779)
• ResizeLayer (done, out_dim and axis, via ResizeLayer out_dim option #797)
• PrefixInTimeLayer (done, out_dim and axis, via PrefixInTimeLayer, PostfixInTimeLayer, axis and out_dim options #799)
• PostfixInTimeLayer (done, out_dim and axis, via PrefixInTimeLayer, PostfixInTimeLayer, axis and out_dim options #799)
• SliceLayer (done, out_dim, via SliceLayer, out_dim option #772)
• SliceNdLayer (done, out_spatial_dim or size, via SliceNdLayer, size can be DimensionTag #771, SliceNdLayer, out_spatial_dim option #838)
• MergeDimsLayer (done, out_dim and axes, via MergeDimsLayer, handle out_dim option #785)
• SplitDimsLayer (done, dims and axis, via SplitDimsLayer, dims can be list of DimensionTag #786)
• UnflattenNdLayer (done, out_dims, in_dim, via UnflattenNdLayer, in_dim, out_dims options #802)
• RepeatLayer (done, out_dim and axis, via RepeatLayer, handle out_dim option #803)
• TimeChunkingLayer (done, out_dim and axis, via TimeChunkingLayer axis option, TimeUnChunkingLayer, dim tags #805)
• RemoveLayer (done, out_dim and axis, via RemoveLayer, out_dim and axis options #806)
• RecLayer with search (done, axis, via RecUnstackLayer, declare_rec_time option, axis optional #751)
• EditDistanceTableLayer (done, out_dim, via EditDistanceTableLayer, handle out_dim #807)
• MaskedComputationLayer (done, out_spatial_dim, via MaskedComputationLayer, out_spatial_dim option #811; further discussion: MaskedComputationLayer is violating the principle that the user should not need to think about rec automatic optimization #769)
• CumConcatLayer (done, out_spatial_dim, via CumConcatLayer #589 and CumConcatLayer, out_spatial_dim instead of new_dim option #760)
• RangeFromLengthLayer (done, out_spatial_dim, via allow_broadcast_all_sources, changed default, new behavior version #759)
• RangeLayer (done, out_spatial_dim, via allow_broadcast_all_sources, changed default, new behavior version #759)
• ScatterNdLayer (done, out_spatial_dim, via ScatterNdLayer, out_spatial_dim option #770)

Layers involving some linear transformation on the features

• LinearLayer (done, out_dim and in_dim, via Generic LayerBase in_dim and out_dim options #765 and test_LinearLayer_in_dim_spatial, allow in_dim not feature dim #783)
• ConvLayer (done, out_spatial_dims, out_dim, in_spatial_dims, in_dim, ConvLayer and PoolLayer, in_dim, in_spatial_dims, out_dim, out_spatial_dims #789)
• RecLayer (done, axis, via RecUnstackLayer, declare_rec_time option, axis optional #751)

Layers with fixed output shape

• ConstantLayer (done, shape, via Explicit output dim tags for RandIntLayer, VariableLayer, ConstantLayer #762)
• VariableLayer (done, shape, via Explicit output dim tags for RandIntLayer, VariableLayer, ConstantLayer #762)
• RandIntLayer (done, shape, via Explicit output dim tags for RandIntLayer, VariableLayer, ConstantLayer #762)

Others

• ReduceOutLayer (done, out_dim, via ReduceOutLayer, handle out_dim #808)
• StackLayer (done, out_spatial_dim, via StackLayer, out_spatial_dim option #809)
• SplitLayer (done, out_dims, via SplitLayer out_dims option #761)
• TileLayer (done, out_dims, via TileLayer, out_dims option #810)
• ConcatLayer (done, out_dim, via ConcatLayer, add out_dim #824)
• GatingLayer (done, out_dim, via GatingLayer, handle out_dim #839)

[albertz]

Even for simple layers like LinearLayer, ConvLayer, RecLayer etc where some transformation happens, a new dim tag is introduced for the new feature dim.

[albertz]

We should be a bit more systematic here by specifying the list of layers which should be extended/modified by such feature. This should directly be edited into the main issue description. I will start on this but please extend. Or if you are unsure about some layer, just ask in a comment.

[albertz]

Btw, about Sisyphus, or in general simpler ways to specify new dim tags, and also on the dim tag description:

I'm not sure on any of these.

About Sisyphus: Whatever technical limitation there might be in Sisyphus, this could be fixed. E.g. if the issue is that DimensionTag is not pickle-able, or the __repr__ does not work for serialization, we can surely fix that. So I don't really count that as an argument.

On using DimensionTag.description for identification: Currently we rely on that. We even use it for DimensionTag.is_equal. And we also use it for the stag:... syntax when specifying axes in the config, via Data.get_axes_by_tag_name/Data.get_axis_by_tag_name.

Note my recently added comment on that in DimensionTag.is_equal:

# In principle, we would want to check for identity (self is other).
# We currently use the description because the identity would not be the same
# in case of template construction where a dim tag is once created for a template layer,
# and then later again for the real layer.

I don't like it to rely on a string as mean for identification. This can lead to very strange bugs in rare cases which are probably hard to debug.

Also, the automatic description might not always be unique (sometimes it uses the absolute full layer name, sometimes only the base layer name).

Also, the automatic description might even be changed. It is not a good idea to rely on this. This is also what you describe initially ("internal tag descriptions").

With automatic description, I mean all layers which create new dim tags. So when the proposed changes by this issue are implemented, it would all be explicit, so this is not really a problem anymore.

But anyway, I want to get away from relying on description for the identification. Esp if there is not really a big need for it.

See also #634 for some related discussion.

On simpler ways to write that in the config: I was also thinking about this. This might be a valid point.

I think with the proposed new way to define the network via returnn-common, this becomes probably already much simpler.

Maybe it is also just about the DimensionTag constructor. Writing DimensionTag(kind=DimensionTag.Types.Spatial, ...) might be just too long? Maybe we could simply have a small alias for that, like Spatial(...) or new_spatial(...) which would do the same thing? But is this really helpful, or maybe it obfuscates the code more? We should not just optimize the config for easy writing but also clearness and easy reading (even more so than easy writing).

[albertz]

Related is also to enforce dim tags to be unique (#632).

[albertz]

One alternative to this is to use the DimensionTag.derived_from_tag logic more, together with all dim tags being unique (#632).

E.g. when you used ConvLayer or SplitDimsLayer or so on some input [B,T,D] which used dim tag time_dim = DimensionTag(...), we could allow to specify axis=Derived(time_dim) in other layers. When all tags are unique, this should also be non-ambiguous in most cases (and if not, just throw an error).

[albertz]

Note on your initial example using ReinterpretDataLayer with set_dim_tags: The new dim tag would by intention be different from the old dim tag. If you then only use layer x afterwards, there is no real problem. However, if you mix x_ and x, this is maybe not what was being intended. Also, by intention, the new dim tag does not have derived_from_tag set to the old dim tag because we never want that these are being mixed up (that is the purpose of assigning/creating new dim tags via set_dim_tags).

What you actually want is the same_dim_tags_as logic. Which you can e.g. do via CopyLayer, like:

'x': {'class': 'copy', 'from': 'x_', 'out_type': {'same_dim_tags_as':
  Derived(time_dim): DimensionTag(...)}}}

This would use declare_same_as on the dim tag.

[albertz]

Btw, this is not really either-or. We can have both such Derived(...) logic for matching dim tags and also specify new dim tags explicitly for new layers.

I think we need to discuss each (or most) layers individually. E.g. SplitLayer would be kind of specific. And others as well. So we should also open separate issues for each layer. Maybe you start by doing so only for the layer which you care about for now. I don't think we can anyway do all at once, so we should do it layer-by-layer.

[albertz]

Some further ideas along the line of Derived(...) for specifying axes (via Data.get_axis_from_description):

• Derived(tag: DimensionTag): matches derived tags. This basically sets derived_matches=True for DimensionTag.is_equal.
• SameStaticDim(dim: int): matches dims which have the same static dim
• NewSpatialFromLayer(layer_name: str): We can keep track (automatically) which layer created new dim tags. This would match a newly created (spatial) dim tag.
• FeatureFromLayer(layer_name: str): Gets the feature dim from layers output.

Note that we would enforce in any case that the matching is unique. Once it is not unique, it would throw an error.

For the layers we would probably use the get_layer logic in transform_config_dict. Or we might just assume that they must exist at this point (via other deps). This is still a bit unclear.

[albertz]

Another idea I had (also for returnn-common) was that the user would explicitly specify all output dim tags. So basically replacing what the user anyway writes as comment (like # [B, T, NumHeads, Dim] or so) but as code, as an option to any layer (so LayerBase would handle it), like out_dims={batch_dim, time_dim, num_heads_dim, feature_dim} (a set, without order by intention).

This could not just be used for verification, but we can actually infer some of the options of the layer (e.g. n_out for LinearLayer, etc), and it also can be used to specify the dim tags.

There are some cases where the order is relevant though, such as SplitDimsLayer, where it matters in what order we do the splitting. And there are cases where multiple new spatial dims can be created, e.g. ConvLayer with 2D or 3D conv, where the assignment would be ambiguous. Also, even ConvLayer with 1D conv, it might be ambiguous what is actually the output feature (channel) dim and what is the output spatial (maybe static) dim. So in those cases, we would anyway additionally need to make it explicit with special options and out_dims alone would not be enough.

Or we extend out_dims somehow to make it also non-ambiguous in those cases. But not sure. E.g. for ConvLayer, we could assume that there must be exactly one feature dim in the output, which would make that non-ambiguous. But still not the spatial dims. We anyway also need to fix ConvLayer filter_size and other things to not be a list (which depends on the order of spatial axes) but a dict tag -> filter size (dict[DimensionTag,int]). But this would all be layer-specific again. I'm not sure if maybe out_dims should stay always this same concept and type (i.e. just a set[DimensionTag]).

For ConvLayer specifically, and some other layers, also LinearLayer, we could also allow some dim mapping, like in_to_out_dims: {spatial_in: spatial_out, feature_in: feature_out}.

[Zettelkasten]

One alternative to this is to use the DimensionTag.derived_from_tag logic more, together with all dim tags being unique (#632).

E.g. when you used ConvLayer or SplitDimsLayer or so on some input [B,T,D] which used dim tag time_dim = DimensionTag(...), we could allow to specify axis=Derived(time_dim) in other layers. When all tags are unique, this should also be non-ambiguous in most cases (and if not, just throw an error).

For SplitDimsLayer (because you explicitly mentioned it here), this derived logic would almost never be useful and should fail because there is more than one derived dim (for a non-trivial split), right?
But for many other layers, this makes a lot of sense.
I'm not sure about the other options you proposed apart from Derived, would those really be helpful?

SameStaticDim(dim: int): matches dims which have the same static dim

E.g. this one seems dangerous, because static dims might accidentally match up when the user chooses different hyperparameters.

[albertz]

Also, yet another related thing: It might be helpful to also explicitly specify the input dims or part of it. This is standard in PyTorch, e.g. torch.nn.Linear has in_features, out_features.

This would be partly covered by such in_to_out_dims proposed above.

[albertz]

One alternative to this is to use the DimensionTag.derived_from_tag logic more, together with all dim tags being unique (#632).
E.g. when you used ConvLayer or SplitDimsLayer or so on some input [B,T,D] which used dim tag time_dim = DimensionTag(...), we could allow to specify axis=Derived(time_dim) in other layers. When all tags are unique, this should also be non-ambiguous in most cases (and if not, just throw an error).

For SplitDimsLayer (because you explicitly mentioned it here), this derived logic would almost never be useful and should fail because there is more than one derived dim (for a non-trivial split), right?

No, in those cases there can only be one dynamic dim tag (if at all), and all others must be static. And then it can be derived. So it should always be possible to derive it for SplitDimsLayer. Only not the order.

I'm not sure about the other options you proposed apart from Derived, would those really be helpful?

Currently people often use things like static:0, static:-1 etc, which depends on the order of the static dims. SameStaticDim would instead rely on the dimension, which is already much better.

NewSpatialFromLayer would actually make your proposal obsolete, because instead of introducing new_dim_tag, you could use NewSpatialFromLayer instead.

But as said, this doesn't have to be either-or, we can do both, and also discuss them independently.

SameStaticDim(dim: int): matches dims which have the same static dim

E.g. this one seems dangerous, because static dims might accidentally match up when the user chooses different hyperparameters.

But then it would throw an error, so it cannot accidentally match multiple. And also this probably would not happen so often.