returnn 1.20230403.154105__tar.gz → 1.20230403.212652__tar.gz

{returnn-1.20230403.154105/returnn.egg-info → returnn-1.20230403.212652}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230403.154105
+Version: 1.20230403.212652
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20230403.212652/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230403.212652'
+long_version = '1.20230403.212652+git.6c39555'

{returnn-1.20230403.154105 → returnn-1.20230403.212652}/returnn/frontend/run_ctx.py

@@ -9,11 +9,11 @@ or forwarding loop.
 from __future__ import annotations
 from typing import Optional, Union, Any, Sequence, Dict
 from dataclasses import dataclass
-from returnn.tensor import Tensor, Dim
+from returnn.tensor import Tensor, Dim, TensorDict
 import returnn.frontend as rf
 
 
-__all__ = ["RunCtx", "Loss", "Output", "get_run_ctx", "init_train_step_run_ctx", "init_forward_step_run_ctx"]
+__all__ = ["RunCtx", "Loss", "get_run_ctx", "init_train_step_run_ctx", "init_forward_step_run_ctx"]
 
 
 _run_ctx = None  # type: Optional[RunCtx]
@@ -75,7 +75,7 @@ class RunCtx:
         """
         self.stage = stage
         self.losses = {}  # type: Dict[str, Loss]
-        self.outputs = {}  # type: Dict[str, Output]
+        self.outputs = TensorDict()
 
     def mark_as_loss(
         self,
@@ -141,7 +141,7 @@ class RunCtx:
             custom_inv_norm_factor=custom_inv_norm_factor,
         )
 
-    def mark_as_output(self, tensor: Union[Tensor, Any], name: str, *, shape: Optional[Sequence[int]] = None) -> None:
+    def mark_as_output(self, tensor: Union[Tensor, Any], name: str, *, dims: Optional[Sequence[int]] = None) -> None:
         """
         Mark this as an output.
         This has the effect that RETURNN will in any case construct the corresponding layer.
@@ -153,7 +153,7 @@ class RunCtx:
 
         :param tensor:
         :param name:
-        :param shape: this specifies the order of the dims of the output, such that it is well-defined
+        :param dims: this specifies the order of the dims of the output, such that it is well-defined
             for some external application.
             If not specified, we try to infer BTF or BF as default, if that works, otherwise it will be an error.
         """
@@ -161,7 +161,33 @@ class RunCtx:
         if not isinstance(tensor, Tensor):
             tensor = rf.convert_to_tensor(tensor)
         assert name not in self.outputs
-        self.outputs[name] = Output(tensor=tensor, name=name, shape=shape)
+        if dims is None:
+            # We try some reasonable defaults, specifically: BTF or BF.
+            rem_dims = list(tensor.dims)
+            dims = []
+            if tensor.have_batch_axis():
+                rem_dims.remove(tensor.get_batch_dim_tag())
+                dims.append(tensor.get_batch_dim_tag())
+            if tensor.have_time_axis():
+                rem_dims.remove(tensor.get_time_dim_tag())
+                dims.append(tensor.get_time_dim_tag())
+            dyn_dims = [d for d in dims if d.is_dynamic()]
+            if len(dyn_dims) > 1:
+                raise Exception(
+                    f"Cannot infer order of dims automatically for output {name!r}. Please specify a shape explicitly."
+                )
+            elif len(dyn_dims) == 1:
+                rem_dims.remove(dyn_dims[0])
+                dims.append(dyn_dims[0])
+            if len(rem_dims) > 1:
+                raise Exception(
+                    f"Cannot infer order of dims automatically for output {name!r}. Please specify a shape explicitly."
+                )
+            elif len(rem_dims) == 1:
+                dims.append(rem_dims[0])
+        tensor = tensor.copy_transpose(dims, allow_int=False)
+        tensor = tensor.copy(name=name)
+        self.outputs.data[name] = tensor
 
     def mark_as_default_output(self, tensor: Union[Tensor, Any], *, shape: Optional[Sequence[Dim]] = None) -> None:
         """
@@ -173,7 +199,7 @@ class RunCtx:
         :param tensor:
         :param shape:
         """
-        self.mark_as_output(tensor, "output", shape=shape)
+        self.mark_as_output(tensor, "output", dims=shape)
 
     def total_loss(self) -> Union[Tensor, float]:
         """
@@ -233,16 +259,3 @@ class Loss:
         else:
             loss = self.get_summed_loss()
         return loss * self.scale
-
-
-@dataclass
-class Output:
-    """
-    Output via :func:`RunCtx.mark_as_output`.
-
-    We collect all relevant information here.
-    """
-
-    tensor: Tensor
-    name: str
-    shape: Optional[Sequence[Dim]] = None

{returnn-1.20230403.154105 → returnn-1.20230403.212652}/returnn/tensor/_tensor_extra.py

@@ -651,14 +651,16 @@ class _TensorMixin(_TensorMixinBase):
         assert self.time_dim_axis is not None
         return self.copy_move_axis(self.time_dim_axis, time_dim_axis)
 
-    def copy_transpose(self, perm) -> _t.Tensor:
+    def copy_transpose(self, perm: Sequence[Union[int, Dim]], *, allow_int: bool = True) -> _t.Tensor:
         """
-        :param list[int] perm: permutation of the axes, counted with batch-dim.
-          Maps the new axes to the old axes
+        :param perm: permutation of the axes. Maps the new axes to the old axes
+        :param allow_int: allow int as axis, otherwise only :class:`Dim`
         :return: copy of myself with permuted axes
         """
-        assert len(perm) == self.batch_ndim
-        assert set(perm) == set(range(self.batch_ndim))
+        assert len(perm) == self.batch_ndim, f"{self}: invalid perm {perm!r} length"
+        perm_ = perm
+        perm = [self.get_axis_from_description(a, allow_int=allow_int) for a in perm]
+        assert set(perm) == set(range(self.batch_ndim)), f"{self}: invalid perm {perm_!r} (axes: {perm!r})"
         if all(perm[axis] == axis for axis in range(self.batch_ndim)):
             return self.copy()
 
@@ -3156,7 +3158,7 @@ class _TensorMixin(_TensorMixinBase):
         """
         import returnn.frontend as rf
 
-        rf.get_run_ctx().mark_as_output(self, name=name, shape=shape)
+        rf.get_run_ctx().mark_as_output(self, name=name, dims=shape)
 
     def mark_as_default_output(self: Tensor, *, shape: Optional[Sequence[Dim]] = None) -> None:
         """

{returnn-1.20230403.154105 → returnn-1.20230403.212652}/returnn/tf/frontend_layers/config_entry_points.py

@@ -123,12 +123,12 @@ def get_net_dict(
         # Note that this logic might change.
         root_scope.marked_losses.append(loss_t)
 
-    for out in rf.get_run_ctx().outputs.values():
+    for out in rf.get_run_ctx().outputs.data.values():
         if out.name == "output" and out.name not in root_scope.children:
             layer = root_scope.get_child(out.name)
         else:
             layer = root_scope.get_new_child(suggested_name=out.name)
-        out_t = _utils.copy(out.tensor, name=layer)
+        out_t = _utils.copy(out, name=layer)
         if layer.name != "output":
             out_t.raw_tensor.layer_dict["is_output_layer"] = True
         root_scope.marked_outputs.append(out_t)

{returnn-1.20230403.154105 → returnn-1.20230403.212652}/returnn/torch/data/pipeline.py

@@ -18,17 +18,19 @@ However, having this separate pure PyTorch implementation is useful to allow to
 other PyTorch datasets more directly, including also HuggingFace datasets.
 """
 
+from __future__ import annotations
+from typing import List, Dict
 import sys
 from copy import deepcopy
 
-import numpy as np
+import numpy
 import torch
 import torch.utils.data
 
 from returnn.util.basic import NumbersDict
 
 
-def create_tensor(array: np.ndarray) -> torch.Tensor:
+def create_tensor(array: numpy.ndarray) -> torch.Tensor:
     """
     Adjust non-supported dtypes
 
@@ -36,14 +38,14 @@ def create_tensor(array: np.ndarray) -> torch.Tensor:
     """
     # The only supported PyTorch dtypes are:
     # float64, float32, float16, complex64, complex128, int64, int32, int16, int8, uint8, and bool.
-    if array.dtype == np.uint32:
-        array = np.asarray(array, dtype=np.int64)
+    if array.dtype == numpy.uint32:
+        array = numpy.asarray(array, dtype=numpy.int64)
     return torch.tensor(array)
 
 
-def collate_batch(batch):
+def collate_batch(batch: List[Dict[str, numpy.ndarray]]) -> Dict[str, torch.Tensor]:
     """
-    :param list[dict[str, numpy.ndarray]] batch:
+    :param batch:
     """
     assert isinstance(batch, list)
     assert batch, "batch is empty?"

{returnn-1.20230403.154105 → returnn-1.20230403.212652}/returnn/torch/frontend/_rand.py

@@ -10,7 +10,8 @@ import warnings
 
 def no_grad_trunc_normal_(tensor: torch.Tensor, mean, std, a, b, *, generator=None):
     """
-    Code copied and adopted from torch.nn.init._no_grad_trunc_normal_.
+    Code copied and adopted from torch.nn.init._no_grad_trunc_normal_,
+    to support the extra `generator` argument (https://github.com/pytorch/pytorch/issues/98200).
 
     Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
 

{returnn-1.20230403.154105 → returnn-1.20230403.212652/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20230403.154105
+Version: 1.20230403.212652
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20230403.154105/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230403.154105'
-long_version = '1.20230403.154105+git.1732a0b'