returnn 1.20241030.185827__tar.gz → 1.20241106.124322__tar.gz

{returnn-1.20241030.185827 → returnn-1.20241106.124322}/PKG-INFO

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

returnn-1.20241106.124322/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241106.124322'
+long_version = '1.20241106.124322+git.bf0c605'

{returnn-1.20241030.185827 → returnn-1.20241106.124322}/returnn/frontend/array_.py

@@ -36,6 +36,7 @@ __all__ = [
     "masked_scatter",
     "sequence_mask",
     "pack_padded",
+    "pad_packed",
     "gather",
     "scatter",
     "scatter_argmax",
@@ -627,6 +628,8 @@ def pack_padded(
     Packing means to only store the non-padded frames.
     This uses :func:`masked_select` internally based on the mask of non-masked frames.
 
+    See :func:`pad_packed` for the inverse operation.
+
     :param source:
     :param dims: dims in source to pack. the order defines the format. first dim is major, etc.
         if there are no padded frames, e.g. dims=[B,T] would just result in the [B*T,...] reshaped tensor.
@@ -648,6 +651,14 @@ def pack_padded(
     return rf.masked_select(source, mask=mask, dims=dims, out_dim=out_dim)
 
 
+def pad_packed(source: Tensor, *, in_dim: Dim, dims: Sequence[Dim]) -> Tensor:
+    """
+    Inverse of :func:`pack_padded`, i.e. unpack the sequence, i.e. pad it back to the original length.
+    """
+    mask = rf.sequence_mask(dims, device=source.device)
+    return rf.masked_scatter(source, mask=mask, in_dim=in_dim, dims=dims)
+
+
 # noinspection PyUnusedLocal
 def gather(
     source: Tensor,

{returnn-1.20241030.185827 → returnn-1.20241106.124322}/returnn/learning_rate_control.py

@@ -5,7 +5,7 @@ The base class is :class:`LearningRateControl`.
 
 from __future__ import annotations
 
-from typing import Optional, Any, Dict
+from typing import Optional, Union, Any, Dict
 import typing
 import os
 import returnn.util.basic as util
@@ -350,7 +350,7 @@ class LearningRateControl:
                 relative_error /= learning_rate / self.default_learning_rate
         return relative_error
 
-    def set_epoch_error(self, epoch, error):
+    def set_epoch_error(self, epoch: int, error: Dict[str, Union[float, Dict[str, float]]]):
         """
         :type epoch: int
         :type error: dict[str,float|dict[str,float]]

{returnn-1.20241030.185827 → returnn-1.20241106.124322}/returnn/torch/data/pipeline.py

@@ -28,7 +28,7 @@ import numpy
 import torch
 import torch.utils.data
 
-from returnn.util.basic import NumbersDict
+from returnn.util.basic import NumbersDict, get_fwd_compat_kwargs
 
 
 def create_tensor(array: numpy.ndarray) -> Union[torch.Tensor, numpy.ndarray]:
@@ -59,7 +59,7 @@ def collate_batch(batch: List[Dict[str, numpy.ndarray]]) -> Dict[str, Union[torc
 
     res = {}
     for key in data_keys:
-        if key == "num_seqs":
+        if key in ("num_seqs", "epoch"):
             res[key] = batch[0][key]  # it should always be the same
             continue
         ls = [create_tensor(sample[key]) for sample in batch]
@@ -119,7 +119,7 @@ class ChunkingIterDataPipe(torch.utils.data.IterDataPipe):
 
             if not chunking_data_keys:
                 chunking_data_keys = list(data_dict.keys())  # use all if not configured separately
-                chunking_data_key_black_list = ["seq_tag", "seq_idx", "num_seqs"]
+                chunking_data_key_black_list = ["seq_tag", "seq_idx", "num_seqs", "epoch"]
                 for key in chunking_data_key_black_list:
                     if key in chunking_data_keys:
                         chunking_data_keys.remove(key)
@@ -208,20 +208,66 @@ class BatchingIterDataPipe(torch.utils.data.IterDataPipe):
     def __init__(self, dataset: torch.utils.data.IterableDataset, batch_size=1, max_seqs=None):
         """
         :param dataset: dataset to apply batching to
-        :param int|dict[str,int]|None batch_size: Maximum number of time steps (e.g. audio frames / words) in one
-            batch (padding included).
+        :param int|dict[str,int]|None|function batch_size: Maximum number of time steps (e.g. audio frames / words)
+            in one batch (padding included).
             If given as a dict data_key -> value, sets different individual limits per data key.
             If None, no limit.
-        :param int|None max_seqs: maximum number of sequences in a batch,
-            None means unlimited (also -1 to match TF backend)
+            Can also be a callable with kwargs epoch, seq_idx, epoch_continuous, **_other_kwargs,
+            returning the batch size.
+        :param int|None|function max_seqs: maximum number of sequences in a batch,
+            None means unlimited (also -1 to match TF backend).
+            Can also be a callable with kwargs epoch, seq_idx, epoch_continuous, **_other_kwargs,
+            returning the max seqs.
         """
         super().__init__()
         self._dataset = dataset
-        self._max_batch_size = NumbersDict(sys.maxsize if batch_size is None else batch_size)
-        self._max_seqs = sys.maxsize if (max_seqs is None or max_seqs == -1) else max_seqs
+        self._max_batch_size = self._parse_batch_size(batch_size)
+        self._max_seqs = self._parse_max_seqs(max_seqs)
 
-        assert self._max_batch_size.min_value() > 0
-        assert self._max_seqs > 0
+        if not callable(self._max_batch_size):
+            assert isinstance(self._max_batch_size, NumbersDict) and self._max_batch_size.min_value() > 0
+        if not callable(self._max_seqs):
+            assert isinstance(self._max_seqs, int) and self._max_seqs > 0
+
+    @staticmethod
+    def _parse_batch_size(
+        batch_size: Union[int, Dict[str, int], NumbersDict, None, Callable],
+        *,
+        data_dict: Optional[Dict[str, Any]] = None,
+    ) -> Union[NumbersDict, Callable]:
+        """
+        :param batch_size: see __init__()
+        :return: batch_size
+        """
+        if callable(batch_size):
+            if data_dict:
+                batch_size = batch_size(**BatchingIterDataPipe._get_user_func_kwargs_from_data_dict(data_dict))
+            else:
+                return batch_size
+        return NumbersDict(sys.maxsize if batch_size is None else batch_size)
+
+    @staticmethod
+    def _parse_max_seqs(
+        max_seqs: Union[int, None, Callable], *, data_dict: Optional[Dict[str, Any]] = None
+    ) -> Union[int, Callable]:
+        """
+        :param max_seqs: see __init__()
+        :return: max_seqs
+        """
+        if callable(max_seqs):
+            if data_dict:
+                max_seqs = max_seqs(**BatchingIterDataPipe._get_user_func_kwargs_from_data_dict(data_dict))
+            else:
+                return max_seqs
+        return sys.maxsize if (max_seqs is None or max_seqs == -1) else max_seqs
+
+    @staticmethod
+    def _get_user_func_kwargs_from_data_dict(data_dict: Dict[str, Any]) -> Dict[str, Any]:
+        epoch = int(data_dict["epoch"])
+        seq_idx = int(data_dict["seq_idx"])
+        num_seqs = int(data_dict["num_seqs"])  # >=1 if known, otherwise -1
+        epoch_continuous = (epoch - 1 + (seq_idx + 1) / num_seqs) if num_seqs > 0 else None
+        return {"epoch": epoch, "seq_idx": seq_idx, "epoch_continuous": epoch_continuous, **get_fwd_compat_kwargs()}
 
     def __iter__(self):
         """
@@ -233,7 +279,12 @@ class BatchingIterDataPipe(torch.utils.data.IterDataPipe):
         current_max_sequence_lengths = NumbersDict(0)  # data_key -> length of longest sequence in current batch
 
         for data_dict in self._dataset:
-            if len(current_batch) == self._max_seqs:
+            max_seqs = self._parse_max_seqs(self._max_seqs, data_dict=data_dict)
+            max_batch_size = self._parse_batch_size(self._max_batch_size, data_dict=data_dict)
+            assert isinstance(max_seqs, int) and max_seqs > 0
+            assert isinstance(max_batch_size, NumbersDict) and max_batch_size.min_value() > 0
+
+            if len(current_batch) >= max_seqs:
                 yield current_batch
                 current_batch = []
                 current_max_sequence_lengths = NumbersDict(0)
@@ -246,7 +297,7 @@ class BatchingIterDataPipe(torch.utils.data.IterDataPipe):
             max_sequence_lengths_if_included = NumbersDict.max([current_max_sequence_lengths, sequence_lengths])
             batch_size_if_included = max_sequence_lengths_if_included * (len(current_batch) + 1)  # including padding
 
-            if current_batch and batch_size_if_included.any_compare(self._max_batch_size, (lambda a, b: a > b)):
+            if current_batch and batch_size_if_included.any_compare(max_batch_size, (lambda a, b: a > b)):
                 yield current_batch
                 current_batch = [data_dict]
                 current_max_sequence_lengths = sequence_lengths

{returnn-1.20241030.185827 → returnn-1.20241106.124322}/returnn/torch/data/returnn_dataset_wrapper.py

@@ -67,7 +67,14 @@ class ReturnnDatasetIterDataPipe(torch.utils.data.IterDataPipe):
 
     def reset(self):
         """
-        :return:
+        This is called by PyTorch DataLoader mechanism once we create a new iterator over the DataLoader.
+        This happens at the beginning of each epoch.
+
+        (Note: The mechanism where ``reset()`` is actually called is very obfuscated in PyTorch.
+        As I understand it, there is a IterDataPipe metaclass (_IterDataPipeMeta)
+        which automatically registers a hook on ``__iter__`` via ``hook_iterator``.
+        Deep inside the complex logic of this hook, it calls ``_set_datapipe_valid_iterator_id``
+        which then calls ``reset()``.)
         """
         self._reset_callback()
 
@@ -81,6 +88,8 @@ class ReturnnDatasetIterDataPipe(torch.utils.data.IterDataPipe):
         except Exception:  # might not work for all datasets
             num_seqs = -1
         num_seqs = numpy.array(num_seqs)
+        assert self._dataset.epoch is not None
+        epoch = numpy.array(self._dataset.epoch)
 
         try:
             data_keys = self._dataset.get_data_keys()
@@ -94,6 +103,8 @@ class ReturnnDatasetIterDataPipe(torch.utils.data.IterDataPipe):
                 # It's slightly redundant to have num_seqs in each entry,
                 # but it's difficult to pass this back to the main proc otherwise.
                 data["num_seqs"] = num_seqs
+                # epoch is also redundant, but that's the cleanest/simplest way to pass it on to BatchingIterDataPipe.
+                data["epoch"] = epoch
                 yield data
                 seq_index += 1
 

{returnn-1.20241030.185827 → returnn-1.20241106.124322}/returnn/torch/engine.py

@@ -34,6 +34,7 @@ from returnn.util import NumbersDict
 from returnn.util.basic import hms, NotSpecified
 from returnn.util.result_with_reason import ResultWithReason
 from returnn.util.debug import debug_shell
+from returnn.util.math import simplify_and_format_number
 from returnn.forward_iface import ForwardCallbackIface
 
 from .updater import Updater
@@ -125,6 +126,7 @@ class Engine(EngineBase):
         self._log_memory_usage = config.bool("torch_log_memory_usage", False)
         self._log_batch_size = config.bool("log_batch_size", False) and log.verbose[5]
         self._calculate_exp_loss = config.bool("calculate_exp_loss", False)
+        self._log_grad_norm = _parse_log_grad_norm(config)
         self._reset_dev_memory_caches = config.bool("reset_dev_memory_caches", False)
         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)
@@ -383,6 +385,7 @@ class Engine(EngineBase):
                 del num_seqs_
                 if num_seqs is not None:
                     assert last_seq_idx < num_seqs
+                epoch_continuous = (self.epoch - 1 + (last_seq_idx + 1) / num_seqs) if num_seqs is not None else None
 
                 # clear the gradients when every gradient accumulation loop starts
                 if zero_grad_next_step:
@@ -415,7 +418,10 @@ class Engine(EngineBase):
 
                 if accum_grad_multiple_step_dyn:
                     accum_grad_multiple_step = accum_grad_multiple_step_dyn(
-                        epoch=self.epoch, global_train_step=self.global_train_step
+                        epoch=self.epoch,
+                        epoch_continuous=epoch_continuous,
+                        global_train_step=self.global_train_step,
+                        **util.get_fwd_compat_kwargs(),
                     )
                 cur_count_grad_accum += 1
                 perform_update_step = cur_count_grad_accum >= accum_grad_multiple_step
@@ -429,6 +435,12 @@ class Engine(EngineBase):
                     else:
                         total_loss.raw_tensor.backward()
 
+                if self._log_grad_norm and perform_update_step:
+                    key = f"grad_norm:p{simplify_and_format_number(self._log_grad_norm)}"
+                    assert key not in losses_dict
+                    inv_norm_factors_dict[key] = 1.0  # once per update step
+                    losses_dict[key] = _get_total_grad_norm(self._pt_model, p=self._log_grad_norm)
+
                 # only update the weights when every gradient accumulation loop ends
                 if perform_update_step:
                     self._updater.step(grad_scaler=self._grad_scaler)
@@ -469,9 +481,7 @@ class Engine(EngineBase):
                 step_idx += 1
                 self.global_train_step += 1
                 self._updater.set_current_train_step(
-                    global_train_step=self.global_train_step,
-                    epoch=self.epoch,
-                    epoch_continuous=(self.epoch - 1 + (last_seq_idx + 1) / num_seqs) if num_seqs is not None else None,
+                    global_train_step=self.global_train_step, epoch=self.epoch, epoch_continuous=epoch_continuous
                 )
         except Exception as exc:
             help_on_torch_exception(exc, step_idx=step_idx, model=self._orig_model, extern_data=extern_data)
@@ -480,8 +490,8 @@ class Engine(EngineBase):
         elapsed = time.monotonic() - epoch_start_time
         elapsed_computation_percentage = elapsed_computation_time / elapsed
         print(
-            "Trained %i steps, %s elapsed (%.1f%% computing time)"
-            % (step_idx, hms(elapsed), (elapsed_computation_percentage * 100.0)),
+            "Epoch %i: Trained %i steps, %s elapsed (%.1f%% computing time)"
+            % (self.epoch, step_idx, hms(elapsed), (elapsed_computation_percentage * 100.0)),
             file=log.v3,
         )
 
@@ -501,7 +511,7 @@ class Engine(EngineBase):
         if self._do_save():
             self.learning_rate_control.save()
 
-        print(f"Total train loss:", _format_score(dict(accumulated_losses_dict)), file=log.v3)
+        print(f"Epoch {self.epoch}: Total train loss:", _format_score(dict(accumulated_losses_dict)), file=log.v3)
 
         self._maybe_report_dev_memory_stats()
 
@@ -532,8 +542,6 @@ class Engine(EngineBase):
         self._reset_dev_memory_stats()
 
         eval_dump_str = []
-        score_keys = None
-        error_keys = None
 
         for dataset_name, dataset in self.eval_datasets.items():
             if skip_already_evaluated and self._is_dataset_evaluated(name=dataset_name):
@@ -575,10 +583,6 @@ class Engine(EngineBase):
                     self._run_step(extern_data, train_func=True)
                     train_ctx = rf.get_run_ctx()
 
-                    if score_keys is None:
-                        score_keys = set(name for name, loss in train_ctx.losses.items() if not loss.as_error)
-                        error_keys = set(name for name, loss in train_ctx.losses.items() if loss.as_error)
-
                     losses_dict = NumbersDict(
                         {
                             name: (
@@ -615,14 +619,7 @@ class Engine(EngineBase):
                 self.learning_rate_control.save()
 
             # Same format as the TF engine.
-            eval_dump_str += [
-                "%s: score %s error %s"
-                % (
-                    dataset_name,
-                    _format_score({name: accumulated_losses_dict[name] for name in score_keys}),
-                    _format_score({name: accumulated_losses_dict[name] for name in error_keys}),
-                )
-            ]
+            eval_dump_str += ["%s: %s" % (dataset_name, _format_score(dict(accumulated_losses_dict)))]
 
             if self._torch_distributed_ctx:
                 assert self._torch_distributed_ctx.rank() == 0
@@ -630,7 +627,11 @@ class Engine(EngineBase):
                 torch.distributed.broadcast(_has_data, src=0)
 
         if not self._torch_distributed_ctx or self._torch_distributed_ctx.rank() == 0:
-            print(" ".join(eval_dump_str) if eval_dump_str else "(No evaluations.)", file=log.v1)
+            print(
+                f"Epoch {self.epoch} evaluation:",
+                " ".join(eval_dump_str) if eval_dump_str else "(No evaluations.)",
+                file=log.v1,
+            )
 
         self._maybe_report_dev_memory_stats()
 
@@ -684,7 +685,7 @@ class Engine(EngineBase):
         batch_size = self.config.typed_value("batch_size", -1)
         batch_size = self.config.typed_value(f"batch_size_{'train' if train else 'dev'}", batch_size)
         assert batch_size != -1, f"batch_size or batch_size_{'train' if train else 'dev'} not defined in config"
-        max_seqs = self.config.int("max_seqs", -1)
+        max_seqs = self.config.typed_value("max_seqs", -1)
         batches_dataset = data_pipeline.BatchingIterDataPipe(wrapped_dataset, batch_size=batch_size, max_seqs=max_seqs)
 
         loader_opts = self.config.typed_value("torch_dataloader_opts") or {}
@@ -1286,9 +1287,9 @@ def _print_process(
     if log.verbose[5]:  # report every minibatch
         info = [report_prefix, "step %i" % step]
         if eval_info:  # Such as score.
-            info += ["%s %s" % (k, _format_value(v)) for k, v in eval_info.items()]
+            info += ["%s %s" % (k, _format_score_value(v)) for k, v in eval_info.items()]
         if batch_size_info:
-            info += ["%s %s" % (k, _format_value(v)) for k, v in batch_size_info.items()]
+            info += ["%s %s" % (k, _format_score_value(v)) for k, v in batch_size_info.items()]
         if log_memory_usage_device:
             dev = torch.device(log_memory_usage_device)
             if dev.type == "cuda":
@@ -1324,11 +1325,11 @@ def _format_score(score: Dict[str, float]) -> str:
     if not score:
         return "None"
     if len(score) == 1:
-        return _format_value(list(score.values())[0])
-    return " ".join(["%s %s" % (key.split(":", 2)[-1], _format_value(score[key])) for key in score.keys()])
+        return _format_score_value(list(score.values())[0])
+    return " ".join(["%s %s" % (k, _format_score_value(v)) for k, v in score.items()])
 
 
-def _format_value(v: Any) -> str:
+def _format_score_value(v: Any) -> str:
     if isinstance(v, float):
         if abs(v) > 1.0e3 or abs(v) < 1.0e-3:
             return f"{v:.3e}"
@@ -1422,3 +1423,39 @@ def _set_torch_default_dtype_ctx_mgr(dtype: torch.dtype):
         yield
     finally:
         torch.set_default_dtype(old_dtype)
+
+
+def _parse_log_grad_norm(config: Config) -> Optional[Union[int, float]]:
+    log_grad_norm = config.opt_typed_value("log_grad_norm", False)
+    if isinstance(log_grad_norm, str):
+        if log_grad_norm.lower() in ["true", "false", "none"]:
+            log_grad_norm = {"true": True, "false": False, "none": None}[log_grad_norm.lower()]
+        else:
+            raise ValueError(f"Invalid value for log_grad_norm: {log_grad_norm!r}")
+    if log_grad_norm is None:
+        pass
+    elif isinstance(log_grad_norm, bool):
+        if log_grad_norm:
+            log_grad_norm = 2
+        else:
+            log_grad_norm = None
+    elif isinstance(log_grad_norm, (int, float)):
+        assert log_grad_norm > 0, f"log_grad_norm {log_grad_norm} > 0 expected"  # otherwise fine...
+    else:
+        raise TypeError(f"Invalid type for log_grad_norm: {log_grad_norm!r} type {type(log_grad_norm)}")
+    return log_grad_norm
+
+
+def _get_total_grad_norm(model: torch.nn.Module, p: float) -> float:
+    return float(
+        torch.norm(
+            torch.stack(
+                [
+                    param.grad.norm(p=p).detach().cpu()
+                    for param in model.parameters()
+                    if param.requires_grad and param.grad is not None
+                ]
+            ),
+            p=p,
+        ).item()
+    )

{returnn-1.20241030.185827 → returnn-1.20241106.124322}/returnn/util/basic.py

@@ -5,7 +5,7 @@ Various generic utilities, which are shared across different backend engines.
 """
 
 from __future__ import annotations
-from typing import Optional, Union, Any, Generic, TypeVar, Iterable, Tuple, Dict, List, Callable
+from typing import Optional, Union, Any, Generic, TypeVar, Iterable, Tuple, Dict, List, Set, Callable
 
 import subprocess
 from subprocess import CalledProcessError
@@ -554,12 +554,11 @@ def describe_torch_version() -> str:
     return "%s (%s in %s)" % (version, git_info, tdir)
 
 
-def get_tensorflow_version_tuple():
+def get_tensorflow_version_tuple() -> Tuple[int, ...]:
     """
     :return: tuple of ints, first entry is the major version
-    :rtype: tuple[int]
     """
-    import tensorflow as tf
+    import tensorflow as tf  # noqa
     import re
 
     return tuple([int(re.sub("(-rc[0-9]|-dev[0-9]*)", "", s)) for s in tf.__version__.split(".")])
@@ -1963,9 +1962,9 @@ class NumbersDict:
         self.value = broadcast_value
         self.max = self._max_error
 
-    def copy(self):
+    def copy(self) -> NumbersDict:
         """
-        :rtype: NumbersDict
+        :return: copy
         """
         return NumbersDict(self)
 
@@ -1982,11 +1981,10 @@ class NumbersDict:
             numbers_dict={k: const_number for k in numbers_dict.dict.keys()},
         )
 
-    def copy_like(self, numbers_dict):
+    def copy_like(self, numbers_dict: NumbersDict) -> NumbersDict:
         """
-        :param NumbersDict numbers_dict:
+        :param numbers_dict:
         :return: copy of self with same keys as numbers_dict as far as we have them
-        :rtype: NumbersDict
         """
         if self.value is not None:
             return NumbersDict(
@@ -1999,11 +1997,11 @@ class NumbersDict:
             )
 
     @property
-    def keys_set(self):
+    def keys_set(self) -> Set[str]:
         """
         Also see :func:`keys_union` if you want to have a deterministic order.
 
-        :rtype: set[str]
+        :return: set of keys
         """
         return set(self.dict.keys())
 
@@ -2020,29 +2018,32 @@ class NumbersDict:
                     res.append(key)
         return res
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str):
         if self.value is not None:
             return self.dict.get(key, self.value)
         return self.dict[key]
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: str, value):
         self.dict[key] = value
 
-    def __delitem__(self, key):
+    def __delitem__(self, key: str):
         del self.dict[key]
 
-    def get(self, key, default=None):
+    def __contains__(self, item: str):
+        return item in self.dict
+
+    def get(self, key: str, default=None):
         """
-        :param str key:
+        :param key:
         :param T default:
         :rtype: object|T
         """
         # Keep consistent with self.__getitem__. If self.value is set, this will always be the default value.
         return self.dict.get(key, self.value if self.value is not None else default)
 
-    def pop(self, key, *args):
+    def pop(self, key: str, *args):
         """
-        :param str key:
+        :param key:
         :param T args: default, or not
         :rtype: object|T
         """
@@ -2055,22 +2056,21 @@ class NumbersDict:
         # which would only make sense for our values, not the dict keys.
         raise Exception("%s.__iter__ is undefined" % self.__class__.__name__)
 
-    def keys(self):
+    def keys(self) -> Iterable[str]:
         """
         :rtype: set[str]
         """
         return self.dict.keys()
 
-    def values(self):
+    def values(self) -> List[Any]:
         """
-        :rtype: list[object]
+        :return: values: dict values + self.value
         """
         return list(self.dict.values()) + ([self.value] if self.value is not None else [])
 
-    def items(self):
+    def items(self) -> Iterable[Tuple[str, Any]]:
         """
         :return: dict items. this excludes self.value
-        :rtype: str[(str,object)]
         """
         return self.dict.items()
 
@@ -2080,9 +2080,9 @@ class NumbersDict:
         """
         return self.value is not None or key in self.dict
 
-    def has_values(self):
+    def has_values(self) -> bool:
         """
-        :rtype: bool
+        :return: any values in self.dict or self.value
         """
         return bool(self.dict) or self.value is not None
 
@@ -2186,12 +2186,12 @@ class NumbersDict:
     def __neg__(self):
         return self.unary_op(op=lambda a: -a)
 
-    def __bool__(self):
+    def __bool__(self) -> bool:
         return any(self.values())
 
     __nonzero__ = __bool__  # Python 2
 
-    def elem_eq(self, other, result_with_default=True):
+    def elem_eq(self, other, result_with_default: bool = True) -> NumbersDict:
         """
         Element-wise equality check with other.
         Note about broadcast default value: Consider some key which is neither in self nor in other.
@@ -2202,8 +2202,8 @@ class NumbersDict:
           You can control the behavior via result_with_default.
 
         :param NumbersDict|T other:
-        :param bool result_with_default:
-        :rtype: NumbersDict
+        :param result_with_default:
+        :return: new NumbersDict with bool values
         """
 
         def op(a, b):
@@ -2223,19 +2223,17 @@ class NumbersDict:
             res.value = None
         return res
 
-    def __eq__(self, other):
+    def __eq__(self, other) -> bool:
         """
         :param NumbersDict|T other:
         :return: whether self == other elemwise. see self.elem_eq
-        :rtype: bool
         """
         return all(self.elem_eq(other).values())
 
-    def __ne__(self, other):
+    def __ne__(self, other) -> bool:
         """
         :param NumbersDict|T other:
         :return: not (self == other)
-        :rtype: bool
         """
         return not (self == other)
 
@@ -2244,11 +2242,10 @@ class NumbersDict:
         # and it would just confuse.
         raise Exception("%s.__cmp__ is undefined" % self.__class__.__name__)
 
-    def any_compare(self, other, cmp):
+    def any_compare(self, other, cmp) -> bool:
         """
         :param NumbersDict other:
         :param ((object,object)->True) cmp:
-        :rtype: True
         """
         for key in self.keys():
             if key in other.keys():
@@ -2281,11 +2278,11 @@ class NumbersDict:
         return min(*args)
 
     @classmethod
-    def max(cls, items):
+    def max(cls, items) -> NumbersDict:
         """
         Element-wise maximum for item in items.
+
         :param list[NumbersDict|int|float] items:
-        :rtype: NumbersDict
         """
         assert items
         if len(items) == 1:
@@ -2295,11 +2292,10 @@ class NumbersDict:
         return cls.max([items[0], cls.max(items[1:])])
 
     @classmethod
-    def min(cls, items):
+    def min(cls, items) -> NumbersDict:
         """
         Element-wise minimum for item in items.
         :param list[NumbersDict|int|float] items:
-        :rtype: NumbersDict
         """
         assert items
         if len(items) == 1:
@@ -2325,7 +2321,7 @@ class NumbersDict:
         """
         return min(self.values())
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         if self.value is None and not self.dict:
             return "%s()" % self.__class__.__name__
         if self.value is None and self.dict: