returnn 1.20250113.120418__tar.gz → 1.20250113.193158__tar.gz

{returnn-1.20250113.120418/returnn.egg-info → returnn-1.20250113.193158}/PKG-INFO

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

returnn-1.20250113.193158/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250113.193158'
+long_version = '1.20250113.193158+git.777b035'

returnn-1.20250113.193158/returnn/frontend/_cache.py

@@ -0,0 +1,208 @@
+"""
+Cache, to store some data.
+See :class:`Cache`.
+
+One use case example is :func:`sinusoidal_positional_encoding` and :func:`relative_positional_encoding`.
+"""
+
+from __future__ import annotations
+from typing import Optional, Union, Any, Type, Callable, Tuple, Dict
+from weakref import ref
+import tree
+from returnn.util.lru_cache import lru_cache
+from returnn.tensor import Tensor, Dim
+import returnn.frontend as rf
+from returnn.frontend._backend import global_backend, get_backend_by_raw_tensor_type, Backend
+
+
+class Cache:
+    """
+    Cache, intended for internal use of RF functions.
+
+    One use case example is :func:`sinusoidal_positional_encoding` and :func:`relative_positional_encoding`.
+
+    There are some specific properties we must take care of:
+
+    - Lifetime of values: For graph-based backends, it can only stay alive for the current run ctx.
+      (For eager-based backends, there is no such restriction.)
+    - Size: Put some limit, use LRU logic.
+    - Dims: Use only weakrefs. Some Dim should not stay alive just because of the cache.
+    - Scalar dynamic Dims in eager mode, or static dims: Instead of the Dim, use the dim value for the key
+      (and map the output to the Dim).
+    - Tensor as keys: Use weakrefs. Also don't check by value but by identity.
+    """
+
+    def __init__(self, max_size: int):
+        # Use lru_cache here, but not via a decorator,
+        # as we want custom set/get logic.
+        # Also, we want the lru_cache to be local to this Cache instance,
+        # not shared over all instances of this class.
+        self._lru_cache = lru_cache(max_size)(_lru_cache_dummy_func)
+
+    def get(self, key, default=None):
+        """
+        :param key:
+        :param default:
+        :return: entry in cache or default
+        """
+        key_transformed = _transform_key(key)
+        key_transformed_orig, value = self._lru_cache.cache_peek(key_transformed, fallback=(None, None))
+        if key_transformed_orig is None:
+            return default
+
+        assert len(key_transformed_orig) == len(key_transformed)
+        dim_map = {}  # orig -> new
+        for key_item_orig, key_item in zip(key_transformed_orig, key_transformed):
+            if isinstance(key_item_orig, DimWrapper):
+                assert isinstance(key_item, DimWrapper)
+                dim_orig = key_item_orig.dim_ref()
+                dim = key_item.dim_ref()
+                assert isinstance(dim_orig, Dim) and isinstance(dim, Dim)
+                dim_map[dim_orig] = dim
+
+        # noinspection PyShadowingNames
+        def _map_output(output):
+            if isinstance(output, Dim):
+                return dim_map.get(output, output)
+            if isinstance(output, Tensor):
+                if any(dim in dim_map for dim in output.dims):
+                    out_raw = output.raw_tensor
+                    for axis, dim in enumerate(output.dims):
+                        if dim in dim_map:
+                            output = output.copy_template_replace_dim_tag(axis=axis, new_dim_tag=dim_map[dim])
+                    output.raw_tensor = out_raw
+            return output
+
+        return tree.map_structure(_map_output, value)
+
+    def set(self, key, value):
+        """
+        :param key:
+        :param value:
+        """
+
+        def _finalize_callback(*_args):
+            self._lru_cache.cache_pop(key_transformed, fallback=None)
+
+        key_backend = _get_backend(key)
+        value_backend = _get_backend(value)
+        if key_backend != value_backend:
+            raise ValueError(f"key and value have different backends: {key_backend} != {value_backend}")
+        key_transformed = _transform_key(key, finalize_callback=_finalize_callback)
+        self._lru_cache.cache_set(key_transformed, result=(key_transformed, value))
+
+
+def _lru_cache_dummy_func(*_args, **_kwargs):
+    raise Exception("This should not be called.")
+
+
+def _transform_key(
+    key: Any, *, finalize_callback: Optional[Callable] = None, collected_dim_map: Optional[Dict[Dim, DimWrapper]] = None
+) -> Tuple[Union[Type[Backend], ref[rf.RunCtx], _KeyItemType], ...]:
+    backend = _get_backend(key)
+    keys_flat = [backend]
+    if not backend.executing_eagerly():
+        # See comment above: If graph-mode, the cached value becomes invalid
+        # when the current run ctx goes out of scope.
+        keys_flat.append(ref(rf.get_run_ctx(), finalize_callback))
+    if collected_dim_map is None:
+        collected_dim_map = {}
+    keys_flat += [
+        _transform_key_item(key, finalize_callback=finalize_callback, collected_dim_map=collected_dim_map)
+        for key in tree.flatten(key)
+    ]
+    return tuple(keys_flat)
+
+
+def _transform_key_item(
+    key: Any, *, finalize_callback: Optional[Callable] = None, collected_dim_map: Dict[Dim, DimWrapper]
+) -> _KeyItemType:
+    if isinstance(key, Tensor):
+        return TensorWrapper(key, finalize_callback=finalize_callback)
+    if isinstance(key, Dim):
+        if key in collected_dim_map:
+            return collected_dim_map[key]
+        dim_wrapper = DimWrapper(key, finalize_callback=finalize_callback)
+        collected_dim_map[key] = dim_wrapper
+        return dim_wrapper
+    if not isinstance(key, _RawTypes):
+        raise TypeError(f"unexpected type {type(key)}")
+    return key
+
+
+def _get_backend(*args) -> Type[Backend]:
+    args_flat = tree.flatten(args)
+    for arg in args_flat:
+        if isinstance(arg, Tensor) and arg.raw_tensor is not None:
+            return get_backend_by_raw_tensor_type(type(arg.raw_tensor))
+    return global_backend.__class__
+
+
+class TensorWrapper:
+    """
+    Wraps :class:`Tensor`.
+    Using weakref for the tensor, including also ``raw_tensor``.
+    Equality is given if the identity is the same, for the Tensor itself and the raw_tensor.
+    No value of the tensor is checked.
+    """
+
+    def __init__(self, value: Tensor, *, finalize_callback):
+        self.value_ref = ref(value, finalize_callback)
+        self.raw_value_ref = ref(value.raw_tensor, finalize_callback)
+        self._hash = id(value)
+
+    def __eq__(self, other):
+        if isinstance(other, TensorWrapper):
+            return self.value_ref() is other.value_ref() and self.raw_value_ref() is other.raw_value_ref()
+        return False
+
+    def __hash__(self):
+        return self._hash
+
+
+class DimWrapper:
+    """
+    Wraps :class:`Dim`.
+    Using weakref for the dim.
+    If the size is scalar and known, equality is given when the size is equal (and dim tag is ignored)
+    """
+
+    def __init__(self, dim: Dim, *, finalize_callback):
+        self.dim_value = _dim_value_for_key(dim)
+        # finalize_callback only needed when we don't use the dim value.
+        self.dim_ref = ref(dim, finalize_callback if self.dim_value is None else None)
+        self.dyn_size_ref = (
+            # E.g. consider the batch dim or data spatial dim which would be reset each step.
+            # We need some ref to the dyn size, and finalize this key when it goes out of scope.
+            # This is only needed when there is no info on the static size (or eager scalar dyn size).
+            ref(dim.dyn_size_ext.raw_tensor, finalize_callback)
+            if self.dim_value is None and dim.dyn_size_ext and dim.dyn_size_ext.raw_tensor is not None
+            else None
+        )
+        self._hash = hash(dim) if self.dim_value is None else hash(self.dim_value)
+
+    def __eq__(self, other):
+        if isinstance(other, DimWrapper):
+            if self.dim_value is not None:
+                return self.dim_value == other.dim_value
+            return self.dim_ref() == other.dim_ref() and self.dyn_size_ref() is other.dyn_size_ref()
+        return False
+
+    def __hash__(self):
+        return self._hash
+
+
+def _dim_value_for_key(dim: Dim) -> Optional[int]:
+    if dim.size is not None:
+        return dim.size
+    if dim.dyn_size_ext and not dim.dyn_size_ext.dims:
+        if dim.dyn_size_ext.raw_tensor is not None:
+            # noinspection PyProtectedMember
+            if dim.dyn_size_ext._raw_backend.executing_eagerly():
+                return int(dim.get_dim_value())
+    return None
+
+
+# For now... we might extend it by some more types.
+_KeyItemType = Union[None, str, bool, int, float, TensorWrapper, DimWrapper]
+_RawTypes = (type(None), str, bool, int, float)

{returnn-1.20250113.120418 → returnn-1.20250113.193158}/returnn/frontend/attention.py

@@ -4,10 +4,10 @@ Attention
 
 from __future__ import annotations
 from typing import Tuple, Union, Optional, Sequence
-import weakref
 import logging
 from returnn.tensor import Tensor, Dim, single_step_dim
 import returnn.frontend as rf
+from returnn.frontend._cache import Cache
 
 
 __all__ = [
@@ -892,7 +892,7 @@ def _make_indices(
     return indices, out_spatial_dim
 
 
-_relative_positional_encoding_cache = weakref.WeakKeyDictionary()  # run ctx -> (spatial_dim, feat_dim) -> enc
+_relative_positional_encoding_cache = Cache(128)
 
 
 def relative_positional_encoding(
@@ -924,10 +924,10 @@ def relative_positional_encoding(
     """
     if not dtype:
         dtype = rf.get_default_float_dtype()
-    cache = _relative_positional_encoding_cache.setdefault(rf.get_run_ctx(), {})
     cache_key = (query_spatial_dim, key_value_spatial_dim, feat_dim, query_offset, dtype)
-    if cache_key in cache:
-        return cache[cache_key]
+    cache_entry = _relative_positional_encoding_cache.get(cache_key)
+    if cache_entry is not None:
+        return cache_entry
     import math
 
     with rf.control_flow_ctx(None):
@@ -946,11 +946,11 @@ def relative_positional_encoding(
             allow_missing_implicit_dims=True,
         )
         emb.feature_dim = feat_dim
-        cache[cache_key] = emb, out_spatial_dim
+        _relative_positional_encoding_cache.set(cache_key, (emb, out_spatial_dim))
         return emb, out_spatial_dim
 
 
-_sinusoidal_positional_encoding_cache = weakref.WeakKeyDictionary()  # run ctx -> (spatial_dim, feat_dim) -> enc
+_sinusoidal_positional_encoding_cache = Cache(128)  # (spatial_dim, feat_dim) -> enc
 
 
 def sinusoidal_positional_encoding(
@@ -982,10 +982,10 @@ def sinusoidal_positional_encoding(
         dtype = rf.get_default_float_dtype()
     if not device:
         device = rf.get_default_device()
-    cache = _sinusoidal_positional_encoding_cache.setdefault(rf.get_run_ctx(), {})
     cache_key = (spatial_dim, feat_dim, offset, base, dtype, device)
-    if cache_key in cache:
-        return cache[cache_key]
+    cache_entry = _sinusoidal_positional_encoding_cache.get(cache_key)
+    if cache_entry is not None:
+        return cache_entry
     import math
 
     with rf.control_flow_ctx(None):
@@ -1012,7 +1012,7 @@ def sinusoidal_positional_encoding(
             {spatial_dim, feat_dim} if spatial_dim != single_step_dim else {feat_dim}, allow_missing_implicit_dims=True
         )
         emb.feature_dim = feat_dim
-        cache[cache_key] = emb
+        _sinusoidal_positional_encoding_cache.set(cache_key, emb)
         return emb
 
 

returnn-1.20250113.193158/returnn/util/lru_cache.py

@@ -0,0 +1,309 @@
+"""
+:func:`lru_cache`, copied from Python functools, slightly adapted,
+and extended by functions to check whether some key is cached or not.
+"""
+
+from __future__ import annotations
+from typing import Dict, Any
+from functools import update_wrapper
+from threading import RLock
+from collections import namedtuple
+
+
+def lru_cache(maxsize: int = 128, typed: bool = False):
+    """Least-recently-used cache decorator.
+
+    If *maxsize* is set to None, the LRU features are disabled and the cache
+    can grow without bound.
+
+    If *typed* is True, arguments of different types will be cached separately.
+    For example, f(3.0) and f(3) will be treated as distinct calls with
+    distinct results.
+
+    Arguments to the cached function must be hashable.
+
+    Use f.cache_len() to see the current size of the cache.
+    Use f.cache_set(*args, result, **kwargs) to set a value in the cache directly.
+    Use f.cache_peek(*args, update_statistics=False, fallback=None, **kwargs)
+    to peek the cache, without ever calling the user function.
+    View the cache statistics named tuple (hits, misses, maxsize, currsize)
+    with f.cache_info().
+    Clear the cache and statistics with f.cache_clear().
+    Remove the oldest entry from the cache with f.cache_pop_oldest().
+    Take out some entry from the cache with f.cache_pop(*args, fallback=not_specified, **kwargs).
+    Set the maximum cache size to a new value with f.cache_set_maxsize(new_maxsize).
+    Access the underlying function with f.__wrapped__.
+
+    See:  https://en.wikipedia.org/wiki/Cache_replacement_policies#Least_recently_used_(LRU)
+
+    """
+
+    # Users should only access the lru_cache through its public API:
+    #       cache_info, cache_clear, and f.__wrapped__
+    # The internals of the lru_cache are encapsulated for thread safety and
+    # to allow the implementation to change (including a possible C version).
+
+    if isinstance(maxsize, int):
+        assert maxsize >= 0
+    elif callable(maxsize) and isinstance(typed, bool):
+        # The user_function was passed in directly via the maxsize argument
+        user_function, maxsize = maxsize, 128
+        return _lru_cache_wrapper(user_function, maxsize, typed)
+    elif maxsize is not None:
+        raise TypeError("Expected first argument to be an integer, a callable, or None")
+
+    # noinspection PyShadowingNames
+    def _decorating_function(user_function):
+        return _lru_cache_wrapper(user_function, maxsize, typed)
+
+    return _decorating_function
+
+
+_CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"])
+
+
+def _lru_cache_wrapper(user_function, maxsize: int, typed: bool):
+    # Constants shared by all lru cache instances:
+    make_key = _make_key  # build a key from the function arguments
+    # noinspection PyPep8Naming
+    PREV, NEXT, KEY, RESULT = 0, 1, 2, 3  # names for the link fields
+
+    cache: Dict[Any, list] = {}
+    hits = misses = 0
+    full = False
+    cache_get = cache.get  # bound method to lookup a key or return None
+    cache_len = cache.__len__  # get cache size without calling len()
+    lock = RLock()  # because linkedlist updates aren't threadsafe
+    root = []  # root of the circular doubly linked list
+    root[:] = [root, root, None, None]  # initialize by pointing to self
+
+    assert maxsize >= 0
+
+    def wrapper(*args, **kwds):
+        """
+        User-facing wrapper function.
+        """
+        # Size limited caching that tracks accesses by recency
+        nonlocal root, hits, misses, full
+        key = make_key(args, kwds, typed)
+        with lock:
+            link = cache_get(key)
+            if link is not None:
+                # Move the link to the front of the circular queue
+                link_prev, link_next, _key, result = link
+                link_prev[NEXT] = link_next
+                link_next[PREV] = link_prev
+                last = root[PREV]
+                last[NEXT] = root[PREV] = link
+                link[PREV] = last
+                link[NEXT] = root
+                hits += 1
+                return result
+            misses += 1
+        result = user_function(*args, **kwds)
+        if maxsize > 0:
+            _cache_insert(key, result)
+        return result
+
+    def _cache_insert(key, result):
+        nonlocal root, full
+        with lock:
+            if key in cache:
+                # Getting here means that this same key was added to the
+                # cache while the lock was released.  Since the link
+                # update is already done, we need only return the
+                # computed result and update the count of misses.
+                pass
+            elif full:
+                # Use the old root to store the new key and result.
+                oldroot = root
+                oldroot[KEY] = key
+                oldroot[RESULT] = result
+                # Empty the oldest link and make it the new root.
+                # Keep a reference to the old key and old result to
+                # prevent their ref counts from going to zero during the
+                # update. That will prevent potentially arbitrary object
+                # clean-up code (i.e. __del__) from running while we're
+                # still adjusting the links.
+                root = oldroot[NEXT]
+                oldkey = root[KEY]
+                root[KEY] = root[RESULT] = None
+                # Now update the cache dictionary.
+                del cache[oldkey]
+                # Save the potentially reentrant cache[key] assignment
+                # for last, after the root and links have been put in
+                # a consistent state.
+                cache[key] = oldroot
+            else:
+                # Put result in a new link at the front of the queue.
+                last = root[PREV]
+                link = [last, root, key, result]
+                last[NEXT] = root[PREV] = cache[key] = link
+                # Use the cache_len bound method instead of the len() function
+                # which could potentially be wrapped in an lru_cache itself.
+                full = cache_len() >= maxsize
+
+    def cache_info():
+        """Report cache statistics"""
+        with lock:
+            return _CacheInfo(hits, misses, maxsize, cache_len())
+
+    def cache_clear():
+        """Clear the cache and cache statistics"""
+        nonlocal hits, misses, full
+        with lock:
+            for link in cache.values():
+                link.clear()  # make GC happy
+            cache.clear()
+            root[:] = [root, root, None, None]
+            hits = misses = 0
+            full = False
+
+    def cache_parameters():
+        """
+        :return: parameters (maxsize, typed) of the cache as dict
+        """
+        return {"maxsize": maxsize, "typed": typed}
+
+    def cache_set(*args, result, **kwargs):
+        """
+        Sets a value in the cache directly.
+        """
+        nonlocal root, full
+        if maxsize > 0:
+            key = make_key(args, kwargs, typed)
+            _cache_insert(key, result)
+
+    def cache_peek(*args, update_statistics: bool = True, fallback=None, **kwargs):
+        """
+        Peeks the cache without ever calling the user function.
+        """
+        nonlocal hits, misses
+        key = make_key(args, kwargs, typed)
+        with lock:
+            link = cache_get(key)
+            if link is not None:
+                if update_statistics:
+                    hits += 1
+                return link[RESULT]
+            if update_statistics:
+                misses += 1
+            return fallback
+
+    not_specified = object()
+
+    def cache_pop(*args, fallback=not_specified, **kwargs):
+        """
+        Removes the entry from the cache.
+        """
+        nonlocal hits, misses
+        key = make_key(args, kwargs, typed)
+        with lock:
+            link = cache_get(key)
+            if link is not None:
+                # Take out link.
+                link[PREV][NEXT] = link[NEXT]
+                link[NEXT][PREV] = link[PREV]
+                oldkey = link[KEY]
+                oldvalue = link[RESULT]
+                link.clear()
+                del cache[oldkey]
+                return oldvalue
+            if fallback is not_specified:
+                raise KeyError("key not found")
+            return fallback
+
+    def cache_pop_oldest(*, fallback=not_specified):
+        """
+        Removes the oldest entry from the cache.
+        """
+        nonlocal root, full
+        with lock:
+            if not cache:
+                if fallback is not_specified:
+                    raise KeyError("cache is empty")
+                return fallback
+            assert cache
+            # Take out oldest link.
+            link: list = root[NEXT]
+            link[NEXT][PREV] = root
+            root[NEXT] = link[NEXT]
+            oldkey = link[KEY]
+            oldvalue = link[RESULT]
+            link.clear()
+            del cache[oldkey]
+            full = cache_len() >= maxsize
+            return oldvalue
+
+    def cache_set_maxsize(new_maxsize: int):
+        """
+        Resets the maxsize.
+        If the new maxsize is smaller than the current cache size, the oldest entries are removed.
+        """
+        nonlocal maxsize, full
+        assert new_maxsize >= 0
+        with lock:
+            maxsize = new_maxsize
+            while cache_len() > maxsize:
+                cache_pop_oldest()
+            full = cache_len() >= maxsize
+
+    wrapper.cache_info = cache_info
+    wrapper.cache_clear = cache_clear
+    wrapper.cache_parameters = cache_parameters
+    wrapper.cache_set = cache_set
+    wrapper.cache_peek = cache_peek
+    wrapper.cache_pop = cache_pop
+    wrapper.cache_len = cache_len
+    wrapper.cache_pop_oldest = cache_pop_oldest
+    wrapper.cache_set_maxsize = cache_set_maxsize
+
+    update_wrapper(wrapper, user_function)
+
+    return wrapper
+
+
+def _make_key(args, kwds, typed, *, _kwd_mark=(object(),), _fasttypes=(int, str), _tuple=tuple, _type=type, _len=len):
+    """Make a cache key from optionally typed positional and keyword arguments
+
+    The key is constructed in a way that is flat as possible rather than
+    as a nested structure that would take more memory.
+
+    If there is only a single argument and its data type is known to cache
+    its hash value, then that argument is returned without a wrapper.  This
+    saves space and improves lookup speed.
+
+    """
+    # All of code below relies on kwds preserving the order input by the user.
+    # Formerly, we sorted() the kwds before looping.  The new way is *much*
+    # faster; however, it means that f(x=1, y=2) will now be treated as a
+    # distinct call from f(y=2, x=1) which will be cached separately.
+    key = args
+    if kwds:
+        key += _kwd_mark
+        for item in kwds.items():
+            key += item
+    if typed:
+        key += _tuple(_type(v) for v in args)  # noqa
+        if kwds:
+            key += _tuple(_type(v) for v in kwds.values())  # noqa
+    elif _len(key) == 1 and _type(key[0]) in _fasttypes:
+        return key[0]
+    return _HashedSeq(key)
+
+
+class _HashedSeq(list):
+    """
+    This class guarantees that hash() will be called no more than once
+    per element.  This is important because the lru_cache() will hash
+    the key multiple times on a cache miss.
+    """
+
+    __slots__ = "hashvalue"
+
+    def __init__(self, tup, *, _hash=hash):
+        super().__init__(tup)
+        self.hashvalue = _hash(tup)
+
+    def __hash__(self):
+        return self.hashvalue

{returnn-1.20250113.120418 → returnn-1.20250113.193158/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250113.120418 → returnn-1.20250113.193158}/returnn.egg-info/SOURCES.txt

@@ -157,6 +157,7 @@ returnn/extern/graph_editor/transform.py
 returnn/extern/graph_editor/util.py
 returnn/frontend/__init__.py
 returnn/frontend/_backend.py
+returnn/frontend/_cache.py
 returnn/frontend/_numpy_backend.py
 returnn/frontend/_random_journal.py
 returnn/frontend/_utils.py
@@ -320,6 +321,7 @@ returnn/util/debug_helpers.py
 returnn/util/file_cache.py
 returnn/util/fsa.py
 returnn/util/literal_py_to_pickle.py
+returnn/util/lru_cache.py
 returnn/util/math.py
 returnn/util/multi_proc_non_daemonic_spawn.py
 returnn/util/native_code_compiler.py

returnn-1.20250113.120418/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250113.120418'
-long_version = '1.20250113.120418+git.0428e25'