returnn 1.20240907.111018__tar.gz → 1.20240908.14935__tar.gz

{returnn-1.20240907.111018 → returnn-1.20240908.14935}/PKG-INFO

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

returnn-1.20240908.14935/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240908.014935'
+long_version = '1.20240908.014935+git.7390bea'

{returnn-1.20240907.111018 → returnn-1.20240908.14935}/returnn/frontend/array_.py

@@ -10,6 +10,7 @@ from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
 from ._backend import Backend, global_backend, get_backend_by_raw_tensor_type
 from .types import RawTensorTypes
+from . import _utils
 
 T = TypeVar("T")
 
@@ -829,7 +830,9 @@ def where(
 
     :return: true_ if cond else false_, elemwise.
     """
-    cond = rf.convert_to_tensor(cond)
+    if not isinstance(cond, Tensor):
+        backend = _utils.get_backend_from_tensors(true_, false_)
+        cond = rf.convert_to_tensor(cond, _backend=backend)
     # noinspection PyProtectedMember
     return cond._raw_backend.where(cond, true_, false_, allow_broadcast_all_sources=allow_broadcast_all_sources)
 

{returnn-1.20240907.111018 → returnn-1.20240908.14935}/returnn/tensor/_dim_extra.py

@@ -805,7 +805,7 @@ class _DimMixin:
         :return: dyn_size_ext on the device
         """
         assert self.dyn_size_ext
-        if not device:
+        if not device or device == "cpu":
             return self.dyn_size_ext
 
         import returnn.frontend as rf
@@ -1900,18 +1900,22 @@ class _DimMixin:
                 res.append(tag)
         return res
 
-    def get_size_tensor(self) -> _t.Tensor:
+    def get_size_tensor(self, *, device: Optional[str] = None) -> _t.Tensor:
         """
         :return: size tensor, or dyn_size_ext if defined
         :rtype: _t.Tensor
         """
         if self.dyn_size_ext:
-            return self.dyn_size_ext
+            if not device or device == "cpu":
+                return self.dyn_size_ext
+            return self.get_dyn_size_ext_for_device(device)
 
         import returnn.frontend as rf
 
+        if device is None:
+            device = "cpu"
         assert self.size is not None
-        return rf.convert_to_tensor(self.size, name="%s:size" % self.description, device="cpu")
+        return rf.convert_to_tensor(self.size, name="%s:size" % self.description, device=device)
 
     def get_dim_value(self) -> Union[int, _t.RawTensorType]:
         """

{returnn-1.20240907.111018 → returnn-1.20240908.14935}/returnn/torch/engine.py

@@ -794,7 +794,7 @@ class Engine(EngineBase):
                         print("(No relevant parameters matching.)", file=log.v3)
                     continue
                 print(f"Pre-load weights for key '{preload_key}' from {opts['filename']}", file=log.v3)
-                preload_model_state = torch.load(opts["filename"])
+                preload_model_state = torch.load(opts["filename"], map_location=self._device)
                 if opts.get("checkpoint_key", "model") is not None:
                     # This can be used if an external checkpoint saves a checkpoint a different structure that just the
                     # model state dict. E.g., if a checkpoint is created using

{returnn-1.20240907.111018 → returnn-1.20240908.14935}/returnn/torch/frontend/_backend.py

@@ -868,11 +868,18 @@ class TorchBackend(Backend[torch.Tensor]):
             name = name or "raw_tensor"
         else:
             name = name or "const"
-            value = torch.tensor(
-                value,
-                dtype=TorchBackend.as_dtype_raw(dtype),
-                device=device or rf.get_default_device(),
-            )
+            if isinstance(value, (bool, int, float, bool, complex, numpy.number)):
+                # torch.full avoids a device sync.
+                # https://github.com/pytorch/pytorch/issues/120996#issuecomment-2319976284
+                value = torch.full(
+                    (), value, dtype=TorchBackend.as_dtype_raw(dtype), device=device or rf.get_default_device()
+                )
+            else:
+                value = torch.tensor(
+                    value,
+                    dtype=TorchBackend.as_dtype_raw(dtype),
+                    device=device or rf.get_default_device(),
+                )
         assert isinstance(value, torch.Tensor)
         return Tensor(name, dims=dims, dtype=dtype, sparse_dim=sparse_dim, raw_tensor=value)
 
@@ -914,21 +921,40 @@ class TorchBackend(Backend[torch.Tensor]):
         - Tensor.__getitem__
         - torch.embedding
         """
+        axis_int = source.get_axis_from_description(axis, allow_int=False)
         if isinstance(indices, Tensor):
             indices: Tensor[torch.Tensor]
         elif isinstance(indices, int):
+            if indices == 0:  # small opt for this case, as we can simplify a bit, and clip_to_valid is not needed
+                out = Tensor(
+                    "gather",
+                    dims=source.dims[:axis_int] + source.dims[axis_int + 1 :],
+                    dtype=source.dtype,
+                    sparse_dim=source.sparse_dim,
+                )
+                if source.feature_dim and source.feature_dim in out.dims:
+                    out.feature_dim = source.feature_dim
+                if axis_int == 0:
+                    out.raw_tensor = source.raw_tensor[indices]
+                else:
+                    out.raw_tensor = source.raw_tensor[(slice(None),) * axis_int + (indices,)]
+                return out
             indices = Tensor(
                 "indices_int",
                 dims=(),
                 dtype=rf.get_default_array_index_dtype(),
-                raw_tensor=torch.tensor(indices, dtype=TorchBackend.as_dtype_raw(rf.get_default_array_index_dtype())),
+                raw_tensor=torch.full(
+                    (),
+                    indices,
+                    dtype=TorchBackend.as_dtype_raw(rf.get_default_array_index_dtype()),
+                    device=source.raw_tensor.device,
+                ),
             )
         else:
             raise TypeError(f"Unsupported type for indices: {type(indices)}")
-        axis_int = source.get_axis_from_description(axis, allow_int=False)
         if clip_to_valid:
             if axis.dyn_size_ext:
-                indices = rf.clip_by_value(indices, 0, axis.dyn_size_ext - 1)
+                indices = rf.clip_by_value(indices, 0, axis.get_dyn_size_ext_for_device(indices.device) - 1)
             else:
                 indices = indices.copy()
                 indices.raw_tensor = torch.clamp(indices.raw_tensor, 0, source.raw_tensor.shape[axis_int] - 1)

{returnn-1.20240907.111018 → returnn-1.20240908.14935}/returnn.egg-info/PKG-INFO

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

returnn-1.20240907.111018/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240907.111018'
-long_version = '1.20240907.111018+git.0ee25d9'