Apply an elementwise binary op to two caches layer-by-layer.
(
cache: DynamicCache,
other: DynamicCache,
op,
)
| 440 | |
| 441 | |
| 442 | def _elementwise_binary_op( |
| 443 | cache: DynamicCache, |
| 444 | other: DynamicCache, |
| 445 | op, |
| 446 | ) -> DynamicCache: |
| 447 | """Apply an elementwise binary op to two caches layer-by-layer.""" |
| 448 | _ensure_same_layout(cache, other) |
| 449 | base_stack = _stack_cache_tensors(cache) |
| 450 | other_stack = _stack_cache_tensors(other) |
| 451 | if base_stack is not None and other_stack is not None: |
| 452 | result = _copy_cache(cache) |
| 453 | key_stack = op(base_stack[0], other_stack[0]) |
| 454 | value_stack = op(base_stack[1], other_stack[1]) |
| 455 | _assign_stack_to_cache(result, key_stack, value_stack) |
| 456 | _set_seen_tokens(result, key_stack.shape[-2]) |
| 457 | return result |
| 458 | |
| 459 | result = type(cache)() |
| 460 | if _is_layered_cache(cache): |
| 461 | result.layers = [] |
| 462 | else: |
| 463 | result.key_cache = [] |
| 464 | result.value_cache = [] |
| 465 | for idx in range(_get_layer_count(cache)): |
| 466 | key_a, value_a = _get_layer_kv(cache, idx) |
| 467 | key_b, value_b = _get_layer_kv(other, idx) |
| 468 | if _layer_is_empty(key_a): |
| 469 | new_key = _clone_tensor_or_empty(key_b) |
| 470 | new_value = _clone_tensor_or_empty(value_b) |
| 471 | elif _layer_is_empty(key_b): |
| 472 | new_key = _clone_tensor_or_empty(key_a) |
| 473 | new_value = _clone_tensor_or_empty(value_a) |
| 474 | else: |
| 475 | new_key = op(key_a, key_b) |
| 476 | new_value = op(value_a, value_b) |
| 477 | if _is_layered_cache(cache): |
| 478 | layer = copy.deepcopy(cache.layers[idx]) |
| 479 | layer.keys = new_key |
| 480 | layer.values = new_value |
| 481 | if hasattr(layer, "is_initialized"): |
| 482 | layer.is_initialized = not _layer_is_empty(new_key) |
| 483 | if hasattr(layer, "dtype") and isinstance(new_key, torch.Tensor): |
| 484 | layer.dtype = new_key.dtype |
| 485 | if hasattr(layer, "device") and isinstance(new_key, torch.Tensor): |
| 486 | layer.device = new_key.device |
| 487 | result.layers.append(layer) |
| 488 | else: |
| 489 | result.key_cache.append(new_key) |
| 490 | result.value_cache.append(new_value) |
| 491 | _set_seen_tokens(result, _safe_seq_len(cache)) |
| 492 | return result |
| 493 | |
| 494 | |
| 495 | def _split_cache(cache: DynamicCache, sizes: Sequence[int]) -> List[DynamicCache]: |
no test coverage detected