grunet使用例子

import time

import torch
import torch.nn as nn


import torch.nn.functional as F

import torch
import torch.nn as nn
from torch import Tensor
from typing import Optional

from torch.nn import Linear


class MultiHeadAttention(nn.Module):
    def __init__(self, n_state: int, n_head: int):
        super().__init__()
        self.n_head = n_head
        self.query = Linear(n_state, n_state)
        self.key = Linear(n_state, n_state, bias=False)
        self.value = Linear(n_state, n_state)
        self.out = Linear(n_state, n_state)

    def forward(self, x: Tensor, xa: Optional[Tensor] = None, mask: Optional[Tensor] = None, kv_cache: Optional[dict] = None, ):
        q = self.query(x)

        if kv_cache i