Move core `nn` functions to RETURNN

[albertz]

As a result of rwth-i6/returnn#1120 and rwth-i6/returnn#1264, we want to move core functions over to RETURNN, as this should become the common compatibility layer for different backends such as TensorFlow and PyTorch.

These are mostly the functions in nn/_generated_layers.py, but this in turn also means to move over some of the other code logic, like the naming logic, etc.

[albertz]

On RETURNN side, the new frontend API is supposed to cover exactly those core nn functions. It already started in development, although it is quite minimal so far, and mostly targets TF low-level and PyTorch low-level.

Moving nn core functions to RETURNN means to use the TF high-level layer-based frontend for Tensor and to implement those frontend API functions, in returnn/tf/frontend_layers.

There is already the compare function with a small inconsistency, which does not break too much code though, so this should be fine.

This also means to rewrite nn.Tensor, to now use the RETURNN Tensor class.

The RETURNN Tensor class would also be extended by all the math ops like __add__, which just call the frontend API, like we have it currently in nn.Tensor.

[albertz]

About replacing nn.Tensor: We need to go through all the current attribs and methods, check for their usages, and make sure the same thing works with RETURNN Tensor. One most obvious question: What would be the raw tensor type? nn.NameCtx? It seems like NameCtx lacks a few things which is in nn.Tensor currently, such as:

• layer_dict
• is_ref
• parent_modules
• extra_dependencies
• remove_unused_cleanup_hooks
• debug_layer
• require_global_access (class attrib)
• Does it handle the subclass nn.Parameter right?
  • Note that it uses an unnamed dummy NameCtx initially and later reassigns that.
  • __copy__/__deepcopy__, modified from base Tensor, returns a real copy for Parameter
  • auxiliary added by Parameter
  • _initial added by Parameter
• Check the other nn.Tensor methods.
  • We can either move them somewhere else, or make them available in the RETURNN Tensor class.
  • get_dependencies
  • _sis_hash
  • _get_name_in_ctx
  • get_abs_name
  • __add__ etc -> add to RETURNN
  • dims_set -> add to RETURNN
  • remaining_dims -> add to RETURNN
  • batch_dim -> remove
  • __copy__ / __deepcopy__, no-op for Tensor -> removed. remove Tensor __copy__ and __deepcopy__ #253
  • mark_as_loss -> handle via train ctx context manager? See Frontend API and PyTorch backend returnn#1120 (comment)
  • mark_as_output -> unclear, see Frontend API and PyTorch backend returnn#1120 (comment)
  • mark_as_default_output -> unclear, see Frontend API and PyTorch backend returnn#1120 (comment)
  • feature_dim used often -> add to RETURNN? unclear. See Tensor.feature_dim and Tensor.feature_dim_axis returnn#1273
  • _assign_parent_name_ctx. see below. moved to NameCtx._auto_assign_parent
  • _replace_by. see below.

• Does it handle the subclass PrevTensorRef right?

Can we move those things over to NameCtx? Let's just try to do that step by step and see.

We need to be careful about cases where tensor.name_ctx is reassigned. This can happen via Tensor._replace_by and NameCtx.move_layer_ref_here.

Tensor._replace_by usage is only in prepare_for_config_serialization, only this code:

                if root_mod_call.layer_ref is not None:
                    assert not self.layer_ref  # not sure. maybe just reset?
                    assert root_mod_call.layer.layer_dict["class"] == "subnetwork"
                    sub_out = root_mod_call.children.pop("output")
                    assert sub_out.layer.layer_dict["class"] == "copy"
                    sub_real_out = sub_out.layer.layer_dict["from"]
                    assert isinstance(sub_real_out, nn.Tensor)
                    # noinspection PyProtectedMember
                    sub_out.layer._replace_by(sub_real_out)
                    # noinspection PyProtectedMember
                    root_mod_call.layer._replace_by(sub_real_out)

Tensor._assign_parent_name_ctx is called in
NameCtx._remove_unused_and_assign_parents_and_handle_subnets. _assign_parent_name_ctx can probably be moved to NameCtx.

NameCtx.move_layer_ref_here is only in PrevTensorRef.assign_new_cur_layer_name_ctx:

    def assign_new_cur_layer_name_ctx(self, cur_layer_name_ctx: nn.NameCtx):
        """
        Changes self.name_ctx to new name_ctx.
        """
        prev_layer_name = f"prev:{cur_layer_name_ctx.name}"
        assert prev_layer_name not in cur_layer_name_ctx.parent.children
        prev_layer_name_ctx = cur_layer_name_ctx.parent.get_child(prev_layer_name)
        prev_layer_name_ctx.move_layer_ref_here(self)
        assert self.name_ctx is prev_layer_name_ctx
        self.cur_layer_name_ctx = cur_layer_name_ctx

And this is only called via Loop.assign()._map_ref_to_name_ctx():

                    prev_ref.assign_new_cur_layer_name_ctx(layer_ref.name_ctx)