Merge branch 'master' into moritz-shard-mgpu

Author: NeoLegends

.github/workflows/black.yml .github/workflows/ruff.yml.github/workflows/black.yml renamed to .github/workflows/ruff.yml

@@ -1,4 +1,4 @@
-name: black
+name: ruff
 
 on:
   push:
@@ -11,15 +11,14 @@ on:
       - master
 
 jobs:
-  check-black-formatting:
+  check-ruff-formatting:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
     - uses: actions/setup-python@v5
       with:
         python-version: 3.8
         cache: 'pip'
-        cache-dependency-path: '.github/workflows/black.yml'
-    - run: pip install black==22.3.0
-    - run: black --diff .
-    - run: black --check .
+        cache-dependency-path: '.github/workflows/ruff.yml'
+    - run: pip install ruff==0.11.8
+    - run: ruff format --diff .

CONTRIBUTING.md

@@ -18,7 +18,7 @@ General rules when contributing to the code of RETURNN:
   Our code style uses most common Python conventions.
   If you are not an expert in Python, use PyCharm,
   and follow [our PyCharm configuration guide](https://github.com/rwth-i6/returnn/wiki/PyCharm-Configuration).
-  Apply [black](https://black.readthedocs.io/).
+  Apply [ruff](https://github.com/astral-sh/ruff).
 * Make sure all [tests](https://returnn.readthedocs.io/en/latest/advanced/test_suite.html) pass.
 * At the time being, we want to support earlier versions of TF 1
   (consider at least TF 1.8, but maybe even TF 1.4)

__init__.py

@@ -7,7 +7,6 @@
 We want to support the same code.
 """
 
-
 from __future__ import annotations
 import os
 import sys

demos/demo-tf-search-compiled-graph.py

@@ -8,12 +8,12 @@
 
 # No RETURNN dependency needed for the basic search. Just TF itself.
 
-import typing
 import os
 import json
 import argparse
 import tensorflow as tf
 import numpy
+from typing import List, Optional, Tuple
 
 
 class Hyp:
@@ -26,7 +26,7 @@ def __init__(self, idx):
         :param int idx: hyp idx (to identify it in a beam)
         """
         self.idx = idx
-        self.source_idx = None  # type: typing.Optional[int]  # source hyp idx
+        self.source_idx: Optional[int] = None  # source hyp idx
         self.score = 0.0
         self.seq = []  # label seq
 
@@ -91,7 +91,6 @@ def make_initial_feed_dict():
         # Now loop over decoder steps.
         max_dec_len = 100  # TODO better default... depending on input len. or configurable...
         for i in range(max_dec_len):
-
             # Loop over all stochastic variables.
             for stochastic_var in info["stochastic_var_order"]:
                 assert isinstance(stochastic_var, str)
@@ -108,9 +107,7 @@ def make_initial_feed_dict():
                 # TODO: length norm here?
 
                 # Select new hypotheses.
-                best_possibilities = sorted(all_possibilities)[
-                    : args.beam_size
-                ]  # type: typing.List[typing.Tuple[float,int,Hyp]]
+                best_possibilities: List[Tuple[float, int, Hyp]] = sorted(all_possibilities)[: args.beam_size]
                 assert len(best_possibilities) == args.beam_size
                 hyps = [
                     hyp.expand(idx=i, label=label, score=score)
@@ -121,8 +118,9 @@ def make_initial_feed_dict():
                 session.run(
                     info["state_vars"]["stochastic_var_scores_%s" % stochastic_var] + "/Assign...?",  # TODO...
                     feed_dict={
-                        info["state_vars"]["stochastic_var_scores_%s" % stochastic_var]
-                        + "/Initial...?": [[hyp.seq[-1] for hyp in hyps]]  # TODO...
+                        info["state_vars"]["stochastic_var_scores_%s" % stochastic_var] + "/Initial...?": [
+                            [hyp.seq[-1] for hyp in hyps]
+                        ]  # TODO...
                     },
                 )
 

docs/conf.py

@@ -77,7 +77,7 @@
 ]
 
 # See https://github.com/rtfd/readthedocs.org/issues/283
-mathjax_path = "https://cdn.mathjax.org/mathjax/latest/MathJax.js?" "config=TeX-AMS-MML_HTMLorMML"
+mathjax_path = "https://cdn.mathjax.org/mathjax/latest/MathJax.js?config=TeX-AMS-MML_HTMLorMML"
 
 # see https://stackoverflow.com/q/12206334/562769
 numpydoc_show_class_members = False
@@ -151,6 +151,7 @@
 # If true, keep warnings as "system message" paragraphs in the built documents.
 # keep_warnings = False
 
+
 # Resolve function for the linkcode extension.
 def linkcode_resolve(domain, info):
     def find_source():

pyproject.toml

@@ -13,5 +13,9 @@ extend-exclude = '''
 )/
 '''
 
+[tool.ruff]
+line-length = 120
+target-version = "py38"  # https://github.com/rwth-i6/returnn/issues/1326
+
 [build-system]
 requires = ["setuptools", "numpy"]

requirements-dev

@@ -1,2 +1,2 @@
-black==22.3.0
 pytest
+ruff==0.11.8

returnn/datasets/basic.py

@@ -20,7 +20,7 @@
 import numpy
 import functools
 import typing
-from typing import TYPE_CHECKING, Optional, Any, Tuple, Union, Type, Dict, Sequence, List, Callable
+from typing import TYPE_CHECKING, Optional, Any, Set, Tuple, Union, Type, Dict, Sequence, List, Callable
 
 from returnn.log import log
 from returnn.engine.batch import Batch, BatchSetGenerator
@@ -147,12 +147,10 @@ def __init__(
         :param shard_index: local shard index, when sharding is enabled
         """
         self.name = name or ("dataset_id%s" % id(self))
-        self.lock = None  # type: Optional[RLock]  # Used when manipulating our data potentially from multiple threads.
-        self.rnd_seq_drop = None  # type: typing.Optional[Random]
+        self.lock: Optional[RLock] = None  # Used when manipulating our data potentially from multiple threads.
+        self.rnd_seq_drop: Optional[Random] = None
         self.num_inputs = 0  # usually not used, but num_outputs instead, which is more generic
-        self.num_outputs = (
-            None
-        )  # type: typing.Optional[typing.Dict[str,typing.Tuple[int,int]]]  # tuple is num-classes, len(shape).  # nopep8
+        self.num_outputs: Optional[Dict[str, Tuple[int, int]]] = None  # tuple is num-classes, len(shape).
         self.window = window
         self.seq_ordering = seq_ordering  # "default", "sorted" or "random". See self.get_seq_order_for_epoch().
         self.fixed_random_seed = fixed_random_seed
@@ -165,10 +163,10 @@ def __init__(
         self._seq_order_seq_lens_file = seq_order_seq_lens_file
         self._seq_order_seq_lens_by_idx = None
         # There is probably no use case for combining the two, so avoid potential misconfiguration.
-        assert (
-            self.partition_epoch == 1 or self.repeat_epoch == 1
-        ), "Combining partition_epoch and repeat_epoch is prohibited."
-        self.labels = {}  # type: typing.Dict[str,typing.List[str]]
+        assert self.partition_epoch == 1 or self.repeat_epoch == 1, (
+            "Combining partition_epoch and repeat_epoch is prohibited."
+        )
+        self.labels: Dict[str, List[str]] = {}
         self.weights = {}
         self._num_timesteps = 0
         self._num_seqs = 0
@@ -219,8 +217,8 @@ def __repr__(self):
             getattr(self, "epoch", "<unknown>"),
         )
 
-    _getnewargs_exclude_attrs = set()  # type: typing.Set[str]
-    _getnewargs_remap = {}  # type: typing.Dict[str,str]
+    _getnewargs_exclude_attrs: Set[str] = set()
+    _getnewargs_remap: Dict[str, str] = {}
 
     @staticmethod
     def _create_from_reduce(cls, kwargs, state) -> Dataset:
@@ -697,12 +695,13 @@ def get_seq_order_for_epoch(
             )
             old_seq_index = seq_index
             seq_index = [i for i in seq_index if all_seq_tags[i] in self.seq_tags_filter]
-            assert (
-                seq_index
-            ), "%s: empty after applying seq_list_filter_file. Example filter tags: %r, used tags: %r" % (
-                self,
-                sorted(self.seq_tags_filter)[:3],
-                [all_seq_tags[i] for i in old_seq_index[:3]],
+            assert seq_index, (
+                "%s: empty after applying seq_list_filter_file. Example filter tags: %r, used tags: %r"
+                % (
+                    self,
+                    sorted(self.seq_tags_filter)[:3],
+                    [all_seq_tags[i] for i in old_seq_index[:3]],
+                )
             )
         return seq_index
 
@@ -773,9 +772,9 @@ def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
         """
         self.epoch = epoch
         self.rnd_seq_drop = Random(self._get_random_seed_for_epoch(epoch=epoch))
-        assert (
-            self.num_shards == 1 or self.supports_sharding()
-        ), f"{self}: does not support sharding, but got num_shards == {self.num_shards}"
+        assert self._num_shards == 1 or self.supports_sharding(), (
+            f"{self}: does not support sharding, but got num_shards == {self._num_shards}"
+        )
         return False
 
     def finish_epoch(self, *, free_resources: bool = False):
@@ -1007,16 +1006,16 @@ def get_complete_frac(self, sorted_seq_idx: int, *, allow_only_lr_suitable: bool
             except Exception:  # also not always available
                 num_seqs = None  # ignore
 
-        if math.isinf(num_seqs):
+        if num_seqs is not None and math.isinf(num_seqs):
             if allow_only_lr_suitable:
                 # cannot compute meaningful complete_frac for infinite num_seqs
                 return None
             else:
                 num_seqs = None
 
-        assert (
-            num_seqs is None or 0 <= sorted_seq_idx < num_seqs
-        ), f"{self}: invalid seq indices: 0 <= seq_idx ({sorted_seq_idx}) < num_seqs ({num_seqs}) violated"
+        assert num_seqs is None or 0 <= sorted_seq_idx < num_seqs, (
+            f"{self}: invalid seq indices: 0 <= seq_idx ({sorted_seq_idx}) < num_seqs ({num_seqs}) violated"
+        )
         return self.generic_complete_frac(sorted_seq_idx, num_seqs)
 
     @property

returnn/datasets/cached.py

@@ -46,9 +46,10 @@ def __init__(self, cache_byte_size=0, **kwargs):
         self._index_map = range(len(self._seq_index))  # sorted seq idx -> seq_index idx
         self._tag_idx = {}  # type: typing.Dict[str,int]  # map of tag -> real-seq-idx. call _update_tag_idx
         self.targets = {}
-        self.target_keys = (
-            []
-        )  # the keys for which we provide data; we may have labels for additional keys in self.labels
+        # the keys for which we provide data;
+        # we may have labels for additional keys in self.labels
+        self.target_keys = []
+
         self.timestamps = None
 
     def initialize(self):

returnn/datasets/distrib_files.py

@@ -442,8 +442,7 @@ def _distribute_evenly_by_size(
             # We need to decide where to add this file, to the current or the next sub epoch.
             if not files_per_bin[bin_idx] or (
                 # Better to add this file to the current sub epoch?
-                abs((size_taken + size) - avg_size_per_sub_epoch)
-                <= abs(size_taken - avg_size_per_sub_epoch)
+                abs((size_taken + size) - avg_size_per_sub_epoch) <= abs(size_taken - avg_size_per_sub_epoch)
             ):
                 files_per_bin[bin_idx].append(f_tree)
                 size_taken = 0