returnn 1.20240910.114102__tar.gz → 1.20240917.180949__tar.gz

{returnn-1.20240910.114102 → returnn-1.20240917.180949}/PKG-INFO

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

returnn-1.20240917.180949/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240917.180949'
+long_version = '1.20240917.180949+git.6964c19'

{returnn-1.20240910.114102 → returnn-1.20240917.180949}/returnn/frontend/run_ctx.py

@@ -9,6 +9,7 @@ or forwarding loop.
 from __future__ import annotations
 from typing import Optional, Union, Any, Sequence, Dict
 from dataclasses import dataclass
+from contextlib import contextmanager
 from returnn.tensor import Tensor, Dim, TensorDict, batch_dim
 import returnn.frontend as rf
 from . import _backend
@@ -126,6 +127,25 @@ class RunCtx:
         """
         return self._train_flag
 
+    @contextmanager
+    def train_flag_ctx(self, train_flag: Union[bool, Tensor]):
+        """
+        Context manager to temporarily set the train_flag.
+
+        Usage example e.g. to disable dropout for some code::
+
+            with rf.get_run_ctx().train_flag_ctx(False):
+                ...
+
+        :param train_flag: whether we are in training mode
+        """
+        old_train_flag = self.train_flag
+        self._train_flag = train_flag
+        try:
+            yield
+        finally:
+            self._train_flag = old_train_flag
+
     @property
     def step(self) -> Union[int, Tensor]:
         """

{returnn-1.20240910.114102 → returnn-1.20240917.180949}/returnn/tensor/_dim_extra.py

@@ -136,6 +136,9 @@ class _DimExtra:
         if self.kind is not None:
             self.kind = {v.name: v for v in DimTypes.Types}[self.kind]
 
+    def __sis_state__(self):
+        raise ValueError(f"{self}: currently not expected to be part of the Sisyphus state/hash")
+
 
 class _DimMixin:
     name: Optional[str]
@@ -2280,7 +2283,7 @@ class _DimMixin:
 
     def _cache_dim_math_get(
         self: Dim, op_kind: str, operand: Union[Dim, int]
-    ) -> Tuple[Dict[Tuple[str, Union[Dim, int]], Dim], Tuple[str, Union[Dim, int]], Optional[Dim]]:
+    ) -> Tuple[_CacheDimMath, Tuple[str, Union[Dim, int]], Optional[Dim]]:
         same_base = self.get_same_base()
         # noinspection PyProtectedMember
         extra = same_base._make_extra()
@@ -2290,6 +2293,9 @@ class _DimMixin:
         cache_key = (op_kind, operand)
         return cache, cache_key, cache.get(cache_key, None)
 
+    def __sis_state__(self):
+        raise ValueError(f"{self}: currently not expected to be part of the Sisyphus state/hash")
+
 
 def _make_constant_static_dim(value, kind=None):
     """

{returnn-1.20240910.114102 → returnn-1.20240917.180949}/returnn/torch/frontend/_backend.py

@@ -980,7 +980,12 @@ class TorchBackend(Backend[torch.Tensor]):
             index_ext_dims = list(source.dims)
             index_ext_dims[axis_int] = index_own_dims_flat
             assert indices.dims == tuple(index_ext_dims)
-            out_raw = torch.gather(source.raw_tensor, dim=axis_int, index=indices.raw_tensor.type(torch.int64))
+            out_raw = torch.gather(
+                source.raw_tensor,
+                dim=axis_int,
+                # torch.gather wants indices as int64, and the indices should be the same device as the source.
+                index=indices.raw_tensor.type(torch.int64).to(source.raw_tensor.device),
+            )
             if len(index_own_dims) == 1:
                 pass  # nothing to do
             elif len(index_own_dims) == 0:

{returnn-1.20240910.114102 → returnn-1.20240917.180949}/returnn.egg-info/PKG-INFO

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

returnn-1.20240910.114102/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240910.114102'
-long_version = '1.20240910.114102+git.ec1fa03'