Repository: huggingface/transformers
File: src/transformers/models/zamba/modeling_zamba.py (Python)

Code before:
@auto_docstring
class ZambaModel(ZambaPreTrainedModel):
    """
    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`ZambaDecoderLayer`]

    Args:
        config: ZambaConfig
    """

    def __init__(self, config: ZambaConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers_block_type = config.layers_block_type
        layers = []
        self._tied_weights_keys = None
        for layer_id, layer_type in enumerate(self.layers_block_type):
            mamba = ZambaMambaDecoderLayer(config, layer_idx=layer_id)
            if layer_type == "hybrid":
                linear = nn.Linear(self.config.hidden_size, self.config.hidden_size, bias=False)
                layers.append(ZambaHybridLayer(ZambaAttentionDecoderLayer(config), linear, mamba))
                if self._tied_weights_keys is None:
                    self._tied_weights_keys = {
                        rf"layers.(?!{layer_id}\.)\d+.shared_transf": f"layers.{layer_id}.shared_transf"
                    }
            else:
                layers.append(mamba)
        print()
        self.layers = nn.ModuleList(layers)

        self._attn_implementation = config._attn_implementation
        self.final_layernorm = ZambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()

    @auto_docstring
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: ZambaHybridDynamicCache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,

Review comment (suggestion):
it was not matching `layer.25.shared_transform` when we need to tie layer 25 to layer 2. Adding a "dot" amkes sure that the pattern matches all layers except for "2"

Code after:
@auto_docstring
class ZambaModel(ZambaPreTrainedModel):
    """
    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`ZambaDecoderLayer`]

    Args:
        config: ZambaConfig
    """

    def __init__(self, config: ZambaConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers_block_type = config.layers_block_type
        layers = []
        self._tied_weights_keys = None
        for layer_id, layer_type in enumerate(self.layers_block_type):
            mamba = ZambaMambaDecoderLayer(config, layer_idx=layer_id)
            if layer_type == "hybrid":
                linear = nn.Linear(self.config.hidden_size, self.config.hidden_size, bias=False)
                layers.append(ZambaHybridLayer(ZambaAttentionDecoderLayer(config), linear, mamba))
                if self._tied_weights_keys is None:
                    self._tied_weights_keys = {
                        rf"layers.(?!{layer_id}\.)\d+.shared_transf": f"layers.{layer_id}.shared_transf"
                    }
            else:
                layers.append(mamba)
        self.layers = nn.ModuleList(layers)

        self._attn_implementation = config._attn_implementation
        self.final_layernorm = ZambaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()

    @auto_docstring
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: ZambaHybridDynamicCache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        cache_position: torch.LongTensor | None = None,