mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
345bea01dc
This is superior to lru_cache because (1) it's more explicit and (2) it doesn't leak the original function after it's been forced. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/132407 Approved by: https://github.com/albanD
29 lines
625 B
Python
29 lines
625 B
Python
from typing import Callable, Generic, Optional, TypeVar
|
|
|
|
|
|
R = TypeVar("R")
|
|
|
|
|
|
class Thunk(Generic[R]):
|
|
"""
|
|
A simple lazy evaluation implementation that lets you delay
|
|
execution of a function. It properly handles releasing the
|
|
function once it is forced.
|
|
"""
|
|
|
|
f: Optional[Callable[[], R]]
|
|
r: Optional[R]
|
|
|
|
__slots__ = ["f", "r"]
|
|
|
|
def __init__(self, f: Callable[[], R]):
|
|
self.f = f
|
|
self.r = None
|
|
|
|
def force(self) -> R:
|
|
if self.f is None:
|
|
return self.r # type: ignore[return-value]
|
|
self.r = self.f()
|
|
self.f = None
|
|
return self.r
|