Repository: rwth-i6/returnn

Issue:
Wanted behavior for `deepcopy(dim_tag)`?
Do we want:
```
tag == deepcopy(tag)
```
or not?

In returnn_common (and similarly in PyTorch), it is common do do sth like this:
```
class MultiLayer(Module):
  def __init__(self, out_dim: nn.Dim, num_layers: int, layer: Optional[Module] = None):
    if not layer:
      layer = nn.Linear(out_dim)
    self.layers = nn.Sequential(deepcopy(layer) for _ in range(num_layers)
```

That `deepcopy` on such layer has a reference to the `out_dim`.

The current behavior has `tag != deepcopy(tag)`, although I think this is actually buggy for other reasons.

Related is also how we actually define the equality of dim tags (#634). Although, maybe not so much, as we can keep the `deepcopy` logic as we want it.

I'm somewhat undecided about what behavior we want, or we expect. I tend to prefer `tag == deepcopy(tag)`, also because this simplifies the code in returnn_common.



Discussion:
@Zettelkasten opinions?

To get `tag == deepcopy(tag)`, I think we can simply add this to `Dim`:
```
def __copy__(self):
  return self

def __deepcopy__(self, memo):
  return self
```

Although I'm not really sure if this can cause other problems in other situations. Maybe for pickling or so.

Although for textual serialization of dim tags, like what we need for i6 experiments (which just does a pretty print of the net dict), we anyway need to handle dim tags a bit special, because `__repr__` does not work, even if we adopt it to be valid Python code because we don't want to create the same instance multiple times. So we need some external separate cache of dim tags, have somehow unique identifiers, and look them up in the cache. But I think all of this is unrelated now.

Hm I'm also not sure.
Would it be too bad if we had some `dim.copy()` with a different behavior than `dim.__copy__()`?
I think if I explicitly call `dim.copy()` in a config (e.g. to do an outer product), I'd expect it to be different to the old one.
But, at the same time, it would be super confusing to have two dims with the same name, so I would always rename it when I copy it like this (i.e. have some `dim.copy(dim.description + '-other')` or so).

Otherwise I definitely agree that `deepcopy(module)` is a very valid use case where you do not want the dim tag to be changed.

What do you think of having `__copy__` and `__deepcopy__` such that it does nothing, but to add `Dim.copy(description: str)` which actually returns a copy (and furthermore asserts that you do not create two dim tags with the same description).
Would that be too confusing?
> Hm I'm also not sure. Would it be too bad if we had some `dim.copy()` with a different behavior than `dim.__copy__()`? I think if I explicitly call `dim.copy()` in a config (e.g. to do an outer product), I'd expect it to be different to the old one. But, at the same time, it would be super confusing to have two dims with the same name, so I would always rename it when I copy it like this (i.e. have some `dim.copy(dim.description + '-other')` or so).

Yes exactly. Having a separate explicit `copy` function would make it clear that this is really copying a separate instance (otherwise there is anyway no point in copying it). Maybe it would have a mandatory `description` argument because as you say, it would be confusing with the same description. But also, I cannot really think of any use case of such `copy` function right now anyway.

> What do you think of having `__copy__` and `__deepcopy__` such that it does nothing, but to add `Dim.copy(description: str)` which actually returns a copy (and furthermore asserts that you do not create two dim tags with the same description). Would that be too confusing?

I think this is fine. But actually we don't really need the `copy` function for now. I think I would leave that away, unless we really need it at some point. In `__copy__` and `__deepcopy__`, there should be a good comment and explanation on the behavior of course, maybe just with a link to this issue.

Okay, sure, let's add `copy` separately.
I actually could use `dim_other = dim.copy(dim.description + 'other')` in my config though, I have it e.g. here where I want to create a diagonal mask where `i < j`, I do this:
```
other_window_offset_dim = Dim(
  kind=Dim.Types.Spatial, description="other-window-offset", dimension=window_offset_dim.dimension)
```
and then range into both `window_offset_dim` and `other_window_offset_dim` before comparing the two.
Actually we already have `Dim.copy`:
```
def copy(self, kind=None):
  """
  :param Entity|None kind: if set, overwrites self.kind
  :return: copy, maybe as new kind. setting same_as to self
  :rtype: Dim
  """
  tag = Dim(
    kind=kind or self.kind, description=self.description,
    dimension=self.dimension, dyn_size_ext=self.dyn_size_ext,
    batch=self.batch,
    src_data=self.src_data, src_axis=self.src_axis)
  tag.same_as = self  # not declare_same_as, none of the extra checks needed
  tag._same_as_tb = traceback.extract_stack()
  return tag
```
I.e. not as you suggested, it sets `same_as = self`.

It is only used in one single place though, in `Data.copy_add_dim_by_tag`:
```
# Note: if dim_tag is feature, but we are sparse, we just make it spatial
if self.sparse and dim_tag.is_feature_dim():
  dim_tag = dim_tag.copy(kind=Dim.Types.Spatial)
```
(Not sure if this is really needed anyway...)

Maybe it should have a new mandatory option `same_as_self: bool` to avoid any confusion.


Resolving patch:
diff --git a/returnn/tf/util/data.py b/returnn/tf/util/data.py
--- a/returnn/tf/util/data.py
+++ b/returnn/tf/util/data.py
@@ -145,19 +145,49 @@ def short_repr(self):
         desc += "{ctx=%s}" % self.control_flow_ctx.repr_inner()
     return desc
 
-  def copy(self, kind=None):
+  def __copy__(self):
     """
+    Normally we would not want to get a new tag with ``tag != copy(tag)``.
+    https://github.com/rwth-i6/returnn/issues/860
+
+    See :func:`Dim.copy` if you explicitly want a copy.
+
+    :return: self
+    :rtype: Dim
+    """
+    return self
+
+  def __deepcopy__(self, memo=None):
+    """
+    Normally we would not want to get a new tag with ``tag != deepcopy(tag)``.
+    https://github.com/rwth-i6/returnn/issues/860
+
+    See :func:`Dim.copy` if you explicitly want a copy.
+
+    :param memo:
+    :return: self
+    :rtype: Dim
+    """
+    return self
+
+  def copy(self, same_as_self, description=None, kind=None):
+    """
+    :param bool same_as_self:
+    :param str|None description: new description
     :param Entity|None kind: if set, overwrites self.kind
     :return: copy, maybe as new kind. setting same_as to self
     :rtype: Dim
     """
+    if not same_as_self:
+      assert description is not None, "%s copy with not same_as_self should have a new description" % self
     tag = Dim(
-      kind=kind or self.kind, description=self.description,
+      kind=kind or self.kind, description=description or self.description,
       dimension=self.dimension, dyn_size_ext=self.dyn_size_ext,
       batch=self.batch,
       src_data=self.src_data, src_axis=self.src_axis)
-    tag.same_as = self  # not declare_same_as, none of the extra checks needed
-    tag._same_as_tb = traceback.extract_stack()
+    if same_as_self:
+      tag.same_as = self  # not declare_same_as, none of the extra checks needed
+      tag._same_as_tb = traceback.extract_stack()
     return tag
 
   def _can_use_in_ctx(self, ctx):
@@ -3203,7 +3233,7 @@ def copy_add_dim_by_tag(self, dim_tag, unbroadcast=False, axis=None):
     data_opts = self.get_kwargs()
     # Note: if dim_tag is feature, but we are sparse, we just make it spatial
     if self.sparse and dim_tag.is_feature_dim():
-      dim_tag = dim_tag.copy(kind=Dim.Types.Spatial)
+      dim_tag = dim_tag.copy(same_as_self=True, kind=Dim.Types.Spatial)
     if not unbroadcast and dim_tag.dimension != 1:
       dim_tag = Dim(
         kind=dim_tag.kind, description="%s_dummy_dim1" % (dim_tag.description or "unnamed"), dimension=1)


Test patch:
diff --git a/tests/test_TFUtil.py b/tests/test_TFUtil.py
--- a/tests/test_TFUtil.py
+++ b/tests/test_TFUtil.py
@@ -1138,6 +1138,14 @@ def test_Data_copy_move_axis_time_to_end():
   assert d2.shape == (None, 4, None) and d2.feature_dim_axis == 2 and d2.time_dim_axis == 3
 
 
+def test_Dim_copy():
+  # https://github.com/rwth-i6/returnn/issues/860
+  import copy
+  a = SpatialDim("a")
+  assert a == copy.copy(a)
+  assert a == copy.deepcopy(a)
+
+
 def test_dim_math_basics():
   a = SpatialDim("a")
   b = SpatialDim("b")