optimum-rbln 0.8.0.post2__py3-none-any.whl → 0.8.1__py3-none-any.whl

optimum/rbln/transformers/models/siglip/configuration_siglip.py

@@ -18,12 +18,20 @@ from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNSiglipVisionModelConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNSiglipVisionModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized SigLIP vision models for image encoding in multimodal tasks.
+    """
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
         image_size: Optional[int] = None,
         interpolate_pos_encoding: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -33,6 +41,7 @@ class RBLNSiglipVisionModelConfig(RBLNModelConfig):
                 a tuple/list (height, width), or a dictionary with 'height' and 'width' keys.
             interpolate_pos_encoding (Optional[bool]): Whether to interpolate the position encoding.
             output_hidden_states: (Optional[bool]): Whether to return hidden states.
+            output_attentions: (Optional[bool]): Whether to return attentions.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -46,6 +55,7 @@ class RBLNSiglipVisionModelConfig(RBLNModelConfig):
         self.image_size = image_size
         self.interpolate_pos_encoding = interpolate_pos_encoding or False
         self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     @property
     def image_width(self):

optimum/rbln/transformers/models/siglip/modeling_siglip.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
 from transformers import SiglipVisionConfig, SiglipVisionModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling
-from transformers.models.siglip.modeling_siglip import SiglipVisionModelOutput
 
 from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
@@ -34,11 +33,18 @@ if TYPE_CHECKING:
 
 
 class _SiglipVisionModel(torch.nn.Module):
-    def __init__(self, model: SiglipVisionModel, interpolate_pos_encoding: bool, output_hidden_states: bool):
+    def __init__(
+        self,
+        model: SiglipVisionModel,
+        interpolate_pos_encoding: bool,
+        output_hidden_states: bool,
+        output_attentions: bool,
+    ):
         super().__init__()
         self.vision_model = model.vision_model
         self.interpolate_pos_encoding = interpolate_pos_encoding
         self.output_hidden_states = output_hidden_states
+        self.output_attentions = output_attentions
 
     def forward(self, inp):
         enc_out = self.vision_model(
@@ -46,16 +52,25 @@ class _SiglipVisionModel(torch.nn.Module):
             output_hidden_states=self.output_hidden_states,
             return_dict=False,
             interpolate_pos_encoding=self.interpolate_pos_encoding,
+            output_attentions=self.output_attentions,
         )
         return tuple(x for x in enc_out if x is not None)
 
 
 class RBLNSiglipVisionModel(RBLNModel):
+    """
+    RBLN optimized SigLIP vision model.
+
+    This class provides hardware-accelerated inference for SigLIP vision models
+    on RBLN devices, supporting image encoding for multimodal vision-language tasks.
+    """
+
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNSiglipVisionModelConfig) -> torch.nn.Module:
         wrapper_cfg = {
             "interpolate_pos_encoding": rbln_config.interpolate_pos_encoding,
             "output_hidden_states": rbln_config.output_hidden_states,
+            "output_attentions": rbln_config.output_attentions,
         }
         return _SiglipVisionModel(model, **wrapper_cfg).eval()
 
@@ -81,8 +96,10 @@ class RBLNSiglipVisionModel(RBLNModel):
         if rbln_config.image_size is None:
             raise ValueError("`rbln_image_size` should be specified!")
 
+        if rbln_config.output_attentions is None:
+            rbln_config.output_attentions = getattr(model_config, "output_attentions", False)
         if rbln_config.output_hidden_states is None:
-            rbln_config.output_hidden_states = model_config.output_hidden_states
+            rbln_config.output_hidden_states = getattr(model_config, "output_hidden_states", False)
 
         rbln_compile_config = RBLNCompileConfig(
             input_info=[
@@ -104,43 +121,74 @@ class RBLNSiglipVisionModel(RBLNModel):
 
     def forward(
         self,
-        pixel_values: Optional[torch.FloatTensor] = None,
+        pixel_values: torch.Tensor,
         return_dict: bool = None,
+        output_attentions: bool = None,
+        output_hidden_states: bool = None,
         interpolate_pos_encoding: bool = False,
-        **kwargs,
-    ) -> Union[Tuple, SiglipVisionModelOutput]:
-        if len(kwargs) > 0 and any(kwargs.values()):
-            logger.warning(f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__}.")
+        **kwargs: Dict[str, Any],
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
+            logger.warning(
+                f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."
+            )
+
+        output_attentions = output_attentions if output_attentions is not None else self.rbln_config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.rbln_config.output_hidden_states
+        )
+
+        if output_attentions != self.rbln_config.output_attentions:
+            raise ValueError(
+                f"Variable output_attentions {output_attentions} is not equal to rbln_config.output_attentions {self.rbln_config.output_attentions} "
+                f"Please compile again with the correct argument."
+            )
+
+        if output_hidden_states != self.rbln_config.output_hidden_states:
+            raise ValueError(
+                f"Variable output_hidden_states {output_hidden_states} is not equal to rbln_config.output_hidden_states {self.rbln_config.output_hidden_states} "
+                f"Please compile again with the correct argument."
+            )
 
         if interpolate_pos_encoding != self.rbln_config.interpolate_pos_encoding:
             raise ValueError(
-                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding}"
+                f"Variable interpolate_pos_encoding {interpolate_pos_encoding} is not equal to rbln_config.interpolate_pos_encoding {self.rbln_config.interpolate_pos_encoding} "
                 f"Please compile again with the correct argument."
             )
+
         output = super().forward(pixel_values, return_dict=return_dict)
         return output
 
     def _prepare_output(self, output, return_dict):
-        """
-        Prepare model output based on return_dict flag.
-        This method can be overridden by subclasses to provide task-specific output handling.
-        """
+        # Prepare model output based on return_dict flag.
+        # This method can be overridden by subclasses to provide task-specific output handling.
+
         if not return_dict:
             return (output,) if not isinstance(output, (tuple, list)) else output
         else:
-            last_hidden_state = (
-                output[0]
-                if self.rbln_config.interpolate_pos_encoding or self.rbln_config.output_hidden_states
-                else output
-            )
-            pooler_output = output[1] if self.rbln_config.interpolate_pos_encoding else None
+            last_hidden_state = output.pop(0) if isinstance(output, (tuple, list)) else output
+            vision_config = self.config.vision_config if hasattr(self.config, "vision_config") else self.config
+            pooler_output = output.pop(0) if getattr(vision_config, "vision_use_head", True) else None
+
             if self.rbln_config.output_hidden_states:
-                hidden_states = (output[2:] if self.rbln_config.interpolate_pos_encoding else output[1:],)
+                hidden_states = ()
+                num_hidden_layers = vision_config.num_hidden_layers
+                for _ in range(num_hidden_layers + 1):
+                    hidden_states += (output.pop(0),)
             else:
                 hidden_states = None
 
+            if self.rbln_config.output_attentions:
+                attentions = ()
+                num_hidden_layers = vision_config.num_hidden_layers
+                for _ in range(num_hidden_layers):
+                    attentions += (output.pop(0),)
+            else:
+                attentions = None
+
             return BaseModelOutputWithPooling(
                 last_hidden_state=last_hidden_state,
                 pooler_output=pooler_output,
                 hidden_states=hidden_states,
+                attentions=attentions,
             )

optimum/rbln/transformers/models/t5/configuration_t5.py

@@ -17,8 +17,18 @@ from ..seq2seq import RBLNModelForSeq2SeqLMConfig
 
 
 class RBLNT5EncoderModelConfig(RBLNTransformerEncoderForFeatureExtractionConfig):
-    pass
+    """
+    Configuration class for RBLNT5EncoderModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized T5 encoder models for feature extraction tasks.
+    """
 
 
 class RBLNT5ForConditionalGenerationConfig(RBLNModelForSeq2SeqLMConfig):
-    pass
+    """
+    Configuration class for RBLNT5ForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized T5 models for conditional text generation tasks.
+    """

optimum/rbln/transformers/models/t5/modeling_t5.py

@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Callable
 
 import torch
 from transformers import AutoModelForTextEncoding, T5EncoderModel, T5ForConditionalGeneration
+from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
 
 from ...modeling_generic import RBLNTransformerEncoderForFeatureExtraction
 from ...models.seq2seq import RBLNModelForSeq2SeqLM
@@ -41,8 +42,30 @@ class T5EncoderWrapper(torch.nn.Module):
 
 
 class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
+    """
+    The T5 Model transformer with an encoder-only architecture for feature extraction.
+    This model inherits from [`RBLNTransformerEncoderForFeatureExtraction`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model supports various sizes of the T5EncoderModel. For optimal performance, it is highly recommended to adjust the tensor parallelism setting
+        based on the model size. Please refer to the [Optimum RBLN Overview](../../../optimum_rbln.md) for guidance on choosing the appropriate tensor parallelism size for your model.
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNT5EncoderModel
+
+        model = RBLNT5EncoderModel.from_pretrained(
+            "sentence-transformers/sentence-t5-xxl",
+            export=True,
+            rbln_tensor_parallel_size=4,
+        )
+
+        model.save_pretrained("compiled-sentence-t5-xxl")
+        ```
+    """
+
     auto_model_class = AutoModelForTextEncoding
-    rbln_model_input_names = ["input_ids", "attention_mask"]
+    output_class = BaseModelOutputWithPastAndCrossAttentions
 
     @classmethod
     def wrap_model_if_needed(self, model: "PreTrainedModel", rbln_config: RBLNT5EncoderModelConfig):
@@ -50,18 +73,43 @@ class RBLNT5EncoderModel(RBLNTransformerEncoderForFeatureExtraction):
 
     @classmethod
     def update_rbln_config_using_pipe(
-        cls,
-        pipe: "RBLNDiffusionMixin",
-        rbln_config: "RBLNDiffusionMixinConfig",
-        submodule_name: str,
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
     ) -> "RBLNDiffusionMixinConfig":
-        submodule_config = getattr(rbln_config, submodule_name)
-        submodule_config.max_seq_len = rbln_config.max_seq_len or 256
-        submodule_config.model_input_names = ["input_ids"]
         return rbln_config
 
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        input_dict = {"input_ids": input_ids.long()}
+        if attention_mask is not None:
+            input_dict["attention_mask"] = attention_mask.long()
+
+        output = super().forward(**input_dict, **kwargs)
+        return output
+
 
 class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
+    """
+    The T5 Model transformer with a language modeling head for conditional generation.
+    This model inherits from [`RBLNModelForSeq2SeqLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    Important Note:
+        This model supports various sizes of the T5ForConditionalGeneration. For optimal performance, it is highly recommended to adjust the tensor parallelism setting
+        based on the model size. Please refer to the [Optimum RBLN Overview](../../../optimum_rbln.md) for guidance on choosing the appropriate tensor parallelism size for your model.
+
+
+    Examples:
+        ```python
+        from optimum.rbln import RBLNT5ForConditionalGeneration
+
+        model = RBLNT5ForConditionalGeneration.from_pretrained(
+            "google-t5/t5-11b",
+            export=True,
+            rbln_tensor_parallel_size=4,
+        )
+
+        model.save_pretrained("compiled-sentence-t5-xxl")
+        ```
+    """
+
     support_causal_attn = False
 
     @classmethod

optimum/rbln/transformers/models/t5/t5_architecture.py

@@ -136,10 +136,14 @@ class T5Decoder(Seq2SeqDecoder):
 
 
 class T5Block(Seq2SeqDecoderLayer):
+    def __init__(self, decoder_layer, self_attn):
+        super().__init__(decoder_layer, self_attn, cross_attn=None)
+        self.__post_init__()
+
     def __post_init__(self):
         self.self_attn_layer_norm = self._original_mod.layer[0].layer_norm
         self.encoder_attn_layer_norm = self._original_mod.layer[1].layer_norm
-        self.encoder_attn = T5CrossAttention(self._original_mod.layer[1].EncDecAttention)
+        self.cross_attn = T5CrossAttention(self._original_mod.layer[1].EncDecAttention)
         self.ff_layer = self._original_mod.layer[2]
 
     def pre_self_attn_layer_norm(self, hidden_states):

optimum/rbln/transformers/models/{time_series_transformers → time_series_transformer}/__init__.py

@@ -23,4 +23,4 @@
 
 from ....ops import paged_add_softmax_attn_decode, rbln_cache_update
 from .configuration_time_series_transformer import RBLNTimeSeriesTransformerForPredictionConfig
-from .modeling_time_series_transformers import RBLNTimeSeriesTransformerForPrediction
+from .modeling_time_series_transformer import RBLNTimeSeriesTransformerForPrediction

optimum/rbln/transformers/models/{time_series_transformers → time_series_transformer}/configuration_time_series_transformer.py

@@ -1,16 +1,23 @@
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNTimeSeriesTransformerForPredictionConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNTimeSeriesTransformerForPrediction.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Time Series Transformer models for time series forecasting tasks.
+    """
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
         enc_max_seq_len: Optional[int] = None,
         dec_max_seq_len: Optional[int] = None,
         num_parallel_samples: Optional[int] = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args:

optimum/rbln/transformers/models/{time_series_transformers/modeling_time_series_transformers.py → time_series_transformer/modeling_time_series_transformer.py}

@@ -120,6 +120,17 @@ class RBLNSeq2SeqTSDecoderOutput(ModelOutput):
 
 
 class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
+    """
+    The Time Series Transformer Model with a distribution head on top for time-series forecasting. e.g., for datasets like M4, NN5, or other time series forecasting benchmarks.
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformer-based `TimeSeriesTransformerForPrediction` models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers `TimeSeriesTransformerForPrediction` model into a RBLN transformer model by:
+
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN Compiler.
+    """
+
     auto_model_class = None
     main_input_name = "inputs_embeds"
 
@@ -144,11 +155,6 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         )
 
     def __getattr__(self, __name: str) -> Any:
-        """This is the key method to implement RBLN-TimeSeriesTransformersForPrediction.
-        Returns:
-            Any: TimeSeriesTransformersForPrediction's corresponding method
-        """
-
         def redirect(func):
             return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
 
@@ -188,15 +194,19 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
             if "key_value_states" in name:
                 context.mark_static_address(tensor)
 
-        compiled_decoder = super().compile(
+        compiled_decoder = cls.compile(
             wrapped_model.decoder,
             dec_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=dec_example_inputs,
             compile_context=context,
         )
-        compiled_encoder = super().compile(
+        compiled_encoder = cls.compile(
             wrapped_model.encoder,
             enc_compile_config,
+            create_runtimes=rbln_config.create_runtimes,
+            device=rbln_config.device,
             example_inputs=enc_example_inputs,
             compile_context=context,
         )
@@ -211,10 +221,9 @@ class RBLNTimeSeriesTransformerForPrediction(RBLNModel):
         subfolder: str,
         rbln_config: RBLNTimeSeriesTransformerForPredictionConfig,
     ):
-        """
-        If you are unavoidably running on a CPU rather than an RBLN device,
-        store the torch tensor, weight, etc. in this function.
-        """
+        # If you are unavoidably running on a CPU rather than an RBLN device,
+        # store the torch tensor, weight, etc. in this function.
+
         save_dict = {}
         save_dict["embedder"] = model.model.embedder.state_dict()
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")

optimum/rbln/transformers/models/vit/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configuration_vit import RBLNViTForImageClassificationConfig
+from .modeling_vit import RBLNViTForImageClassification
+
+
+__all__ = ["RBLNViTForImageClassificationConfig", "RBLNViTForImageClassification"]

optimum/rbln/transformers/models/vit/configuration_vit.py

@@ -0,0 +1,24 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...configuration_generic import RBLNModelForImageClassificationConfig
+
+
+class RBLNViTForImageClassificationConfig(RBLNModelForImageClassificationConfig):
+    """
+    Configuration class for RBLNViTForImageClassification.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Vision Transformer (ViT) models for image classification tasks.
+    """

optimum/rbln/transformers/models/vit/modeling_vit.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...modeling_generic import RBLNModelForImageClassification
+
+
+class RBLNViTForImageClassification(RBLNModelForImageClassification):
+    """
+    RBLN optimized Vision Transformer (ViT) model for image classification tasks.
+
+    This class provides hardware-accelerated inference for Vision Transformer models
+    on RBLN devices, supporting image classification with transformer-based architectures
+    that process images as sequences of patches.
+    """

optimum/rbln/transformers/models/wav2vec2/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .configuration_wav2vec import RBLNWav2Vec2ForCTCConfig
+from .configuration_wav2vec2 import RBLNWav2Vec2ForCTCConfig
 from .modeling_wav2vec2 import RBLNWav2Vec2ForCTC

optimum/rbln/transformers/models/wav2vec2/configuration_wav2vec2.py

@@ -0,0 +1,26 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...configuration_generic import RBLNModelForMaskedLMConfig
+
+
+class RBLNWav2Vec2ForCTCConfig(RBLNModelForMaskedLMConfig):
+    """
+    Configuration class for RBLNWav2Vec2ForCTC.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Wav2Vec2 models for Connectionist Temporal Classification (CTC) tasks.
+    """
+
+    rbln_model_input_names = ["input_values"]

optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py

@@ -17,7 +17,7 @@ import torch
 from transformers import AutoModelForMaskedLM, Wav2Vec2ForCTC
 
 from ...modeling_generic import RBLNModelForMaskedLM
-from .configuration_wav2vec import RBLNWav2Vec2ForCTCConfig
+from .configuration_wav2vec2 import RBLNWav2Vec2ForCTCConfig
 
 
 class _Wav2Vec2(torch.nn.Module):

optimum/rbln/transformers/models/whisper/configuration_whisper.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any, Dict
+
 import rebel
 
 from ....configuration_utils import RBLNModelConfig
@@ -22,6 +24,13 @@ logger = get_logger()
 
 
 class RBLNWhisperForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNWhisperForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Whisper models for speech recognition and transcription tasks.
+    """
+
     def __init__(
         self,
         batch_size: int = None,
@@ -29,7 +38,7 @@ class RBLNWhisperForConditionalGenerationConfig(RBLNModelConfig):
         use_attention_mask: bool = None,
         enc_max_seq_len: int = None,
         dec_max_seq_len: int = None,
-        **kwargs,
+        **kwargs: Dict[str, Any],
     ):
         """
         Args: