(fn: Callable[P, T])
| 343 | @functools.wraps(functools.lru_cache) |
| 344 | def lru_cache_unless_export(maxsize=128, typed=False): |
| 345 | def outer_wrapper(fn: Callable[P, T]): |
| 346 | cached = functools.lru_cache(maxsize=maxsize, typed=typed)(fn) |
| 347 | if is_torch_version("<", "2.7.0"): |
| 348 | return cached |
| 349 | |
| 350 | @functools.wraps(fn) |
| 351 | def inner_wrapper(*args: P.args, **kwargs: P.kwargs): |
| 352 | compiler = getattr(torch, "compiler", None) |
| 353 | is_exporting = bool(compiler and hasattr(compiler, "is_exporting") and compiler.is_exporting()) |
| 354 | is_compiling = bool(compiler and hasattr(compiler, "is_compiling") and compiler.is_compiling()) |
| 355 | |
| 356 | # Fallback for older builds where compiler.is_compiling is unavailable. |
| 357 | if not is_compiling: |
| 358 | dynamo = getattr(torch, "_dynamo", None) |
| 359 | if dynamo is not None and hasattr(dynamo, "is_compiling"): |
| 360 | is_compiling = dynamo.is_compiling() |
| 361 | |
| 362 | if is_exporting or is_compiling: |
| 363 | return fn(*args, **kwargs) |
| 364 | return cached(*args, **kwargs) |
| 365 | |
| 366 | return inner_wrapper |
| 367 | |
| 368 | return outer_wrapper |
| 369 |
nothing calls this directly
no test coverage detected
searching dependent graphs…