returnn 1.20241020.5643__tar.gz → 1.20241022.224754__tar.gz

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/LICENSE

@@ -176,4 +176,11 @@ Contains code from PyTorch:
  - Copyright 2016..2023 various developers
  - Various code snippets
 
+Contains code from lion-pytorch:
+ - https://github.com/google/automl/blob/master/lion/lion_pytorch.py
+ - https://github.com/lucidrains/lion-pytorch/
+ - MIT License / Apache License 2.0
+ - Copyright 2023 Google Research, Phil Wang
+ - torch/optim/lion.py
+
 Various code snippets from StackOverflow, which are under Creative Commons / Public Domain.

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/PKG-INFO

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

returnn-1.20241022.224754/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241022.224754'
+long_version = '1.20241022.224754+git.b5db365'

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn/frontend/audio/mel.py

@@ -56,8 +56,8 @@ def mel_filterbank(
         filter_bank_matrix_np = _mel_filter_bank_matrix_np(
             f_min=f_min, f_max=f_max, sampling_rate=sampling_rate, fft_size=fft_length, nr_of_filters=out_dim.dimension
         )
-        filter_bank_matrix_np = filter_bank_matrix_np.astype(x.dtype)
         filter_bank_matrix = rf.convert_to_tensor(filter_bank_matrix_np, dims=(in_dim, out_dim), _backend=backend)
+        filter_bank_matrix = rf.cast(filter_bank_matrix, dtype=x.dtype)
         filter_bank_matrix = rf.copy_to_device(filter_bank_matrix, x.device)
         if backend.executing_eagerly():
             if len(_mel_filter_bank_matrix_cache) > 100:
@@ -191,6 +191,9 @@ def log_mel_filterbank_from_raw(
         fft_length=n_fft,
     )
     power_spectrogram = rf.abs(spectrogram) ** 2.0
+    # stft might have upcasted this to float32 because some PyTorch versions don't support stft on bfloat16.
+    # https://github.com/pytorch/pytorch/issues/117844
+    power_spectrogram = rf.cast(power_spectrogram, dtype=raw_audio.dtype)
     mel_fbank = mel_filterbank(power_spectrogram, in_dim=in_dim_, out_dim=out_dim, sampling_rate=sampling_rate)
     log_mel_fbank = rf.safe_log(mel_fbank, eps=1e-10)
     if log_base != math.e:

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn/frontend/decoder/transformer.py

@@ -101,15 +101,11 @@ class TransformerDecoder(rf.Module):
         if pos_enc is None:
             pass
         elif isinstance(pos_enc, dict):
-            pos_enc = rf.build_from_dict(
-                pos_enc, feat_dim=embed_dim or model_dim, dtype=self.input_embedding.weight.dtype
-            )
+            pos_enc = rf.build_from_dict(pos_enc, feat_dim=embed_dim or model_dim)
         elif isinstance(pos_enc, rf.Module):
             pass
         elif isinstance(pos_enc, FunctionType):
-            pos_enc = functools.partial(
-                pos_enc, feat_dim=embed_dim or model_dim, dtype=self.input_embedding.weight.dtype
-            )
+            pos_enc = functools.partial(pos_enc, feat_dim=embed_dim or model_dim)
         else:
             raise TypeError(f"unexpected pos_enc type {pos_enc!r}")
         self.pos_enc = pos_enc

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn/frontend/dtype.py

@@ -3,9 +3,17 @@ DType helpers
 """
 
 from __future__ import annotations
+from contextlib import contextmanager
 
 
-__all__ = ["get_default_float_dtype", "get_default_int_dtype", "get_default_array_index_dtype", "is_float_dtype"]
+__all__ = [
+    "get_default_float_dtype",
+    "set_default_float_dtype",
+    "set_default_float_dtype_ctx",
+    "get_default_int_dtype",
+    "get_default_array_index_dtype",
+    "is_float_dtype",
+]
 
 
 _default_float_dtype: str = "float32"
@@ -21,6 +29,32 @@ def get_default_float_dtype() -> str:
     return _default_float_dtype
 
 
+def set_default_float_dtype(dtype: str):
+    """
+    Set the default float dtype
+
+    :param dtype: the new default float dtype
+    """
+    global _default_float_dtype
+    assert isinstance(dtype, str)
+    _default_float_dtype = dtype
+
+
+@contextmanager
+def set_default_float_dtype_ctx(dtype: str):
+    """
+    :param dtype: see :func:`get_default_float_dtype`
+    """
+    global _default_float_dtype
+    assert isinstance(dtype, str)
+    old_default_float_dtype = _default_float_dtype
+    try:
+        _default_float_dtype = dtype
+        yield
+    finally:
+        _default_float_dtype = old_default_float_dtype
+
+
 def get_default_int_dtype() -> str:
     """
     https://data-apis.org/array-api/latest/API_specification/data_types.html#default-data-types

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn/tensor/utils.py

@@ -104,6 +104,9 @@ def tensor_fill_random_numpy_(
             x.raw_tensor = rnd.randint(0, 2, size=shape, dtype=x.dtype)
         elif x.dtype.startswith("float"):
             x.raw_tensor = rnd.normal(0.0, 1.0, size=shape).astype(x.dtype)
+        elif x.dtype == "bfloat16":
+            # Numpy does not support bfloat16, will later be casted to bfloat16
+            x.raw_tensor = rnd.normal(0.0, 1.0, size=shape).astype("float32")
         elif x.dtype.startswith("complex"):
             real = rnd.normal(0.0, 1.0, size=shape)
             imag = rnd.normal(0.0, 1.0, size=shape)

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn/torch/data/extern_data.py

@@ -3,7 +3,7 @@ From raw dict to extern_data tensor dict.
 """
 
 from __future__ import annotations
-from typing import Any, Union, Dict, List, Sequence
+from typing import Optional, Any, Union, Dict, List, Sequence
 import numpy
 import torch
 from returnn.tensor import Tensor, TensorDict, Dim
@@ -27,13 +27,18 @@ def raw_dict_to_extern_data(
     *,
     extern_data_template: TensorDict,
     device: Union[str, torch.device],
+    float_dtype: Optional[Union[str, torch.dtype]] = None,
 ) -> TensorDict:
     """
     :param extern_data_raw: This comes out of the DataLoader, via our collate_batch.
     :param extern_data_template: Specified via `extern_data` in the config.
     :param device: E.g. the GPU.
+    :param float_dtype:
     :return: tensor dict, like extern_data_template, but with raw tensors set to Torch tensors, on the right device.
     """
+    if isinstance(float_dtype, str):
+        float_dtype = getattr(torch, float_dtype)
+        assert isinstance(float_dtype, torch.dtype)
     assert isinstance(extern_data_raw, dict) and extern_data_raw
     batch_dim = get_batch_dim_from_extern_data(extern_data_template)
     for dim in _get_dyn_dims_from_extern_data(extern_data_template):
@@ -51,6 +56,8 @@ def raw_dict_to_extern_data(
                     dim.dimension == raw_tensor.shape[i]
                 ), f"shape mismatch for {k}: {raw_tensor.shape} vs {data.batch_shape}"
         if isinstance(raw_tensor, torch.Tensor):
+            if raw_tensor.dtype.is_floating_point and float_dtype:
+                raw_tensor = raw_tensor.to(dtype=float_dtype)
             data.dtype = str(raw_tensor.dtype).split(".")[-1]  # just overwrite for now...
             data.raw_tensor = raw_tensor.to(device)
         elif isinstance(raw_tensor, numpy.ndarray):

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn/torch/engine.py

@@ -4,7 +4,7 @@ Main engine for PyTorch
 
 from __future__ import annotations
 from typing import Optional, Any, Union, Callable, Dict, Set
-from contextlib import nullcontext
+from contextlib import nullcontext, ExitStack, contextmanager
 
 import gc
 import os
@@ -129,6 +129,13 @@ class Engine(EngineBase):
         self._forward_auto_split_batch_on_oom = config.bool("forward_auto_split_batch_on_oom", False)
         self._stop_on_nonfinite_train_score = config.bool("stop_on_nonfinite_train_score", True)
 
+        default_float_dtype = config.value("default_float_dtype", None)
+        if default_float_dtype is not None:
+            assert isinstance(default_float_dtype, str)
+            default_float_dtype = getattr(torch, default_float_dtype)
+            assert isinstance(default_float_dtype, torch.dtype)
+        self._default_float_dtype: Optional[torch.dtype] = default_float_dtype
+
         amp_options = self.config.opt_typed_value("torch_amp")
         grad_scaler_opts = self.config.typed_value("grad_scaler", NotSpecified)
         if amp_options is not None:
@@ -380,7 +387,10 @@ class Engine(EngineBase):
                     cur_count_grad_accum = 0
 
                 extern_data = extern_data_util.raw_dict_to_extern_data(
-                    extern_data_raw, extern_data_template=self.extern_data, device=self._device
+                    extern_data_raw,
+                    extern_data_template=self.extern_data,
+                    device=self._device,
+                    float_dtype=self._default_float_dtype,
                 )
                 self._run_step(extern_data, train_flag=True, train_func=True)
 
@@ -389,7 +399,7 @@ class Engine(EngineBase):
                 losses_dict = NumbersDict(
                     {
                         name: (
-                            float(loss.get_summed_loss().raw_tensor.detach().cpu().numpy())
+                            float(loss.get_summed_loss().raw_tensor.detach().cpu().item())
                             if self._device != "meta"
                             else float("nan")
                         )
@@ -553,7 +563,10 @@ class Engine(EngineBase):
                         torch.distributed.broadcast(_has_data, src=0)
 
                     extern_data = extern_data_util.raw_dict_to_extern_data(
-                        extern_data_raw, extern_data_template=self.extern_data, device=self._device
+                        extern_data_raw,
+                        extern_data_template=self.extern_data,
+                        device=self._device,
+                        float_dtype=self._default_float_dtype,
                     )
 
                     self._run_step(extern_data, train_func=True)
@@ -566,7 +579,7 @@ class Engine(EngineBase):
                     losses_dict = NumbersDict(
                         {
                             name: (
-                                float(loss.get_summed_loss().raw_tensor.detach().cpu().numpy())
+                                float(loss.get_summed_loss().raw_tensor.detach().cpu().item())
                                 if self._device != "meta"
                                 else float("nan")
                             )
@@ -686,6 +699,26 @@ class Engine(EngineBase):
 
         return data_loader
 
+    @contextmanager
+    def _run_ctx_mgr(self):
+        with ExitStack() as stack:
+            if self._use_autocast:
+                stack.enter_context(autocast(device_type=self._device.split(":")[0], dtype=self._autocast_dtype))
+            stack.enter_context(rf.set_default_device_ctx(self._device))
+            if self._default_float_dtype:
+                stack.enter_context(rf.set_default_float_dtype_ctx(str(self._default_float_dtype).split(".")[-1]))
+                stack.enter_context(self._set_torch_default_dtype_ctx_mgr(self._default_float_dtype))
+            yield
+
+    @contextmanager
+    def _set_torch_default_dtype_ctx_mgr(self, dtype: torch.dtype):
+        old_dtype = torch.get_default_dtype()
+        try:
+            torch.set_default_dtype(dtype)
+            yield
+        finally:
+            torch.set_default_dtype(old_dtype)
+
     def _run_step(
         self, extern_data: TensorDict, *, train_flag: bool = False, train_func: bool, _inside_wrapped: bool = False
     ):
@@ -706,11 +739,7 @@ class Engine(EngineBase):
                 expected_outputs=self._forward_step_expected_outputs, step=self.global_train_step, epoch=self.epoch
             )
 
-        with (
-            autocast(device_type=self._device.split(":")[0], dtype=self._autocast_dtype)
-            if self._use_autocast
-            else nullcontext()
-        ), rf.set_default_device_ctx(self._device):
+        with self._run_ctx_mgr():
             sentinel_kw = util.get_fwd_compat_kwargs()
             if train_func:
                 self._train_step_func(model=self._orig_model, extern_data=extern_data, **sentinel_kw)
@@ -893,6 +922,8 @@ class Engine(EngineBase):
                 )
             )
 
+        if self._default_float_dtype:
+            self._pt_model.to(dtype=self._default_float_dtype)
         self._pt_model.to(self._device)
 
         if model_epoch_filename and is_training:
@@ -906,11 +937,7 @@ class Engine(EngineBase):
 
         load_model_post_hooks = self.config.typed_value("load_model_post_hooks")
         if load_model_post_hooks:
-            with (
-                autocast(device_type=self._device.split(":")[0], dtype=self._autocast_dtype)
-                if self._use_autocast
-                else nullcontext()
-            ), rf.set_default_device_ctx(self._device):
+            with self._run_ctx_mgr():
                 sentinel_kw = util.get_fwd_compat_kwargs()
                 for hook in load_model_post_hooks:
                     hook(model=self._orig_model, **sentinel_kw)
@@ -1090,7 +1117,10 @@ class Engine(EngineBase):
             # Currently, this callback interface is intended to also be used by other backends,
             # and then the user can always assume Numpy arrays.
             if isinstance(raw, torch.Tensor):  # might already be numpy array
-                raw = raw.detach().cpu().numpy()
+                raw = raw.detach().cpu()
+                if raw.dtype == torch.bfloat16:
+                    raw = raw.float()
+                raw = raw.numpy()
             y.raw_tensor = raw
             return y
 
@@ -1120,7 +1150,10 @@ class Engine(EngineBase):
                     # Also resets any dyn dims, which might have been set in the prev step.
                     self._forward_step_expected_outputs.reset_content()
                 extern_data = extern_data_util.raw_dict_to_extern_data(
-                    extern_data_raw, extern_data_template=self.extern_data, device=self._device
+                    extern_data_raw,
+                    extern_data_template=self.extern_data,
+                    device=self._device,
+                    float_dtype=self._default_float_dtype,
                 )
                 try:
                     self._run_step(extern_data, train_func=False)
@@ -1224,7 +1257,10 @@ def _to_raw(n: Union[int, float, Tensor]):
     if isinstance(n, (int, float)):
         return n
     if isinstance(n, Tensor):
-        return n.raw_tensor.detach().cpu().numpy()
+        x = n.raw_tensor.detach().cpu()
+        if x.dtype == torch.bfloat16:
+            x = x.float()
+        return x.numpy()
     raise TypeError(f"Unexpected {n} of type {type(n)}")
 
 

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn/torch/frontend/_backend.py

@@ -676,6 +676,9 @@ class TorchBackend(Backend[torch.Tensor]):
         if len(batch_dims) != 1:
             targets_raw = torch.reshape(targets_raw, (batch_n_elems, targets_raw.shape[-1]))  # [B', S]
             targets_lengths = torch.reshape(targets_lengths, (batch_n_elems,))  # [B']
+        if log_probs.dtype == torch.bfloat16:
+            # Currently (PyTorch 2.5), ctc_loss does not support bfloat16.
+            log_probs = log_probs.to(torch.float32)
         loss_raw = torch.nn.functional.ctc_loss(
             log_probs=log_probs,
             targets=targets_raw,
@@ -691,7 +694,7 @@ class TorchBackend(Backend[torch.Tensor]):
             name="ctc_loss",
             dims=batch_dims,
             raw_tensor=loss_raw,
-            dtype=logits.dtype,
+            dtype=TorchBackend.get_dtype_name_raw(loss_raw),
         )
         return loss
 
@@ -2039,6 +2042,12 @@ class TorchBackend(Backend[torch.Tensor]):
                 pad_right = fft_length - frame_length - pad_left
                 window_pt = torch.nn.functional.pad(window_pt, (pad_left, pad_right))
 
+        orig_dtype = x_raw.dtype
+        if orig_dtype == torch.bfloat16:
+            # PyTorch stft does not support bfloat16 currently (PyTorch 2.5):
+            # https://github.com/pytorch/pytorch/issues/117844
+            # (Check back later here whether that's still the case...)
+            x_raw = x_raw.to(torch.float32)
         y_raw = torch.stft(
             x_raw,
             n_fft=fft_length,

returnn-1.20241022.224754/returnn/torch/optim/README.md

@@ -0,0 +1,5 @@
+Here we can put some arbitrary external optimizers.
+It might be copied from some existing code, or our own implementation.
+It might also happen that some of these will be added to later versions of PyTorch.
+So, regarding the user config, the optimizers here should be differentiated
+by having the full module name, e.g. like ``returnn.torch.optim.lion.Lion``.

returnn-1.20241022.224754/returnn/torch/optim/__init__.py

@@ -0,0 +1,3 @@
+"""
+Any custom optimizer
+"""

returnn-1.20241022.224754/returnn/torch/optim/lion.py

@@ -0,0 +1,205 @@
+"""
+Lion optimizer <https://arxiv.org/abs/2302.06675>
+
+Code adapted from https://github.com/lucidrains/lion-pytorch/,
+which is adapted from https://github.com/google/automl/blob/master/lion/lion_pytorch.py.
+"""
+
+from __future__ import annotations
+from typing import Optional, Tuple, Callable
+import inspect
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class Lion(Optimizer):
+    """
+    Lion (Evolved Sign Momentum (Evo_l_ved S_i_gn M_o_me_n_tum)) optimizer <https://arxiv.org/abs/2302.06675>
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-4,
+        betas: Tuple[float, float] = (0.9, 0.99),
+        weight_decay: float = 0.0,
+        use_triton: Optional[bool] = None,
+        decoupled_weight_decay: bool = False,
+    ):
+        assert lr > 0.0
+        assert all([0.0 <= beta <= 1.0 for beta in betas])
+
+        self._init_lr = lr
+        self.decoupled_wd = decoupled_weight_decay
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+
+        super().__init__(params, defaults)
+
+        if use_triton is None:
+            use_triton = bool(triton_update_fn)
+        self.use_triton = use_triton
+
+    @torch.no_grad()
+    def step(self, closure: Optional[Callable] = None):
+        """update step"""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                beta1, beta2 = group["betas"]
+                grad, lr, wd, state, decoupled_wd, init_lr = (
+                    p.grad,
+                    group["lr"],
+                    group["weight_decay"],
+                    self.state[p],
+                    self.decoupled_wd,
+                    self._init_lr,
+                )
+
+                # maybe decoupled weight decay
+
+                if decoupled_wd:
+                    wd /= init_lr
+
+                # init state - exponential moving average of gradient values
+
+                if len(state) == 0:
+                    state["exp_avg"] = torch.zeros_like(p)
+
+                exp_avg = state["exp_avg"]
+
+                if self.use_triton and p.is_cuda:
+                    triton_update_fn(p, grad, exp_avg, lr, wd, beta1, beta2)
+                else:
+                    update_fn(p, grad, exp_avg, lr, wd, beta1, beta2)
+
+        return loss
+
+
+# update functions
+
+
+def update_fn(p, grad, exp_avg, lr, wd, beta1, beta2):
+    """
+    Lion update function
+    """
+    # stepweight decay
+
+    p.data.mul_(1.0 - lr * wd)
+
+    # weight update
+
+    update = exp_avg.clone().mul_(beta1).add(grad, alpha=1.0 - beta1).sign_()
+    p.add_(update, alpha=-lr)
+
+    # decay the momentum running average coefficient
+
+    exp_avg.mul_(beta2).add_(grad, alpha=1.0 - beta2)
+
+
+try:
+    # noinspection PyPackageRequirements
+    import triton
+
+    # noinspection PyPackageRequirements
+    import triton.language as tl
+except ImportError as e:
+    triton = None
+    tl = None
+
+
+# restore_value is not available in older versions of triton
+if triton and "restore_value" in inspect.signature(triton.autotune).parameters:
+    # triton cuda kernel
+
+    # noinspection PyPep8Naming,PyArgumentList
+    @triton.autotune(
+        configs=[
+            triton.Config({"BLOCK_SIZE": 128}, num_warps=4),
+            triton.Config({"BLOCK_SIZE": 1024}, num_warps=8),
+        ],
+        key=["n_elements"],
+        restore_value=["p_ptr", "exp_avg_ptr"],
+    )
+    @triton.jit
+    def _triton_update_fn_kernel(
+        p_ptr,
+        grad_ptr,
+        exp_avg_ptr,
+        lr,
+        wd,
+        beta1,
+        beta2,
+        n_elements,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        pid = tl.program_id(axis=0)
+
+        block_start = pid * BLOCK_SIZE
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+
+        mask = offsets < n_elements
+
+        # offsetted pointers
+
+        offset_p_ptr = p_ptr + offsets
+        offset_grad_ptr = grad_ptr + offsets
+        offset_exp_avg_ptr = exp_avg_ptr + offsets
+
+        # load
+
+        p = tl.load(offset_p_ptr, mask=mask)
+        grad = tl.load(offset_grad_ptr, mask=mask)
+        exp_avg = tl.load(offset_exp_avg_ptr, mask=mask)
+
+        # stepweight decay
+
+        p = p * (1 - lr * wd)
+
+        # diff between momentum running average and grad
+
+        diff = exp_avg - grad
+
+        # weight update
+
+        update = diff * beta1 + grad
+
+        # torch.sign
+
+        can_update = update != 0
+        update_sign = tl.where(update > 0, -lr, lr)
+
+        p = p + update_sign * can_update
+
+        # decay the momentum running average coefficient
+
+        exp_avg = diff * beta2 + grad
+
+        # store new params and momentum running average coefficient
+
+        tl.store(offset_p_ptr, p, mask=mask)
+        tl.store(offset_exp_avg_ptr, exp_avg, mask=mask)
+
+    def triton_update_fn(
+        p: torch.Tensor, grad: torch.Tensor, exp_avg: torch.Tensor, lr: float, wd: float, beta1: float, beta2: float
+    ):
+        """
+        Lion update function using triton kernel
+        """
+        assert all([t.is_cuda for t in (p, grad, exp_avg)])
+        n_elements = p.numel()
+
+        def _grid(meta):
+            return tuple((triton.cdiv(n_elements, meta["BLOCK_SIZE"]),))
+
+        _triton_update_fn_kernel[_grid](p, grad, exp_avg, lr, wd, beta1, beta2, n_elements)
+
+else:
+    triton_update_fn = None

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn/torch/updater.py

@@ -5,7 +5,7 @@ and model param update logic in general.
 
 from __future__ import annotations
 
-from typing import Optional, Union, Any, Type, Sequence, Iterable, Set, Dict, List, Tuple
+from typing import Optional, Union, Any, Type, Callable, Sequence, Iterable, Set, Dict, List, Tuple
 import os
 import gc
 import torch
@@ -38,28 +38,40 @@ def _init_optimizer_classes_dict():
         _OptimizerClassesDict[name.lower()] = cls
 
 
-def get_optimizer_class(class_name) -> Type[torch.optim.Optimizer]:
+def get_optimizer_class(
+    class_name: Union[str, Type[torch.optim.Optimizer], Callable[[], Type[torch.optim.Optimizer]]]
+) -> Type[torch.optim.Optimizer]:
     """
-    :param str|()->torch.optim.Optimizer|type[torch.optim.Optimizer] class_name:
-        Optimizer data, e.g. "adam", torch.optim.Adam...
-    :return: Optimizer class
+    :param class_name: Optimizer class, either as str (e.g. "adam"), as type (torch.optim.Adam) or callable.
+        If str, we support all torch.optim optimizers (ignoring case) (e.g. "adam"),
+        or class names with full module path (e.g. "returnn.torch.optim.lion.Lion").
+    :return: Optimizer class, e.g. torch.optim.Adam
     """
     _init_optimizer_classes_dict()
     if isinstance(class_name, type):
         assert issubclass(class_name, torch.optim.Optimizer)
+        return class_name
     elif callable(class_name):
-        class_name = class_name()
-    else:
-        assert isinstance(class_name, str)
-        assert (
-            class_name.lower() in _OptimizerClassesDict
-        ), "%s not found in the available torch optimizers list: %s." % (
-            class_name.lower(),
-            ", ".join("'%s'" % key for key in _OptimizerClassesDict),
-        )
-        class_name = _OptimizerClassesDict[class_name.lower()]
+        return class_name()
+    elif isinstance(class_name, str):
+        if "." in class_name:
+            import importlib
 
-    return class_name
+            mod_name, class_name_ = class_name.rsplit(".", 1)
+            mod = importlib.import_module(mod_name)
+            return getattr(mod, class_name_)
+
+        if class_name.lower() not in _OptimizerClassesDict:
+            raise ValueError(
+                "Optimizer %r not found in the available torch optimizers list: %s."
+                % (
+                    class_name.lower(),
+                    ", ".join("'%s'" % key for key in _OptimizerClassesDict),
+                )
+            )
+        return _OptimizerClassesDict[class_name.lower()]
+    else:
+        raise TypeError(f"Invalid optimizer class_name {class_name!r} type {type(class_name).__name__}")
 
 
 def _get_class_init_kwargs(optim_class):
@@ -411,7 +423,7 @@ class Updater:
         # If the user specified it as epsilon, parse it as eps for the optimizer
         if "eps" in optim_class_init_kwargs and "epsilon" in opt_kwargs:
             opt_kwargs["eps"] = opt_kwargs.pop("epsilon")
-        if "learning_rate" in opt_kwargs:
+        if "learning_rate" in opt_kwargs or "lr" in opt_kwargs:
             raise ValueError("'learning_rate' should be set outside of the 'optimizer' dict.")
         lr = lr * opt_kwargs.pop("learning_rate_multiplier", 1.0)
         opt_kwargs["lr"] = lr

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn/torch/util/exception_helper.py

@@ -86,7 +86,10 @@ def _help_data_or_array(
     :return: (info,(min,max))
     """
     if isinstance(value, torch.Tensor):
-        value = value.detach().cpu().numpy()
+        value = value.detach().cpu()
+        if value.dtype == torch.bfloat16:
+            value = value.float()
+        value = value.numpy()
     v_minmax = -1, -1
     if isinstance(value, np.ndarray):
         info = "shape %s, dtype %s" % (value.shape, value.dtype)

{returnn-1.20241020.5643 → returnn-1.20241022.224754}/returnn.egg-info/PKG-INFO

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