Extend SliceNdLayer to work with start of shape (B,T) and with dynamic sizes

[robin-p-schmitt]

I would like to extend the SliceNdLayer so that it works if we provide start as a layer of shape (B,T). This is necessary, if start is inside a recurrent unit and is optimized out of the loop. Currently, the layer only works if start is of shape (B,).

Additionally, I would like to make size a layer so that we can extract slices with a dynamic length. Again, this should also work if size is of shape (B,T) because of the reason mentioned above. Currently, the layer only supports size being of type int.

Right now, if start is of shape (B,) and the input is of shape (B,T,...), SliceNdLayer returns shape (B,size,...). The question is, what we should return if start is of shape (B,T). I think it would be useful if it was of shape (B,T,size,...) as it would then be easy to select slices for each time step of the sequences. This works easily if size is a static integer. It gets a bit more complicated, if size also becomes a layer of shape (B,T) because the size placeholder of the size dimension of the output would then depend on both the sequence and the time step, making it 2-dimensional.

I have already tried out some code in my local repository. My current solution for allowing start to be of shape (B,T) is the following: the layer uses the slice_nd function from returnn.tf.util which I assume we don't want to change as other layers might depend on it. The slice_nd function expects start to be of shape (B,) and size to be an integer. Therefore, in the SliceNdLayer I flatten all axes of start into the batch dimension and repeat the entries of the input as many times as to fit the new shape of start. After slice_nd is finished, the returned slices can be reshaped again to match the size (B,T,size,...). This functionality already works.

I also tried out to allow size to be a layer of shape (B,T). Before passing it to the slice_nd function, I extracted the max size to get an integer. After slice_nd is finished, I padded the extracted slices (which are all of the same size) in order to match the sizes from the size layer. I think this should also work in theory, however, so far I am not sure of how to set the dimension tags such that the size placeholder of the size axis is dependent on both the batch and the time dimension. We anyway have to decide if we want to do this or if we want to return another shape than I proposed.

What do you think @albertz ?

[albertz]

@Zettelkasten recently did changes to GatherLayer (GatherNdLayer) which is very similar conceptually. I think it should be consistent to that. And maybe he can also comment here.

[albertz]

Additionally, I would like to make size a layer so that we can extract slices with a dynamic length. Again, this should also work if size is of shape (B,T) because of the reason mentioned above. Currently, the layer only supports size being of type int.

Do you really need that? Only start working on this if you really need that. Don't introduce complexity which is not actually used.

So, do you need size to be a layer?

If yes, what formats of size do you really need to have supported?

• scalar. Should be most simple.
• [B]. Should also be simple.
• [B,T], while T is the same T as in start.
• [B,T2], while T2 is sth different.
• whatever else

Also, when you can separate the changes/extensions to start from those changes/extensions to size, that would be better. Always try to make things (changes) as atomic/minimal as possible (unless that would complicate things a lot). Then first do one PR which extends the support on start. And another separate PR for size (if really needed).

[albertz]

the slice_nd function from returnn.tf.util which I assume we don't want to change as other layers might depend on it.

Surely you can extend such function.

Also, do not make such assumptions which you can easily check. In PyCharm, just rightlick -> "Find usages".

Also, we have a lot of test cases, so it probably would be detected if sth breaks.

[albertz]

Btw, more generic, you should create some simple test cases where you construct the input by hand, and calculate the expected output by hand (or maybe by some simple naive pure Python implementation). And then this should be checked against the output of the layer.

[robin-p-schmitt]

Do you really need that? Only start working on this if you really need that. Don't introduce complexity which is not actually used.

So, in my use case, both the start as well as the size of the slice can potentially change in every time step in some of the sequences. I could also try to avoid the SliceNdLayer and use, e.g., GatherLayer among others.

Edit: Maybe I can try that first and if it gets too complex, we can talk about changing the SliceNdLayer instead.

[robin-p-schmitt]

If yes, what formats of size do you really need to have supported?

• scalar. Should be most simple.
• [B]. Should also be simple.
• [B,T], while T is the same T as in start.
• [B,T2], while T2 is sth different.
• whatever else

I would need size to be of shape (B,T), as each start index belongs to one size.

[albertz]

Btw, I was just checking our implementation of slice_nd. Actually this just used gather_nd under the hood. So maybe for SliceNdLayer, we can simply also call GatherNdLayer internally? Or at least share some code? (Similarly as e.g. GenericAttentionLayer wraps DotLayer internally.) This would probably directly make it as generic as possible, for all the cases you need. (@Zettelkasten opinions?)

[Zettelkasten]

Btw, I was just checking our implementation of slice_nd. Actually this just used gather_nd under the hood. So maybe for SliceNdLayer, we can simply also call GatherNdLayer internally?

I didn't look into the implementation to understand what's going on, but I noticed it used tf.gather_nd. There, the indices have a different format (are of shape [B...,#n-dims]). The format for tf.gather is different. Both GatherLayer and GatherNdLayer use the tf.gather format though (we deprecated GatherNdLayer anyway because of that reason).
So we would need to use the "tf.gather"-like functionality of GatherLayer if we want to use that.
Maybe that's simple to change though, I didn't think further about it.
If it's not a problem, then I totally think sharing the code would be great.

It anyway seems like @robin-p-schmitt needs the most general case here.

[Zettelkasten]

Btw, SliceNdLayer and SliceLayer? Is there a good reason to have both? And not just one, supporting the most general case?
I remember this was also a topic for GatherLayer/GatherNdLayer, and there we ended up deprecating GatherNdLayer.

[albertz]

Btw, SliceNdLayer and SliceLayer? Is there a good reason to have both? And not just one, supporting the most general case?
I remember this was also a topic for GatherLayer/GatherNdLayer, and there we ended up deprecating GatherNdLayer.

Yea. From an API point-of-view, it make sense to have that consistent and unified as well.
W.r.t. the implementation, SliceLayer simply uses slices, which is much more efficient. So when we unify this into a single layer, we should keep the efficient implementation for the simple case.

[robin-p-schmitt]

So we would need to use the "tf.gather"-like functionality of GatherLayer if we want to use that.
Maybe that's simple to change though, I didn't think further about it.
If it's not a problem, then I totally think sharing the code would be great.

I can look at that today, try to modify the SliceNdLayer accordingly and create a draft PR.

[robin-p-schmitt]

So I have been looking at the GatherLayer code and it should be straightforward to handle the case if start from SliceNdLayer is of shape (B,). The problem comes when start is of shape (B,T). Assuming an input of shape (B,T,features), my initial idea was to create an indices tensor of shape (B,T,size) (assuming size is an integer) and then gathering the indices from the time axis of the input, setting the batch_dims parameter of tf.gather to 1. This would then give us the desired output of shape (B,T,size,features).

However, the GatherLayer groups together all common axes of the input and the indices and sets the batch_dims attribute of tf.gather. This means that tf.gather assumes both B and T as batch axes and so I cannot gather from the T axis anymore.

My solution would be to either not use GatherLayer and instead use tf.gather directly like described above or to maybe add an option to GatherLayer to set the batch_dims attribute of tf.gather. Another option could be to copy the time dimension of the input in order to get shape (B,T,T_copy,features) and then use GatherLayer to gather from the T_copy axis (bc B and T are assumed to be batch axes). What do you think @albertz @Zettelkasten ? I guess just adding the batch_dims atribute to the GatherLayer would be the easiest way.

[Zettelkasten]

So I have been looking at the GatherLayer code and it should be straightforward to handle the case if start from SliceNdLayer is of shape (B,). The problem comes when start is of shape (B,T). Assuming an input of shape (B,T,features), my initial idea was to create an indices tensor of shape (B,T,size) (assuming size is an integer) and then gathering the indices from the time axis of the input, setting the batch_dims parameter of tf.gather to 1. This would then give us the desired output of shape (B,T,size,features).

I'm not sure I understand. You want to "slice into" axis T of input with shape (B,T,F)? Then your start tensor should not also have the same T axis, but something else, like some T2. This axis should have an entirely different dim tag. You will be able to rename an axis with ReinterpretDataLayer soon (see #633).
The axes of both tensors can be marked as time_dim_axis, but that should not really mean much. We probably want to remove that in the future anyway (see #586), and entirely rely on dimension tags and their descriptions.

Then, from these inputs, like you said you can construct some indices of shape (B,T2,T_sliced).
Then gather these indices into the T axis of the input (B,T,F), giving shape (B,T2,T_sliced,F). That's what you want, right?

However, the GatherLayer groups together all common axes of the input and the indices and sets the batch_dims attribute of tf.gather. This means that tf.gather assumes both B and T as batch axes and so I cannot gather from the T axis anymore.

I don't think using GatherLayer where the axis we want to gather into is also present in the indices makes a lot of sense.
(But maybe GatherLayer supports this actually, I briefly looked into the code and I don't think the gather axis is counted as common axis. But this is untested / I wouldn't recommend in using it like this.)

My solution would be to either not use GatherLayer and instead use tf.gather directly like described above or to maybe add an option to GatherLayer to set the batch_dims attribute of tf.gather. Another option could be to copy the time dimension of the input in order to get shape (B,T,T_copy,features) and then use GatherLayer to gather from the T_copy axis (bc B and T are assumed to be batch axes). What do you think @albertz @Zettelkasten ? I guess just adding the batch_dims atribute to the GatherLayer would be the easiest way.

I don't think we need to change anything in GatherLayer. Instead, I think we can prepare the inputs as shown before and use GatherLayer. And we should enforce that the "slice" axis of SliceNdLayer (currently always the T axis of the input) is not present in the size/start tensors.

[robin-p-schmitt]

Then, from these inputs, like you said you can construct some indices of shape (B,T2,T_sliced).
Then gather these indices into the T axis of the input (B,T,F), giving shape (B,T2,T_sliced,F). That's what you want, right?

Okay, I didn't know that start/ size and the input shouldn't have any common axes. In this case, you're right and in my case T2 would have the same lengths as T.

I don't think the gather axis is counted as common axis

Actually, I think you're right, I didn't notice that.

And we should enforce that the "slice" axis of SliceNdLayer (currently always the T axis of the input) is not present in the size/start tensors.

Okay, but size and start can have the same axes right? In my case, they would always have the same shape, one describing the starts of the slices and the other defining the length of the slices.

Now concerning the dynamic sizes of the slices. As I mentioned, in my use case, I will have different sizes for each sequence and for each time step. Instead of having an output (B,T2,T_sliced,F), where both T2 and T_sliced are dynamic, we could also have an output of shape (B,T2,size_max,F), such that each slice has the same length (the max length in size) by zero-padding the slices to the same length.

[albertz]

We can also extend the functionality of GatherLayer when needed. Eg an interface similar to DotLayer.