Monkey-patch DynamicCache with convenience methods used by KVCOMM.
()
| 503 | |
| 504 | |
| 505 | def _install_dynamic_cache_extensions() -> None: |
| 506 | """Monkey-patch DynamicCache with convenience methods used by KVCOMM.""" |
| 507 | if getattr(DynamicCache, "_kvcomm_extensions_installed", False): |
| 508 | return |
| 509 | |
| 510 | DynamicCache._normalize_slice_indices = lambda self, start=None, end=None: _normalize_indices(self, start, end) |
| 511 | DynamicCache.slice_ = lambda self, start=None, end=None: _slice_inplace(self, start, end) |
| 512 | DynamicCache.slice = lambda self, start=None, end=None: _slice_functional(self, start, end) |
| 513 | DynamicCache.concat_ = lambda self, other: _concat_inplace(self, _ensure_cache_sequence(other)) |
| 514 | DynamicCache.concat = lambda self, other: _concat_functional(self, _ensure_cache_sequence(other)) |
| 515 | DynamicCache.replace_ = lambda self, start, end, real: _replace_inplace(self, start, end, real) |
| 516 | DynamicCache.replace = lambda self, start, end, real: _replace_functional(self, start, end, real) |
| 517 | DynamicCache.select_indices = lambda self, indices: _select_indices(self, indices) |
| 518 | DynamicCache.to = lambda self, device: _to_device(self, device) |
| 519 | DynamicCache.copy = lambda self: _copy_cache(self) |
| 520 | DynamicCache.split_cache_by_placeholders = lambda self, placeholder_dict: _split_cache_by_placeholders( |
| 521 | self, placeholder_dict |
| 522 | ) |
| 523 | DynamicCache.__add__ = lambda self, other: _elementwise_binary_op(self, other, torch.add) |
| 524 | DynamicCache.__sub__ = lambda self, other: _elementwise_binary_op(self, other, torch.sub) |
| 525 | DynamicCache.split = lambda self, sizes: _split_cache(self, sizes) |
| 526 | DynamicCache._kvcomm_extensions_installed = True |
| 527 | |
| 528 | |
| 529 | _install_dynamic_cache_extensions() |
no test coverage detected