Root module, naming logic API

[albertz]

I think the current naming logic is not straightforward or easy to understand. I don't mean the underlying implementation but the conceptual logic and the API as it is used by the user.

The user currently must explicitly create a new root NameCtx instance.
The "root module" is defined by passing name=root_name_ctx on its first call.

Let's take a step back. What is the naming logic relevant for?

• Final absolute variable names, of our model checkpoint, and also for param sharing (important)
• Layer names (nice to have consistent but not really so important)

The root module defines the absolute parameter names: They are required to be accessible via some attrib chain from the root module. So it is also important that there is only one root module, otherwise there might be conflicting names. There might still be ambiguities which are resolved by the the order of assignment.

So far, this is straightforward. This is also similar as how you would handle it in PyTorch when saving and loading a checkpoint. See #93.

NameCtx corresponds (mostly) to RETURNN layer names. This is why you assign a module call to a NameCtx, and not the module itself. (See the docstring of the naming module for further documentation.)

While it makes sense to keep the layer name hierarchy information (NameCtx) internally, and for other custom layer creation code it might also make sense to expose (absolute or relative) layer names, otherwise this probably should not be the main API for the naming logic. Also, even exposing the layer names could be problematic as there are some optimizations which still might move around layers, so the layer names are only really final at the end. And exposing the layer names as a str is also not really needed. It can just be a nn.Tensor instance and it will automatically derived lazily when constructing the final net dict.

Related older similar issue: Better way to define main network dict (make_root_net_dict) (#44)

---

Proposal

The public API in returnn-common maybe just should expose some function to define the root module. This is probably the only really needed API function. And this is only really needed at the time when we create the net dict.

Another related aspect is to have an explicit root name ctx. In most cases, like most configs or also interactive debugging, this is not really necessary, and it would be fine to just have one default root ctx. We could adopt a similar scheme as in TensorFlow which provides a default graph but you can also easily create your own.

The example from #44 could then look like:

net = Model(...)

time_dim = nn.SpatialDim("time")
in_dim = nn.FeatureDim("input", ...)
data = nn.get_extern_data(nn.Data("data", dim_tags=[nn.batch_dim, time_dim, in_dim]))
out = net(data, axis=time_dim)
assert isinstance(out, nn.Tensor)
out.mark_as_default_output()

config_code = nn.get_returnn_config_serialized(net)

(For Sisyphus or other wrappers, it might make sense to split up get_returnn_config_serialized further into extern data and net dict.)

[albertz]

@JackTemaki @mmz33 Maybe you have some thoughts.

[albertz]

This issue here is about the API but fixing this obviously implies some internal changes. While doing this, it is a good idea to also clean up the logic a bit. E.g. NameCtx has many methods, properties or attribs which are all somewhat related and have partly duplicate or at least similar code.

NameCtx.get_abs_name: Using the NameCtx hierarchy names without any custom logic.
No change needed.

NameCtx.layer_abs_name_scope: This defines the TF abs name scope which this layer (esp VariableLayer) is supposed to have. This is determined via _get_abs_canonical_name. For VariableLayer, it's important that this is well defined and unique via the root module attrib chain (and it should error if this is not the case). For other cases, esp pure functions where there is no parent module, it doesn't matter much.
Change This is unknown without knowing the root module. We can either redefine it to be an "assumed" layer abs name scope based on NameCtx info. Or remove it. Some other code uses it to get a abs name, which can be replaced. Edit Now removed.

NameCtx.layer_abs_name_scope_effective: The actual TF abs name scope of this layer, determined via all parent name ctx and the "name_scope" of each layer, unless some name ctx is marked as virtual, or unless custom_layer_name_scope is set for some name ctx.
Change We want to set "name_scope" only at the end when we really know the root module. So this function must be (re)moved. Edit Now removed.

NameCtx.custom_layer_name_scope: Used by layer_abs_name_scope_effective. Set by Cond, Loop and MaskedComputation to the empty string "", or set by _assign_layer_name_scope.
Change Should be cleaned up.

NameCtx._get_abs_canonical_name: Used by layer_abs_name_scope. Via module.parents_with_attr of self.module and potential parent module, search for any calls (i.e. parent NameCtx) and use layer_abs_name_scope as base, or search for root module (corresponding to root name ctx). Uses a queue for depth-first search.
Change Not needed anymore? New logic in _NamePathCache.get_name_path.

NameCtx._get_suggested_name: Used by _get_unique_name when there is no suggested name. Check parent name scope module, any attrib from there to self.module. Do a depth-first search through the parents via module.parents_with_attr, starting from self.module, until we find self.parent.module.
Change Cleaned up, uses _NamePathCache internally now.

NameCtx._get_unique_name: Used by __init__ or assign_parent to assign the name. Without suggested name, calls _get_suggested_name. Checks self.parent.module attribs for exact match. Makes sure name is unique.
No change needed.

NameCtx._assign_layer_name_scope: Sets "name_scope" in a layer dict (and custom_layer_name_scope) such that layer_abs_name_scope_effective matches layer_abs_name_scope.
Change Move this logic.

Tensor._assign_parent_name_ctx (renamed from Tensor._assign_parent): Used by Tensor.get_name_in_ctx (i.e. only at later point). Via parent_modules, find parent name ctx from their calls, and assigns this parent to self.name_ctx. Also potentially assigns custom_layer_name_scope and "name_scope".
Change Cleanup (move) the custom_layer_name_scope and "name_scope" logic.

Again, as mentioned, much of this should maybe be rewritten.

As we only know the root module at some later point (usually at least; or maybe require it earlier?), we can ultimately also define some parts only lazily at that late point (or really? maybe if all is relative?). Maybe only layer_abs_name_scope and related functions on the name scope ("name_scope") should be used late or even on-the-fly when building the net dict.

---

What cases do we have?

• Params (VariableLayer): must have well-defined name (attr chain) to the root module but this is straightforward (like _get_abs_canonical_name, maybe without the calls logic). Error if there is no attr chain.
• Module call (any module method call which is decorated with nn.scoped, where self is a module): _get_abs_canonical_name with reasonable fallbacks. Should never error.
• Other scoped function call (nn.scoped): It actually will create a dummy module (but this should not matter). Just same logic as module call. _get_abs_canonical_name with reasonable fallbacks. Should never error.
• make_layer: _get_abs_canonical_name with reasonable fallbacks. Should never error.
• Cond/Loop/MaskedComputation: _get_abs_canonical_name with reasonable fallbacks. Should never error.

So, how to reimplement?

• _get_abs_canonical_name for a tensor (param) is a bit different at the leaf end (not a module but a tensor).
• _get_abs_canonical_name should have a relaxed option (never fail, reasonable fallbacks).
• _get_abs_canonical_name and _get_suggested_name might be merged as a generic function: get_name_path_to_root or so. _get_suggested_name would use self.parent.module as the root module.

The relaxed option should first really try to find a direct attr chain and only if that is not found then use any of the fallbacks.

Another extension for the relaxed option: We can look at the stack frames and search for any module method (e.g. also __init__).

Should we reuse subnets of modules? output must be unique but output is not really relevant when we always access the sublayers directly. How would we do this? Check the module calls (or create new entry there).

When do we need the name? Basically only at the very end.

Maybe we can still leave the NameCtx.name attr as-is (for children etc, also relevant for the order) and have a special property/method get_layer_name which is only used when constructing the net dict.