59 lines
No EOL
1.9 KiB
Python
59 lines
No EOL
1.9 KiB
Python
import numpy as np
|
|
import zlib
|
|
|
|
import numpy as np
|
|
import zlib
|
|
import struct
|
|
from typing import Any, Callable
|
|
|
|
def crc32_array(a: np.ndarray) -> int:
|
|
if not a.size: return 0
|
|
ac = np.ascontiguousarray(a)
|
|
mv = memoryview(ac.data).cast("B")
|
|
return zlib.adler32(mv)
|
|
|
|
def crc32_list(l: list[np.ndarray]) -> int:
|
|
if not l: return 0
|
|
h = 0
|
|
for a in l:
|
|
h = zlib.adler32(struct.pack(">I", crc32_array(a)), h)
|
|
return h
|
|
|
|
def ndarray_lru_cache(maxsize=128):
|
|
def decorator(func: Callable):
|
|
cache = {}
|
|
order = []
|
|
def wrapper(*args, **kwargs):
|
|
# Create a key by hashing any ndarray arguments
|
|
key_parts = []
|
|
for arg in args:
|
|
if isinstance(arg, np.ndarray):
|
|
key_parts.append(crc32_array(arg))
|
|
elif isinstance(arg, list) and arg and isinstance(arg[0], np.ndarray):
|
|
key_parts.append(crc32_list(arg))
|
|
else:
|
|
key_parts.append(arg)
|
|
for k, v in sorted(kwargs.items()):
|
|
if isinstance(v, np.ndarray):
|
|
key_parts.append((k, crc32_array(v)))
|
|
elif isinstance(v, list) and v and isinstance(v[0], np.ndarray):
|
|
key_parts.append((k, crc32_list(v)))
|
|
else:
|
|
key_parts.append((k, v))
|
|
|
|
key = tuple(key_parts)
|
|
if key in cache:
|
|
# Move to end for LRU
|
|
if key in order: order.remove(key)
|
|
order.append(key)
|
|
return cache[key]
|
|
|
|
result = func(*args, **kwargs)
|
|
cache[key] = result
|
|
order.append(key)
|
|
if len(order) > maxsize:
|
|
oldest = order.pop(0)
|
|
del cache[oldest]
|
|
return result
|
|
return wrapper
|
|
return decorator |