Function _make_random_keys

tests/test_compress.py:21–23

(seq_k: int, head_dim: int = 128)

def _make_random_keys(seq_k: int, head_dim: int = 128) -> torch.Tensor:
    K = torch.randn(seq_k, head_dim)
    return K.half()


@pytest.mark.parametrize("seq_k", [64, 128, 256, 512])