Repository: huggingface/transformers
File: src/transformers/models/kyutai_speech_to_text/modeling_kyutai_speech_to_text.py (Python)

Code before:

        self.padding_cache = [None] * num_layers

    def _cache_init(self, hidden_states: torch.Tensor, layer_idx: int):
        """
        Initialize the cache for a specific layer.

        Parameters:
            hidden_states (`torch.Tensor`):
                The hidden states to initialize the cache with.
            layer_idx (`int`):
                The index of the layer to initialize the cache for.
        Returns:
            `torch.Tensor`, the initialized cache.
        """
        batch_size, dtype, device = hidden_states.shape[0], hidden_states.dtype, hidden_states.device
        padding, padding_mode, in_channels = (
            self.per_layer_padding[layer_idx],
            self.per_layer_padding_mode[layer_idx],
            self.per_layer_in_channels[layer_idx],
        )

        if padding_mode == "constant":
            current_cache = torch.zeros(batch_size, in_channels, padding, device=device, dtype=dtype)
        elif padding_mode == "replicate":
            current_cache = (
                torch.ones(batch_size, in_channels, padding, device=device, dtype=dtype) * hidden_states[..., :1]
            )
        else:
            raise NotImplementedError(f"Padding mode {padding_mode} not supported")

        return current_cache

    def update(self, hidden_states: torch.Tensor, layer_idx: int):
        """
        Updates the padding cache with the new padding states for the layer `layer_idx` and returns the current cache.

        Parameters:
            hidden_states (`torch.Tensor`):
                The hidden states to be partially cached.
            layer_idx (`int`):
                The index of the layer to cache the states for.
        Returns:
            `torch.Tensor` or `None`, the current padding cache.
        """
        batch_size, dtype, device = hidden_states.shape[0], hidden_states.dtype, hidden_states.device
        padding, in_channels = self.per_layer_padding[layer_idx], self.per_layer_in_channels[layer_idx]

        if self.padding_cache[layer_idx] is None:
            current_cache = self._cache_init(hidden_states, layer_idx)
        else:

Review comment (performance):
Good idea to add `_cache_init` so that different models can initialize differently if needed

How about removing the check [here](https://github.com/huggingface/transformers/blob/main/src/transformers/models/mimi/modeling_mimi.py#L100-L103) inside init? Since you have a similar one in this new method?

Code after:
        """
        Initialize the cache for a specific layer.

        Parameters:
            hidden_states (`torch.Tensor`):
                The hidden states to initialize the cache with.
            layer_idx (`int`):
                The index of the layer to initialize the cache for.
        Returns:
            `torch.Tensor`, the initialized cache.
        """
        batch_size, dtype, device = hidden_states.shape[0], hidden_states.dtype, hidden_states.device
        padding, padding_mode, in_channels = (
            self.per_layer_padding[layer_idx],
            self.per_layer_padding_mode[layer_idx],
            self.per_layer_in_channels[layer_idx],
        )

        if padding_mode == "constant":
            current_cache = torch.zeros(batch_size, in_channels, padding, device=device, dtype=dtype)
        elif padding_mode == "replicate":
            current_cache = (
                torch.ones(batch_size, in_channels, padding, device=device, dtype=dtype) * hidden_states[..., :1]
            )
        else:
            raise NotImplementedError(f"Padding mode {padding_mode} not supported")

        return current_cache

    def update(self, hidden_states: torch.Tensor, layer_idx: int):
        """
        Updates the padding cache with the new padding states for the layer `layer_idx` and returns the current cache.

        Parameters:
            hidden_states (`torch.Tensor`):
                The hidden states to be partially cached.
            layer_idx (`int`):
                The index of the layer to cache the states for.
        Returns:
            `torch.Tensor` or `None`, the current padding cache.
        """
        batch_size, dtype, device = hidden_states.shape[0], hidden_states.dtype, hidden_states.device
        padding, in_channels = self.per_layer_padding[layer_idx], self.per_layer_in_channels[layer_idx]

        if self.padding_cache[layer_idx] is None:
            current_cache = self._cache_init(hidden_states, layer_idx)
        else:
            current_cache = self.padding_cache[layer_idx]

        # update the cache
        if padding > 0: