def __init__(
        self,
        length=5,
        embed_dim=768,
        embedding_key="mean",
        prompt_init="uniform",
        prompt_pool=False,
        prompt_key=False,
        pool_size=None,
        top_k=None,
        batchwise_prompt=False,
        prompt_key_init="uniform",
    ):
        """
        Args:
            length (int): length of the prompt. Default 5.
            embed_dim (int): embedding dimension of the prompt. Default 768.
            embedding_key (str): method to generate embedding to find key \
                                similary. Default "mean".
            prompt_init (str): initialization of the prompt pool. \
                                Default "uniform".
            prompt_pool (bool): use prompt pool or not. Default False.
            prompt_key (bool): use learnable prompt keys. Default False.
            pool_size (int): size of the pool.
            top_k (int): select the top k similar prompts.
            batchwise_prompt (bool): use prompt batchwise. Defalt False.
            prompt_key_init (str): initialization of the key pool. \
                                Default "uniform",
        """
        super().__init__()

        self.length = length
        self.embed_dim = embed_dim
        self.prompt_pool = prompt_pool
        self.embedding_key = embedding_key
        self.prompt_init = prompt_init
        self.prompt_key = prompt_key
        self.pool_size = pool_size
        self.top_k = top_k
        self.batchwise_prompt = batchwise_prompt

        if self.prompt_pool:
            prompt_pool_shape = (pool_size, length, embed_dim)
            if prompt_init == "zero":
                self.prompt = nn.Parameter(torch.zeros(prompt_pool_shape))
            elif prompt_init == "uniform":
                self.prompt = nn.Parameter(torch.randn(prompt_pool_shape))
                nn.init.uniform_(self.prompt, -1, 1)

        # if using learnable prompt keys
        if prompt_key:
            key_shape = (pool_size, embed_dim)
            if prompt_key_init == "zero":
                self.prompt_key = nn.Parameter(torch.zeros(key_shape))
            elif prompt_key_init == "uniform":
                self.prompt_key = nn.Parameter(torch.randn(key_shape))
                nn.init.uniform_(self.prompt_key, -1, 1)
        else:
            # else use mean of prompt as key
            # only compatible with prompt, not prefix
            prompt_mean = torch.mean(self.prompt, dim=1)
            self.prompt_key = prompt_mean

        Args:
            length (int): length of the prompt. Default 5.
            embed_dim (int): embedding dimension of the prompt. Default 768.
            embedding_key (str): method to generate embedding to find key                                 similary. Default "mean".
            prompt_init (str): initialization of the prompt pool.                                 Default "uniform".
            prompt_pool (bool): use prompt pool or not. Default False.
            prompt_key (bool): use learnable prompt keys. Default False.
            pool_size (int): size of the pool.
            top_k (int): select the top k similar prompts.
            batchwise_prompt (bool): use prompt batchwise. Defalt False.
            prompt_key_init (str): initialization of the key pool.                                 Default "uniform",
        
avalanche/models/prompt.py