Function cache_prefix

gpu/model.py:346–366

Take a prefix view of a larger cache. The original cache object remains of identical size and valid after the shrinked alias has been used. This function is useful when a cache was allocated for a larger batch size than what is necessary. Args: cache: the

(cache: list[LayerCache], length: int)

def cache_prefix(cache: list[LayerCache], length: int) -> list[LayerCache]:
    """
    Take a prefix view of a larger cache.

    The original cache object remains of identical size and valid
    after the shrinked alias has been used. This function is useful
    when a cache was allocated for a larger batch size than what is
    necessary.

    Args:
        cache: the cache to take a view in.
        length (int): the desired length

    Returns:
        A view in the input cache object.
    """

    if len(cache) > 0:
        assert cache[0][0].shape[1] >= length

    return [(ck[:, :length], cv[:, :length]) for ck, cv in cache]