Repository: rwth-i6/returnn

Issue:
Error in RangeInAxisLayer
https://github.com/rwth-i6/returnn/blob/8f5142be4ae7ef1770bbe37db82cc0f95def3674/returnn/tf/layers/basic.py#L2025

I think there is an error in the layer. The description says:
`Assume that the input is e.g. (B,T,D), and you specify axis="T", you will get (B=1,T,D=1),
  where the specified axis is filled with ``tf.range``.`

It first sets: `out = tf.range(0, source_shape[axis], dtype=dtype)` (shape (T,))
Then, if we set `keepdims=True`, it sets the output shape as follows:
`out_shape = [
        source_shape[i]
        if (i == axis or i == self.output.batch_dim_axis)
        else 1
        for i in range(self.output.batch_ndim)]`
In the example, this would yield shape `(B,T,1)`.

Then, in the end, it tries: `out = tf.reshape(out, out_shape)`. However, this won't work because tf doesn't support broadcasting in the reshape function and therefore we cannot reshape `(T,)` to `(B,T,1)`.
In my case, tf raised the error: `Input to reshape is a tensor with 3 values, but the requested shape has 12`, because my range axis has size 3 and my batch axis has size 4.

However, even if it would work, it would not return the shape described in the documentation, namely `(1,T,1)`. 

An easy fix would be to just change the output shape:
`out_shape = [
        source_shape[i]
        if i == axis
        else 1
        for i in range(self.output.batch_ndim)]`

What do you think? @albertz @Zettelkasten 


Discussion:
I think the batch dim axis always needs to be unbroadcasted. I think there is code that relies on this (maybe beam search and so, I'm not sure). E.g. see `Data.copy_add_batch_dim`, that also always unbroadcasts the batch dim.

So maybe we should update the docstring?

Is this actually an issue for you?

Also there is `keepdims=False` you could set. Then the output shape is just `(T)`.
I would usually always use `keepdims=False`. This makes much more sense to me than having these "dummy" axes.
I think those were mainly there because some layers didn't support inputs without a batch dim back then. But most now should do that.
I don't really need it as I can also use the `ExpandDimsLayer` afterwards to re-add the axes with `dim=1` (so that other layers can broadcast the dimensions then). I was just wondering if the layer would work at all if setting `keepdims=True` because I think the reshape operation always fails in this case, right? Except for when we choose the batch axis as the range axis.
You should never need to add the axes.


Robin Schmitt ***@***.***> schrieb am Fr., 10. Sept. 2021,
13:01:

> I don't really need it as I can also use the ExpandDimsLayer afterwards
> to re-add the axes with dim=1 (so that other layers can broadcast the
> dimensions then). I was just wondering if the layer would work at all if
> setting keepdims=True because I think the reshape operation always fails
> in this case, right? Except for when we choose the batch axis as the range
> axis.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/rwth-i6/returnn/issues/638#issuecomment-916818172>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAAON7HNMQBBHIKLTW3HBIDUBHQQHANCNFSM5DY2OC3Q>
> .
> Triage notifications on the go with GitHub Mobile for iOS
> <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675>
> or Android
> <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>.
>
>

> I was just wondering if the layer would work at all if setting keepdims=True because I think the reshape operation always fails in this case, right?

Oh yea, I think you're right. So I guess nobody anyway used `RangeInAxisLayer` with `keepdims=True`? I would then maybe suggest just to remove `keepdims=True` entirely... see #639.
Such broadcast axes should never be needed to be added. RETURNN should always handle that correctly internally. In fact, manual broadcast axes by the user could be even confusing for RETURNN (although there is some logic for it to handle that, e.g in `Data.get_common_data`, or `Data.copy_compatible_to`).

For the same reason, there should never be the need to use `ExpandDimsLayer`.
> Such broadcast axes should never be needed to be added. RETURNN should always handle that correctly internally. In fact, manual broadcast axes by the user could be even confusing for RETURNN (although there is some logic for it to handle that, e.g in `Data.get_common_data`, or `Data.copy_compatible_to`).
> 
> For the same reason, there should never be the need to use `ExpandDimsLayer`.

Okay, I didn't know that RETURNN could do that automatically. I will have a look at that.

Resolving patch:
diff --git a/returnn/tf/layers/basic.py b/returnn/tf/layers/basic.py
--- a/returnn/tf/layers/basic.py
+++ b/returnn/tf/layers/basic.py
@@ -2032,12 +2032,12 @@ class RangeInAxisLayer(LayerBase):
   recurrent = True  # if axis=="T", the time-dim order matters
 
   # noinspection PyUnusedLocal
-  def __init__(self, axis, dtype="int32", unbroadcast=False, keepdims=True, sparse=False, **kwargs):
+  def __init__(self, axis, dtype="int32", unbroadcast=False, keepdims=False, sparse=False, **kwargs):
     """
     :param str axis:
     :param str dtype:
-    :param bool unbroadcast:
-    :param bool keepdims:
+    :param bool unbroadcast: DEPRECATED, unsupported, and not needed
+    :param bool keepdims: DEPRECATED, unsupported, and not needed
     :param bool sparse:
     """
     super(RangeInAxisLayer, self).__init__(**kwargs)
@@ -2048,47 +2048,28 @@ def __init__(self, axis, dtype="int32", unbroadcast=False, keepdims=True, sparse
     source_shape = get_shape(source.placeholder)
     out = tf.range(0, source_shape[axis], dtype=dtype)
     if unbroadcast:
-      assert keepdims
+      raise Exception("%s: do not use unbroadcast")
     if keepdims:
-      out_shape = [
-        source_shape[i]
-        if (i == axis or i == self.output.batch_dim_axis)
-        else 1
-        for i in range(self.output.batch_ndim)]
-      out = tf.reshape(out, out_shape)  # add missing axes (keep_dims)
-      if unbroadcast:
-        out = out + tf.zeros(source_shape, dtype=out.dtype)
+      raise Exception("%s: do not use keepdims")
     self.output.placeholder = out
 
   @classmethod
-  def get_out_data_from_opts(cls,
-                             name, sources, axis, dtype="int32", unbroadcast=False, keepdims=True, sparse=False,
-                             **kwargs):
+  def get_out_data_from_opts(cls, name, sources, axis, dtype="int32", sparse=False, **kwargs):
     """
     :param str name:
     :param list[LayerBase] sources:
     :param str axis:
     :param str dtype:
-    :param bool unbroadcast:
-    :param bool keepdims:
     :param bool sparse:
     """
-    from ..util.data import DimensionTag
     assert len(sources) == 1, "%s layer %r requires single source" % (cls, name)
     source = sources[0].output
     axis = source.get_axis_from_description(axis)
-    if keepdims:
-      data_opts = source.get_kwargs(include_special_axes=True)
-      dim_tags = [
-        tag if (i == axis or tag.is_batch_dim() or unbroadcast)
-        else DimensionTag(kind=tag.kind, description="%s_keep%i" % (name, i), dimension=1)
-        for i, tag in enumerate(source.dim_tags)]
-    else:
-      data_opts = source.get_kwargs(include_special_axes=False)
-      dim_tags = [source.dim_tags[axis]]
-      if not dim_tags[0].is_batch_dim():
-        data_opts.pop("batch", None)
-        data_opts.pop("beam", None)
+    data_opts = source.get_kwargs(include_special_axes=False)
+    dim_tags = [source.dim_tags[axis]]
+    if not dim_tags[0].is_batch_dim():
+      data_opts.pop("batch", None)
+      data_opts.pop("beam", None)
     data_opts["name"] = "%s_output" % name
     data_opts["dim_tags"] = dim_tags
     data_opts["dtype"] = dtype


Test patch: