returnn 1.20230718.124003__tar.gz → 1.20230718.183712__tar.gz

{returnn-1.20230718.124003/returnn.egg-info → returnn-1.20230718.183712}/PKG-INFO

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

returnn-1.20230718.183712/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230718.183712'
+long_version = '1.20230718.183712+git.b77467a'

{returnn-1.20230718.124003 → returnn-1.20230718.183712}/returnn/frontend/_backend.py

@@ -653,6 +653,7 @@ class Backend(Generic[T]):
         dims: Sequence[Dim],
         dtype: str,
         sparse_dim: Optional[Dim] = None,
+        device: Optional[str] = None,
         name: Optional[str] = None,
     ) -> Tensor[T]:
         """
@@ -660,6 +661,7 @@ class Backend(Generic[T]):
         :param dims:
         :param dtype:
         :param sparse_dim:
+        :param device:
         :param name:
         :return: tensor
         """

{returnn-1.20230718.124003 → returnn-1.20230718.183712}/returnn/frontend/_utils.py

@@ -79,13 +79,16 @@ def bin_op_out_template(
     :return: out, a, b
     """
     src_dtype = None
+    src_device = None
     if isinstance(a, Tensor):
         src_dtype = a.dtype
+        src_device = a.device
     elif isinstance(b, Tensor):
         src_dtype = b.dtype
-    a = rf.convert_to_tensor(a, dtype=src_dtype, _backend=backend)
+        src_device = b.device
+    a = rf.convert_to_tensor(a, dtype=src_dtype, device=src_device, keep_scalar_on_cpu=allow_scalar, _backend=backend)
     src_dtype = src_dtype or a.dtype
-    b = rf.convert_to_tensor(b, dtype=src_dtype, _backend=backend)
+    b = rf.convert_to_tensor(b, dtype=src_dtype, device=src_device, keep_scalar_on_cpu=allow_scalar, _backend=backend)
     # sanity checks
     # noinspection PyProtectedMember
     assert a._raw_backend == b._raw_backend, "Cannot combine tensors from two different frontends, e.g. TF and PT"

{returnn-1.20230718.124003 → returnn-1.20230718.183712}/returnn/frontend/array_.py

@@ -46,6 +46,8 @@ def convert_to_tensor(
     dtype: Optional[str] = None,
     sparse_dim: Optional[Dim] = None,
     shape: Sequence[Dim] = None,
+    device: Optional[str] = None,
+    keep_scalar_on_cpu: bool = False,
     name: Optional[str] = None,
     _backend: Optional[Type[Backend]] = None,
 ) -> Tensor[T]:
@@ -56,6 +58,8 @@ def convert_to_tensor(
     :param sparse_dim:
     :param shape: alias for dims, for some older code
     :param name:
+    :param device:
+    :param keep_scalar_on_cpu: if the value is already on the CPU, keep it there, even if `device` is sth else
     :param _backend:
     :return: tensor
     """
@@ -81,6 +85,8 @@ def convert_to_tensor(
                 dtype = value.dtype.name
             else:
                 raise ValueError(f"number {value} type {type(value)} needs explicit `dtype` specification")
+        if keep_scalar_on_cpu:
+            device = "cpu"
     elif isinstance(value, numpy.ndarray):
         if _backend is None:
             # Small exception: Do not use the NumpyBackend but the global backend in this case.
@@ -100,7 +106,9 @@ def convert_to_tensor(
             ]
         if dtype is None:
             dtype = value_backend.get_dtype_name_raw(value)
-    return _backend.convert_to_tensor(value=value, dims=dims, dtype=dtype, sparse_dim=sparse_dim, name=name)
+    return _backend.convert_to_tensor(
+        value=value, dims=dims, dtype=dtype, sparse_dim=sparse_dim, device=device, name=name
+    )
 
 
 constant = convert_to_tensor  # alias for some older code

returnn-1.20230718.183712/returnn/frontend/device.py

@@ -0,0 +1,56 @@
+"""
+Device handling.
+"""
+
+
+from __future__ import annotations
+from typing import Optional
+from contextlib import contextmanager
+from returnn.tensor import Tensor
+
+
+__all__ = ["copy_to_device", "get_default_device", "set_default_device_ctx"]
+
+
+_default_device: Optional[str] = None
+
+
+def copy_to_device(x: Tensor, device: Optional[str] = None) -> Tensor:
+    """
+    Copy tensor to device.
+
+    :param x: tensor
+    :param device:
+    :return: tensor on device
+    """
+    if not device:
+        device = get_default_device()
+    if not device:
+        return x
+    if x.raw_tensor is None:
+        return x
+    if x.device == device:
+        return x
+    # noinspection PyProtectedMember
+    return x._raw_backend.copy_to_device(x, device)
+
+
+def get_default_device() -> Optional[str]:
+    """
+    :return: default device, where to put new tensors (via random number generators, constant, range_over_dim, etc)
+    """
+    return _default_device
+
+
+@contextmanager
+def set_default_device_ctx(device: Optional[str]):
+    """
+    :param device: see :func:`get_default_device`
+    """
+    global _default_device
+    old_device = _default_device
+    try:
+        _default_device = device
+        yield
+    finally:
+        _default_device = old_device

{returnn-1.20230718.124003 → returnn-1.20230718.183712}/returnn/tensor/_dim_extra.py

@@ -114,6 +114,8 @@ class _DimExtra:
         # We can have different tag variants per batch info (e.g. with beam), or per control flow ctx.
         # They each have same_as = self. The same_base should have the base (global) batch info.
         self.same_for_batch_ctx = {}  # type: Dict[Tuple[BatchInfo,Optional[ControlFlowContext]],_d.Dim]
+        self.cache_dyn_size_ext_dev = {}  # type: Dict[str,_t.Tensor]  # device -> dyn_size_ext
+        self.cache_seq_mask = {}  # type: Dict[str,_t.Tensor]  # device -> seq_mask
 
     def __getstate__(self):
         d = vars(self).copy()
@@ -349,6 +351,9 @@ class _DimMixin:
         """
         if self.dyn_size_ext:
             self.dyn_size_ext.raw_tensor = None
+        if self._extra:
+            self._extra.cache_dyn_size_ext_dev.clear()
+            self._extra.cache_seq_mask.clear()
 
     def _can_use_in_ctx(self, ctx):
         """
@@ -701,10 +706,23 @@ class _DimMixin:
             )
         self.dyn_size_ext.placeholder = dyn_size
 
-    def get_mask(self: Dim, *, dim_order: Optional[Sequence[Dim]] = None) -> _t.Tensor:
+    def _get_dyn_size_ext_for_device(self: Dim, device: Optional[str]) -> _t.Tensor:
+        if not device:
+            return self.dyn_size_ext
+
+        import returnn.frontend as rf
+
+        self._make_extra()
+        if device in self._extra.cache_dyn_size_ext_dev:
+            return self._extra.cache_dyn_size_ext_dev[device]
+        self._extra.cache_dyn_size_ext_dev[device] = rf.copy_to_device(self.dyn_size_ext, device=device)
+        return self._extra.cache_dyn_size_ext_dev[device]
+
+    def get_mask(self: Dim, *, dim_order: Optional[Sequence[Dim]] = None, device: Optional[str] = None) -> _t.Tensor:
         """
         :param dim_order: if given, the dims of the mask will be in this order.
             This can be useful if the mask is broadcasted against some other tensor.
+        :param str|None device: if given, will move the mask to this device
         :return: if need_masking(), the corresponding mask.
             If this is e.g. the time-dim T of shape [B], then the mask will be of shape [B,T].
             The mask could be used with :func:`masked_select` (``boolean_mask``) or ``where``.
@@ -715,9 +733,17 @@ class _DimMixin:
         # noinspection PyProtectedMember
         backend = self.dyn_size_ext._raw_backend
 
+        if not device:
+            device = rf.get_default_device()
+
+        if self._extra and device in self._extra.cache_seq_mask:
+            return self._extra.cache_seq_mask[device]
+
+        size_ext = self._get_dyn_size_ext_for_device(device)
+
         max_idx = rf.reduce(
-            self.dyn_size_ext,
-            axis=self.dyn_size_ext.dims,
+            size_ext,
+            axis=size_ext.dims,
             mode="max",
             # Masking here is not always possible, e.g. if we have
             # tag = Dim{'self-att-keys'['time:var:extern_data:classes'[B]]}
@@ -731,9 +757,11 @@ class _DimMixin:
         # and this likely produces nan in backprop or elsewhere.
         # Thus, mask size_ext itself, and set the padded values to 1.
         # This assumes that max_idx >= 1.
-        size_ext = self.dyn_size_ext.copy_masked(max_idx)
-        idx_range = backend.range_over_dim(self)
+        size_ext = size_ext.copy_masked(max_idx)
+        with rf.set_default_device_ctx(device):
+            idx_range = backend.range_over_dim(self)
         seq_mask = rf.compare(idx_range, "<", size_ext, allow_broadcast_all_sources=True, dim_order=dim_order)
+        self._make_extra().cache_seq_mask[device] = seq_mask
         return seq_mask
 
     def is_batch_dim(self):
@@ -1053,7 +1081,8 @@ class _DimMixin:
                 assert x.dimension is not None
                 if y is None:
                     if not template_only and backend and not tf:
-                        y = backend.convert_to_tensor(x.dimension, dims=[], dtype=size_dtype, name=y_name)
+                        with rf.set_default_device_ctx(None):
+                            y = backend.convert_to_tensor(x.dimension, dims=[], dtype=size_dtype, name=y_name)
                     else:
                         y = _t.Tensor(
                             name=y_name,

{returnn-1.20230718.124003 → returnn-1.20230718.183712}/returnn/tensor/_tensor_extra.py

@@ -2730,7 +2730,6 @@ class _TensorMixin(_TensorMixinBase):
           if BT or TB major, and axis is T or None.
           In general compatible to placeholder, i.e. same ndim, with broadcast dims.
           We assert here that the axis is dynamic (:func:`is_axis_dynamic`), i.e. we have the size.
-        :rtype: tf.Tensor
         """
         if isinstance(axis, Dim):
             axis = self.get_axis_from_description(axis)
@@ -2781,7 +2780,7 @@ class _TensorMixin(_TensorMixinBase):
         assert 0 <= axis < self.batch_ndim
         assert axis != self.batch_dim_axis
         tag: Dim = self.dim_tags[axis]
-        return tag.get_mask(dim_order=self.dims)
+        return tag.get_mask(dim_order=self.dims, device=self.device)
 
     def get_sequence_lengths_broadcast(self, axis=None):
         """

{returnn-1.20230718.124003 → returnn-1.20230718.183712}/returnn/tensor/tensor_dict.py

@@ -68,6 +68,7 @@ class TensorDict:
             dim.batch = None
             if dim.dyn_size_ext:
                 dim.dyn_size_ext.reset()
+            dim.reset_eager()
 
     def copy_template(self) -> TensorDict:
         """copy template"""

{returnn-1.20230718.124003 → returnn-1.20230718.183712}/returnn/tf/frontend_layers/_backend.py

@@ -431,6 +431,7 @@ class ReturnnLayersBackend(Backend[Layer]):
         dims: Sequence[Dim],
         dtype: str,
         sparse_dim: Optional[Dim] = None,
+        device: Optional[str] = None,
         name: Optional[str] = None,
     ) -> Tensor[Layer]:
         """convert to tensor"""

{returnn-1.20230718.124003 → returnn-1.20230718.183712}/returnn/tf/frontend_low_level/_backend.py

@@ -370,6 +370,7 @@ class TFBackend(Backend[tf.Tensor]):
         dims: Sequence[Dim],
         dtype: str,
         sparse_dim: Optional[Dim] = None,
+        device: Optional[str] = None,
         name: Optional[str] = None,
     ) -> _TT:
         """
@@ -377,6 +378,7 @@ class TFBackend(Backend[tf.Tensor]):
         :param dims:
         :param dtype:
         :param sparse_dim:
+        :param device:
         :param name:
         :return: tensor
         """

{returnn-1.20230718.124003 → returnn-1.20230718.183712}/returnn/torch/frontend/_backend.py

@@ -525,7 +525,11 @@ class TorchBackend(Backend[torch.Tensor]):
         :return: parameter
         """
         assert all(d.is_static() for d in tensor.dims)
-        data = torch.zeros([d.dimension for d in tensor.dims], dtype=TorchBackend.as_dtype_raw(tensor.dtype))
+        data = torch.zeros(
+            [d.dimension for d in tensor.dims],
+            dtype=TorchBackend.as_dtype_raw(tensor.dtype),
+            device=rf.get_default_device(),
+        )
         if tensor.dtype.startswith("int"):
             requires_grad = False
         else:
@@ -645,6 +649,7 @@ class TorchBackend(Backend[torch.Tensor]):
         dims: Sequence[Dim],
         dtype: str,
         sparse_dim: Optional[Dim] = None,
+        device: Optional[str] = None,
         name: Optional[str] = None,
     ) -> Tensor[torch.Tensor]:
         """
@@ -652,6 +657,7 @@ class TorchBackend(Backend[torch.Tensor]):
         :param dims:
         :param dtype:
         :param sparse_dim:
+        :param device:
         :param name:
         :return: tensor
         """
@@ -661,7 +667,11 @@ 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))
+            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)
 
@@ -682,7 +692,9 @@ class TorchBackend(Backend[torch.Tensor]):
             # onnx::ConstantOfShape (via torch.full) must get shape as int64.
             # https://github.com/rwth-i6/returnn/issues/1333#issuecomment-1607236783
             shape = [dim.long() if isinstance(dim, torch.Tensor) else dim for dim in shape]
-        raw_tensor = torch.full(shape, fill_value, dtype=TorchBackend.as_dtype_raw(dtype))
+        raw_tensor = torch.full(
+            shape, fill_value, dtype=TorchBackend.as_dtype_raw(dtype), device=rf.get_default_device()
+        )
         return Tensor(
             "full", dims=dims, sparse_dim=sparse_dim, feature_dim=feature_dim, dtype=dtype, raw_tensor=raw_tensor
         )
@@ -934,7 +946,9 @@ class TorchBackend(Backend[torch.Tensor]):
             sparse_dim=dim,
             dtype=dtype,
         )
-        out.raw_tensor = torch.arange(dim.get_dim_value(), dtype=TorchBackend.as_dtype_raw(out.dtype))
+        out.raw_tensor = torch.arange(
+            dim.get_dim_value(), dtype=TorchBackend.as_dtype_raw(out.dtype), device=rf.get_default_device()
+        )
         return out
 
     @staticmethod
@@ -1084,7 +1098,7 @@ class TorchBackend(Backend[torch.Tensor]):
             out = Tensor(
                 name=f"random_{distribution}", dims=dims, dtype=dtype, sparse_dim=sparse_dim, feature_dim=feature_dim
             )
-            out.raw_tensor = torch.empty(shape, dtype=dtype_)
+            out.raw_tensor = torch.empty(shape, dtype=dtype_, device=rf.get_default_device())
         assert explicit_state is None  # not implemented otherwise
         generator = None  # using the global default from PT
         assert isinstance(static, bool)

{returnn-1.20230718.124003 → returnn-1.20230718.183712/returnn.egg-info}/PKG-INFO

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

returnn-1.20230718.124003/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230718.124003'
-long_version = '1.20230718.124003+git.88cdf02'

returnn-1.20230718.124003/returnn/frontend/device.py

@@ -1,26 +0,0 @@
-"""
-Device handling.
-"""
-
-
-from __future__ import annotations
-from typing import Optional
-from returnn.tensor import Tensor
-
-
-def copy_to_device(x: Tensor, device: Optional[str]) -> Tensor:
-    """
-    Copy tensor to device.
-
-    :param x: tensor
-    :param device:
-    :return: tensor on device
-    """
-    if not device:
-        return x
-    if x.raw_tensor is None:
-        return x
-    if x.device == device:
-        return x
-    # noinspection PyProtectedMember
-    return x._raw_backend.copy_to_device(x, device)