Repository: rwth-i6/returnn

Issue:
Generalize `DotLayer` to allow sparse inputs
This basically would unify `GatherLayer` and `DotLayer`.

This is the same logic as for `LinearLayer`, which also works with either dense or sparse input.

I'm not really suggesting to do this, I just want to raise the question here, whether this makes sense.
Maybe it is more clean however to not do this, and make it more explicit for the user.
Maybe it also was a mistake that `LinearLayer` does this. Other frameworks usually have a separate `Embedding` or `EmbeddingLayer` for the case of sparse inputs.



Discussion:
> Maybe it also was a mistake that `LinearLayer` does this. Other frameworks usually have a separate `Embedding` or `EmbeddingLayer` for the case of sparse inputs.

Why do you think it was a mistake to handle both, did you ever run into issues with this?
Are there any dangerous things that can happen here?
I at first found an `Embedding` layer a bit weird, mainly because of its name (I think `SparseLookup` or so would be a better name, this describes much better what it does. That it is an embedding is just a special interpretation).

If you think about this sparse dim as just another (implicit) dim in the shape, I think it makes a lot of sense. Just like we would probably also add these sparse dims to `out_shape` in some way (#706).
So I would support sparse inputs whenever possible (so also here for `DotLayer`). Also for https://github.com/rwth-i6/returnn_common/issues/38#issuecomment-963024728 I would just have one method that accepts both.

> > Maybe it also was a mistake that `LinearLayer` does this. Other frameworks usually have a separate `Embedding` or `EmbeddingLayer` for the case of sparse inputs.
> 
> Why do you think it was a mistake to handle both, did you ever run into issues with this? Are there any dangerous things that can happen here?

No. I thought it might be maybe unintuitive. But maybe not. Or maybe only for people coming from other frameworks.

> I at first found an `Embedding` layer a bit weird, mainly because of its name (I think `SparseLookup` or so would be a better name, this describes much better what it does. That it is an embedding is just a special interpretation).

Well, I think `Embedding` makes more sense, because it represents an embedding and it implies that it is trainable, it contains the embedding matrix.
`SparseLookup` sounds to me like it would do just a lookup, i.e. an alias to `tf.gather`. I.e. purely functional, not trainable. It would not contain the embedding matrix.

> If you think about this sparse dim as just another (implicit) dim in the shape, I think it makes a lot of sense. Just like we would probably also add these sparse dims to `out_shape` in some way (#706). So I would support sparse inputs whenever possible (so also here for `DotLayer`). Also for [rwth-i6/returnn_common#38 (comment)](https://github.com/rwth-i6/returnn_common/issues/38#issuecomment-963024728) I would just have one method that accepts both.

Yes, I see this viewpoint.

But then `GatherLayer` is also basically obsolete, right?

`GatherLayer` doesn't require the `sparse` flag to be set though. But this is a detail, and maybe the `sparse` flag should have been set actually when it is not. For cases where it is not, you would just use `ReinterpretDataLayer`.

Although it might be a bit unintuitive for the user to realize when `sparse` is actually set and when it is not?

Hm okay, I do see some dangers with sparse tensors now. E.g. if you have a sparse `[B,T]` which is then actually virtually `[B,T,F]`. and then you e.g. reduce the `T` axis, I think it will reduce on the vocab ids currently. But instead it should make a dense `[B,T,F]` where `F` is just the sum of the one-hot vectors before.
And probably also in many other cases. Like also `DotLayer` here.
> But then `GatherLayer` is also basically obsolete, right?
> 
> `GatherLayer` doesn't require the `sparse` flag to be set though. But this is a detail, and maybe the `sparse` flag should have been set actually when it is not. For cases where it is not, you would just use `ReinterpretDataLayer`.
> 
> Although it might be a bit unintuitive for the user to realize when `sparse` is actually set and when it is not?

Hm, yeah, I don't know.
I usually use the indices I pass to `GatherLayer` as positions which have a fixed order. So e.g. I would average + round them.
So here these indices tensors need to be dense, otherwise this should operate on one-hot vectors (which I think it currently does not, but I would argue that that is broken).
So in my case, I would always need to explicitly "mark it as sparse" if I wanted to gather using `DotLayer`.
Maybe that could be the main differentiation, that `GatherLayer` treats the inputs really as "indices", not so much as "one hot vectors" like the sparse flag would do?

(in the code, we can of course clean this up and for both `GatherLayer` and `DotLayer` use the same code)
> Hm okay, I do see some dangers with sparse tensors now. E.g. if you have a sparse `[B,T]` which is then actually virtually `[B,T,F]`. and then you e.g. reduce the `T` axis, I think it will reduce on the vocab ids currently. But instead it should make a dense `[B,T,F]` where `F` is just the sum of the one-hot vectors before. And probably also in many other cases. Like also `DotLayer` here.

Is this only for sparse tensors interpreted as implicit dense tensors, but not for expansions because of implicit dims via `dyn_size_ext`? (#721)

But it's also a matter of interpretation and definition, what the `ReduceLayer` should do in this case.

> Is this only for sparse tensors interpreted as implicit dense tensors, but not for expansions because of implicit dims via `dyn_size_ext`? (#721)

Ah true, it's actually a quite similar issue as #721. In both cases you need to make all those implicit dims explicit.

> But it's also a matter of interpretation and definition, what the `ReduceLayer` should do in this case.

My interpretation was that sparse dims should always behave like they are implicitly dense.
But I also never used sparse dims before apart from embedding/specifying loss targets. We can discuss this (maybe in another issue, this is now much more general than DotLayer).
> So in my case, I would always need to explicitly "mark it as sparse" if I wanted to gather using `DotLayer`.
> Maybe that could be the main differentiation, that `GatherLayer` treats the inputs really as "indices", not so much as "one hot vectors" like the sparse flag would do?

There is some danger in this for `DotLayer` when there are situations to the user when it is not obvious whether the some output is sparse or dense (whether the `sparse` flag is set or not). Although when `reduce` is explicitly specified, maybe not so much. And multiplying int with float would also throw an error.

Note that many (almost all?) layers really ignore the `sparse` flag. So it's not at all consistent that layers use this consistently correct to interpret it as a one-hot dense tensor. `LinearLayer` is one of the rare exceptions. And the losses (`CrossEntropyLoss` etc) also handle that because it's very common there to have sparse targets.

So this viewpoint of interpreting it as a one-hot dense tensor, if we consider all the layers buggy which don't do that correctly, and if we would "fix" this, this could potentially break a lot of configs. And maybe this viewpoint is not always useful, except in some cases. I'm not sure...

When we have #706 and the user somehow specifies the implicit dimensions more explicitly, then this is maybe less an issue.
Although, to add, it is also rare that we actually have tensors with `sparse=True` at all, except for targets. So maybe by now fixing or changing some layers behavior on this, we would not really break anything.
> Although, to add, it is also rare that we actually have tensors with `sparse=True` at all, except for targets. So maybe by now fixing or changing some layers behavior on this, we would not really break anything.

Extending on this thought, and getting back to the issue, so you say we should do this? I tend to agree.

Ok, no further comments here, so the conclusion: This should be implemented.

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
@@ -6277,6 +6277,16 @@ class DotLayer(LayerBase):
   I, J, K can also be multiple axes from the sources.
   The var-dims don't need to exist.
   All other axes (shared...) are expected to match.
+
+  You should try to avoid having the same dims in both sources when they are not reduced
+  such that you would end up having some dim twice in the output, e.g. (shared..., I, I).
+  You should avoid this because the dim order should never matter
+  (https://github.com/rwth-i6/returnn/wiki/RETURNN-principles).
+  If you need to perform such an operation, you can use :class:`ReinterpretDataLayer`
+  to introduce a new dim tag.
+
+  The reduce dim can also be the sparse dim of one of the sources.
+  In this case, it behaves like :class:`GatherLayer`.
   """
   layer_class = "dot"
 
@@ -6332,9 +6342,14 @@ def __init__(self,
     b_out = self.sources[1].output.copy()
     a_reduce_axes = a_out.get_axes_from_description(red1, allow_int=axis_desc_allow_int)
     b_reduce_axes = b_out.get_axes_from_description(red2, allow_int=axis_desc_allow_int)
-    assert a_reduce_axes and b_reduce_axes, "%s: sources %r, red1 %r, red2 %r" % (self, self.sources, red1, red2)
-    assert len(a_reduce_axes) == len(b_reduce_axes), (
-      "%s: sources %r, red1 %r, red2 %r, reduce axes must match in count" % (self, self.sources, red1, red2))
+    if red1 == a_out.sparse_dim:
+      assert len(b_reduce_axes) == 1 and len(a_reduce_axes) == 0
+    elif red2 == b_out.sparse_dim:
+      assert len(a_reduce_axes) == 1 and len(b_reduce_axes) == 0
+    else:
+      assert a_reduce_axes and b_reduce_axes, "%s: sources %r, red1 %r, red2 %r" % (self, self.sources, red1, red2)
+      assert len(a_reduce_axes) == len(b_reduce_axes), (
+        "%s: sources %r, red1 %r, red2 %r, reduce axes must match in count" % (self, self.sources, red1, red2))
     if BehaviorVersion.get() >= 3 and (var1 is NotSpecified or var2 is NotSpecified):
       assert var1 is NotSpecified and var2 is NotSpecified
       a_var_axes, b_var_axes = self._auto_var_axes(a_out, b_out, a_reduce_axes, b_reduce_axes)
@@ -6354,8 +6369,8 @@ def __init__(self,
     assert all(b_axis in map_a_to_b_rem_axes.values() for b_axis in b_rem_axes)
     b_rem_axes = [map_a_to_b_rem_axes[a_axis] for a_axis in a_rem_axes]
 
-    transpose_a = bool(a_var_axes and a_reduce_axes[0] < a_var_axes[0])
-    transpose_b = bool(b_var_axes and b_reduce_axes[0] > b_var_axes[0])
+    transpose_a = bool(a_var_axes and a_reduce_axes and a_reduce_axes[0] < a_var_axes[0])
+    transpose_b = bool(b_var_axes and b_reduce_axes and b_reduce_axes[0] > b_var_axes[0])
     # For A, if not transpose_a, we must reorder the axes as: a_rem_axes + a_var_axes + a_reduce_axes.
     # For A, if transpose_a, we must reorder the axes as: a_rem_axes + a_reduce_axes + a_var_axes.
     # For B, if not transpose_b, we must reorder the axes as: b_rem_axes + b_reduce_axes + b_var_axes.
@@ -6378,16 +6393,17 @@ def __init__(self,
     b_var_dims = [b_shape[i] for i in b_var_axes]
     a_reduce_dims = [a_shape[i] for i in a_reduce_axes]
     b_reduce_dims = [b_shape[i] for i in b_reduce_axes]
-    assert len(a_reduce_axes) == len(b_reduce_axes)
-    assert all([
-      a_out.dim_tags[i1] == b_out.dim_tags[i2] or d1 == d2
-      for (d1, d2, i1, i2) in zip(a_reduce_dims, b_reduce_dims, a_reduce_axes, b_reduce_axes)])
+    assert len(a_reduce_axes) == len(b_reduce_axes) or not a_reduce_axes or not b_reduce_axes
+    if len(a_reduce_axes) == len(b_reduce_axes):
+      assert all([
+        a_out.dim_tags[i1] == b_out.dim_tags[i2] or d1 == d2
+        for (d1, d2, i1, i2) in zip(a_reduce_dims, b_reduce_dims, a_reduce_axes, b_reduce_axes)])
     a_var_dim = prod(a_var_dims)
     b_var_dim = prod(b_var_dims)
     a_reduce_dyn_axes = [i for i in a_reduce_axes if a_out.batch_shape[i] is None]
     b_reduce_dyn_axes = [i for i in b_reduce_axes if b_out.batch_shape[i] is None]
-    assert len(a_reduce_dyn_axes) == len(b_reduce_dyn_axes)
-    if a_reduce_dyn_axes:
+    assert len(a_reduce_dyn_axes) == len(b_reduce_dyn_axes) or not a_reduce_axes or not b_reduce_axes
+    if a_reduce_dyn_axes and b_reduce_dyn_axes:
       a_pad, b_pad = get_padding_info_dict_ref(a), get_padding_info_dict_ref(b)
       a_pad_values = [a_pad.get(a_out.dim_tags[i], None) for i in a_reduce_dyn_axes]
       b_pad_values = [b_pad.get(b_out.dim_tags[i], None) for i in b_reduce_dyn_axes]
@@ -6413,8 +6429,8 @@ def __init__(self,
           self._info_reduce_mask = "mask-source-1"
     else:
       self._info_reduce_mask = "none-dynamic"
-    a_reduce_dim = prod(a_reduce_dims)
-    b_reduce_dim = prod(b_reduce_dims)
+    a_reduce_dim = prod(a_reduce_dims) if a_reduce_dims else None
+    b_reduce_dim = prod(b_reduce_dims) if b_reduce_dims else None
     if debug:
       print("%s, red1=%r, red2=%r, var1=%r, var2=%r:" % (self, red1, red2, var1, var2), file=log.v3)
       print(" ", "a:", a_out, a, file=log.v3)
@@ -6431,23 +6447,43 @@ def __init__(self,
             tf_compat.v1.assert_equal(
               a_reduce_dim, b_reduce_dim, data=[a_shape, b_shape, a_reduce_dims, b_reduce_dims], summarize=100)]):
         a = tf.identity(a)
-    if not transpose_a:
-      a = tf.transpose(a, a_rem_axes + a_var_axes + a_reduce_axes)
-      a = tf.reshape(a, a_rem_dims + [a_var_dim, a_reduce_dim])
+    if a_reduce_axes:
+      if not transpose_a:
+        a = tf.transpose(a, a_rem_axes + a_var_axes + a_reduce_axes)
+        a = tf.reshape(a, a_rem_dims + [a_var_dim, a_reduce_dim])
+      else:
+        a = tf.transpose(a, a_rem_axes + a_reduce_axes + a_var_axes)
+        a = tf.reshape(a, a_rem_dims + [a_reduce_dim, a_var_dim])
     else:
-      a = tf.transpose(a, a_rem_axes + a_reduce_axes + a_var_axes)
-      a = tf.reshape(a, a_rem_dims + [a_reduce_dim, a_var_dim])
-    if not transpose_b:
-      b = tf.transpose(b, b_rem_axes + b_reduce_axes + b_var_axes)
-      b = tf.reshape(b, b_rem_dims + [b_reduce_dim, b_var_dim])
+      a = tf.transpose(a, a_rem_axes + a_var_axes)  # no reduce axis
+      a = tf.reshape(a, a_rem_dims + [a_var_dim])
+    if b_reduce_axes:
+      if not transpose_b:
+        b = tf.transpose(b, b_rem_axes + b_reduce_axes + b_var_axes)
+        b = tf.reshape(b, b_rem_dims + [b_reduce_dim, b_var_dim])
+      else:
+        b = tf.transpose(b, b_rem_axes + b_var_axes + b_reduce_axes)
+        b = tf.reshape(b, b_rem_dims + [b_var_dim, b_reduce_dim])
     else:
-      b = tf.transpose(b, b_rem_axes + b_var_axes + b_reduce_axes)
-      b = tf.reshape(b, b_rem_dims + [b_var_dim, b_reduce_dim])
-    # `res` will be of shape: a_rem_dims + [a_var_dim, b_var_dim]
-    res = tf.matmul(a, b, transpose_a=transpose_a, transpose_b=transpose_b)
-    if not b_var_dims and add_var2_if_empty:
-      b_var_dims.append(1)
-      b_var_axes.append(None)
+      b = tf.transpose(b, b_rem_axes + b_var_axes)  # no reduce axis
+      b = tf.reshape(b, b_rem_dims + [b_var_dim])
+    if a_reduce_axes and b_reduce_axes:
+      # `res` will be of shape: a_rem_dims + [a_var_dim, b_var_dim]
+      res = tf.matmul(a, b, transpose_a=transpose_a, transpose_b=transpose_b)
+      if not b_var_dims and add_var2_if_empty:
+        b_var_dims.append(1)
+        b_var_axes.append(None)
+    else:
+      # one of the inputs is sparse
+      assert len(a_rem_dims) == len(b_rem_dims)  # batch dims
+      batch_dims = len(a_rem_dims)
+      if red1 == a_out.sparse_dim:
+        res = tf.gather(params=b, indices=a, batch_dims=batch_dims, axis=batch_dims + (1 if transpose_b else 0))
+      elif red2 == b_out.sparse_dim:
+        res = tf.gather(params=a, indices=b, batch_dims=batch_dims, axis=batch_dims + (0 if transpose_a else 1))
+      else:
+        raise ValueError("%s: unexpected reduce %s, %s for inputs %s" % (self, red1, red2, self.sources))
+      # result_shape(p_shape, i_shape, axis=0) = p_shape[:axis] + i_shape + p_shape[axis+1:]
     res = tf.reshape(res, a_rem_dims + a_var_dims + b_var_dims)
     self.output.placeholder = res
 
@@ -6590,9 +6626,14 @@ def get_out_data_from_opts(cls, name, sources,
     b_out = sources[1].output.copy()
     assert not a_out.beam or not b_out.beam or a_out.beam == b_out.beam
     b_reduce_axes = b_out.get_axes_from_description(red2, allow_int=axis_desc_allow_int)
-    assert a_reduce_axes and b_reduce_axes, "%s: sources %r, red1 %r, red2 %r" % (name, sources, red1, red2)
-    assert len(a_reduce_axes) == len(b_reduce_axes), (
-      "%s: sources %r, red1 %r, red2 %r, reduce axes must match in count" % (name, sources, red1, red2))
+    if red1 == a_out.sparse_dim:
+      assert len(b_reduce_axes) == 1 and len(a_reduce_axes) == 0
+    elif red2 == b_out.sparse_dim:
+      assert len(a_reduce_axes) == 1 and len(b_reduce_axes) == 0
+    else:
+      assert a_reduce_axes and b_reduce_axes, "%s: sources %r, red1 %r, red2 %r" % (name, sources, red1, red2)
+      assert len(a_reduce_axes) == len(b_reduce_axes), (
+        "%s: sources %r, red1 %r, red2 %r, reduce axes must match in count" % (name, sources, red1, red2))
     if BehaviorVersion.get() >= 3 and (var1 is NotSpecified or var2 is NotSpecified):
       assert var1 is NotSpecified and var2 is NotSpecified
       a_var_axes, b_var_axes = cls._auto_var_axes(a_out, b_out, a_reduce_axes, b_reduce_axes)
@@ -6644,7 +6685,7 @@ def find_axis(a_axis, b_axis):
       name="%s_output" % name,
       dim_tags=dim_tags,
       time_dim_axis=time_dim_axis,
-      dtype=a_out.dtype,
+      dtype=a_out.dtype if not a_out.sparse else b_out.dtype,
       batch=BatchInfo.get_common_batch_info([src.batch for src in (a_out, b_out)]),
       beam=SearchBeam.get_combined_beam(a_out.beam, b_out.beam))
 


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -4394,6 +4394,32 @@ def test_DotLayer_self_att_dyn_size_ext():
   DotLayer(network=net, name="dot", sources=[a_lay, b_lay], output=dot, **dot_kwargs)  # just check that it builds.
 
 
+def test_DotLayer_sparse_input():
+  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
+  time_dim = SpatialDim("time")
+  classes_dim = FeatureDim("classes", dimension=5)
+  embed_dim = FeatureDim("embed", dimension=3)
+  config = Config({
+    "extern_data": {
+      "src": {"dim_tags": [batch_dim, time_dim], "sparse_dim": classes_dim},
+      "embed": {"dim_tags": [classes_dim, embed_dim]},
+    },
+    "network": {
+      "output": {"class": "dot", "from": ["data:src", "data:embed"], "reduce": classes_dim},
+    },
+    "debug_runtime_sanity_checks": True,
+  })
+
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    net.construct_from_dict(config.typed_dict["network"])
+    layer = net.get_default_output_layer()
+    assert layer.output.dim_tags == (batch_dim, time_dim, embed_dim)
+
+    feed_dict = make_feed_dict(net.extern_data)
+    session.run(layer.output.placeholder, feed_dict=feed_dict)
+
+
 def test_subnet_load_on_init():
   import tempfile
   model_tmp_dir = tempfile.mkdtemp("tmp-checkpoint")