optimum-rbln 0.8.1a6__py3-none-any.whl → 0.8.1a7__py3-none-any.whl

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -43,6 +43,13 @@ class _TextEncoder(torch.nn.Module):
 
 
 class RBLNCLIPTextModel(RBLNModel):
+    """
+    RBLN optimized CLIP text encoder model.
+
+    This class provides hardware-accelerated inference for CLIP text encoders
+    on RBLN devices, supporting text encoding for multimodal tasks.
+    """
+
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPTextModelConfig) -> torch.nn.Module:
         return _TextEncoder(model).eval()
@@ -95,7 +102,12 @@ class RBLNCLIPTextModel(RBLNModel):
 
 
 class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
-    pass
+    """
+    RBLN optimized CLIP text encoder model with projection layer.
+
+    This class extends RBLNCLIPTextModel with a projection layer for
+    multimodal embedding alignment tasks.
+    """
 
 
 class _VisionEncoder(torch.nn.Module):
@@ -109,6 +121,13 @@ class _VisionEncoder(torch.nn.Module):
 
 
 class RBLNCLIPVisionModel(RBLNModel):
+    """
+    RBLN optimized CLIP vision encoder model.
+
+    This class provides hardware-accelerated inference for CLIP vision encoders
+    on RBLN devices, supporting image encoding for multimodal tasks.
+    """
+
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNCLIPVisionModelConfig) -> torch.nn.Module:
         return _VisionEncoder(model).eval()
@@ -182,6 +201,13 @@ class RBLNCLIPVisionModel(RBLNModel):
 
 
 class RBLNCLIPVisionModelWithProjection(RBLNCLIPVisionModel):
+    """
+    RBLN optimized CLIP vision encoder model with projection layer.
+
+    This class extends RBLNCLIPVisionModel with a projection layer for
+    multimodal embedding alignment tasks.
+    """
+
     def forward(
         self,
         pixel_values: Optional[torch.FloatTensor] = None,

optimum/rbln/transformers/models/decoderonly/modeling_decoderonly.py

@@ -78,7 +78,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
                 torch.ones(1, 1, self.rbln_config.prefill_chunk_size, self.rbln_config.prefill_chunk_size), diagonal=1
             )
 
-    def get_block_tables(self, cache_position: torch.Tensor, batch_idx: int = None):
+    def get_block_tables(self, cache_position: torch.Tensor, batch_idx: int = None) -> torch.Tensor:
         """
         Manages and returns the KV cache block tables.
         Updates the block tables based on the given cache_position, allocating new blocks or reusing existing ones as needed.
@@ -88,7 +88,7 @@ class RBLNRuntimeModel(RBLNPytorchRuntime):
             batch_idx (int, optional): Specific batch index, used when phase is 'prefill'.
 
         Returns:
-            torch.Tensor: Updated block tables.
+            Updated block tables.
         """
 
         NO_BLOCKS_ERROR = (
@@ -458,6 +458,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     This class serves as the foundation for various decoder-only architectures like GPT, LLaMA, etc.
 
     The class provides core functionality for:
+
     1. Converting pre-trained transformer models to RBLN-optimized format
     2. Handling the compilation process for RBLN devices
     3. Managing inference operations for causal language modeling
@@ -532,7 +533,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     @classmethod
     def save_torch_artifacts(
         cls,
-        model: "PreTrainedModel",
+        model: PreTrainedModel,
         save_dir_path: Path,
         subfolder: str,
         rbln_config: RBLNDecoderOnlyModelForCausalLMConfig,
@@ -566,7 +567,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     def get_quantized_model(
         cls,
         model_id: str,
-        config: Optional["PretrainedConfig"] = None,
+        config: Optional[PretrainedConfig] = None,
         use_auth_token: Optional[Union[bool, str]] = None,
         revision: Optional[str] = None,
         force_download: bool = False,
@@ -605,16 +606,15 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         return model
 
     def __getattr__(self, __name: str) -> Any:
-        """
-        Special method to delegate attribute access to the original Huggingface LM class.
-        This method is called when an attribute is not found in the current instance's dictionary.
-        It enables transparent access to the original model's attributes and methods while maintaining
-        proper method binding.
-
-        The method implements a delegation pattern that:
-        1. For methods: Creates a wrapper that properly binds 'self' to method calls
-        2. For other attributes: Returns them directly from the original class
-        """
+        # Special method to delegate attribute access to the original Huggingface LM class.
+        # This method is called when an attribute is not found in the current instance's dictionary.
+        # It enables transparent access to the original model's attributes and methods while maintaining
+        # proper method binding.
+
+        # The method implements a delegation pattern that:
+
+        # 1. For methods: Creates a wrapper that properly binds 'self' to method calls
+        # 2. For other attributes: Returns them directly from the original class
 
         def redirect(func):
             return lambda *pargs, **kwargs: func(self, *pargs, **kwargs)
@@ -627,7 +627,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     @classmethod
     def get_pytorch_model(
         cls, *args, rbln_config: Optional[RBLNDecoderOnlyModelForCausalLMConfig] = None, **kwargs
-    ) -> "PreTrainedModel":
+    ) -> PreTrainedModel:
         if rbln_config and rbln_config.quantization:
             model = cls.get_quantized_model(*args, **kwargs)
         else:
@@ -636,7 +636,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
         return model
 
     @classmethod
-    def wrap_model_if_needed(cls, model: "PreTrainedModel", rbln_config: "RBLNDecoderOnlyModelForCausalLMConfig"):
+    def wrap_model_if_needed(cls, model: PreTrainedModel, rbln_config: "RBLNDecoderOnlyModelForCausalLMConfig"):
         wrapper_cfg = {
             "max_seq_len": rbln_config.max_seq_len,
             "attn_impl": rbln_config.attn_impl,
@@ -654,7 +654,7 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
 
     @classmethod
     @torch.inference_mode()
-    def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
+    def get_compiled_model(cls, model: PreTrainedModel, rbln_config: RBLNDecoderOnlyModelForCausalLMConfig):
         wrapped_model = cls.wrap_model_if_needed(model, rbln_config)
 
         rbln_compile_configs = rbln_config.compile_cfgs
@@ -679,9 +679,11 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
                     quantization.maybe_set_quantization_env()
                 original_linear = torch.nn.functional.linear
                 torch.nn.functional.linear = torch.ops.rbln_custom_ops.linear
-                compiled_model = RBLNModel.compile(
+                compiled_model = cls.compile(
                     wrapped_model,
                     compile_config,
+                    create_runtimes=rbln_config.create_runtimes,
+                    device=rbln_config.device,
                     example_inputs=example_inputs,
                     compile_context=compile_context,
                 )
@@ -973,8 +975,8 @@ class RBLNDecoderOnlyModelForCausalLM(RBLNModel):
     def _update_rbln_config(
         cls,
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]] = None,
-        model: Optional["PreTrainedModel"] = None,
-        model_config: Optional["PretrainedConfig"] = None,
+        model: Optional[PreTrainedModel] = None,
+        model_config: Optional[PretrainedConfig] = None,
         rbln_config: Optional[RBLNDecoderOnlyModelForCausalLMConfig] = None,
     ) -> RBLNDecoderOnlyModelForCausalLMConfig:
         if rbln_config.max_seq_len is None:

optimum/rbln/transformers/models/distilbert/configuration_distilbert.py

@@ -16,4 +16,9 @@ from ...configuration_generic import RBLNModelForQuestionAnsweringConfig
 
 
 class RBLNDistilBertForQuestionAnsweringConfig(RBLNModelForQuestionAnsweringConfig):
-    ""
+    """
+    Configuration class for RBLNDistilBertForQuestionAnswering.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized DistilBERT models for question answering tasks.
+    """

optimum/rbln/transformers/models/distilbert/modeling_distilbert.py

@@ -16,4 +16,12 @@ from ...modeling_generic import RBLNModelForQuestionAnswering
 
 
 class RBLNDistilBertForQuestionAnswering(RBLNModelForQuestionAnswering):
+    """
+    RBLN optimized DistilBERT model for question answering tasks.
+
+    This class provides hardware-accelerated inference for DistilBERT models
+    on RBLN devices, supporting extractive question answering tasks where
+    the model predicts start and end positions of answers in text.
+    """
+
     rbln_model_input_names = ["input_ids", "attention_mask"]

optimum/rbln/transformers/models/dpt/configuration_dpt.py

@@ -16,4 +16,9 @@ from ...configuration_generic import RBLNModelForDepthEstimationConfig
 
 
 class RBLNDPTForDepthEstimationConfig(RBLNModelForDepthEstimationConfig):
-    pass
+    """
+    Configuration class for RBLNDPTForDepthEstimation.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized DPT (Dense Prediction Transformer) models for depth estimation tasks.
+    """

optimum/rbln/transformers/models/dpt/modeling_dpt.py

@@ -17,4 +17,9 @@ from ...modeling_generic import RBLNModelForDepthEstimation
 
 
 class RBLNDPTForDepthEstimation(RBLNModelForDepthEstimation):
-    pass
+    """
+    RBLN optimized DPT model for depth estimation tasks.
+
+    This class provides hardware-accelerated inference for DPT (Dense Prediction Transformer)
+    models on RBLN devices, supporting monocular depth estimation from single images.
+    """

optimum/rbln/transformers/models/gemma3/modeling_gemma3.py

@@ -326,7 +326,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
         attention_mask: torch.Tensor,
         position_ids: torch.Tensor,
         token_type_ids: Optional[torch.Tensor] = None,
-    ):
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, torch.Tensor]:
         """
         Pads inputs, attention_mask, and position_ids so image token groups (256 tokens with token_type_ids == 1)
         start at multiples of prefill_chunk_size (256). Returns padded tensors and total padded length.
@@ -338,7 +338,7 @@ class RBLNGemma3RuntimeModel(RBLNRuntimeModel):
             token_type_ids: (1, seq_len) tensor, 0 for text, 1 for image.
 
         Returns:
-            Tuple: (inputs_padded, attention_mask_padded, position_ids_padded, padded_len, token_type_ids_padded).
+            (inputs_padded, attention_mask_padded, position_ids_padded, padded_len, token_type_ids_padded).
         """
 
         if token_type_ids is None:
@@ -816,9 +816,11 @@ class RBLNGemma3ForCausalLM(RBLNDecoderOnlyModelForCausalLM):
                     quantization.maybe_set_quantization_env()
                 original_linear = torch.nn.functional.linear
                 torch.nn.functional.linear = torch.ops.rbln_custom_ops.linear
-                compiled_model = RBLNModel.compile(
+                compiled_model = cls.compile(
                     wrapped_model,
                     compile_config,
+                    create_runtimes=rbln_config.create_runtimes,
+                    device=rbln_config.device,
                     example_inputs=example_inputs,
                     compile_context=compile_context,
                 )

optimum/rbln/transformers/models/llava_next/configuration_llava_next.py

@@ -18,6 +18,14 @@ from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNLlavaNextForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNLlavaNextForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized LLaVA-Next models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
     submodules = ["vision_tower", "language_model"]
 
     def __init__(

optimum/rbln/transformers/models/qwen2_5_vl/configuration_qwen2_5_vl.py

@@ -19,6 +19,14 @@ from ..decoderonly.configuration_decoderonly import RBLNDecoderOnlyModelForCausa
 
 
 class RBLNQwen2_5_VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausalLMConfig):
+    """
+    Configuration class for RBLNQwen2_5_VLForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Qwen2.5-VL models for multimodal conditional generation tasks
+    that combine vision and language processing capabilities.
+    """
+
     submodules = ["visual"]
 
     def __init__(
@@ -37,6 +45,14 @@ class RBLNQwen2_5_VLForConditionalGenerationConfig(RBLNDecoderOnlyModelForCausal
 
 
 class RBLNQwen2_5_VisionTransformerPretrainedModelConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNQwen2_5_VisionTransformerPretrainedModel.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized Qwen2.5-VL vision transformer models with window-based attention
+    mechanisms for processing images and videos.
+    """
+
     def __init__(self, max_seq_lens: Union[int, List[int]] = None, **kwargs: Dict[str, Any]):
         """
         Args:

optimum/rbln/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -54,6 +54,14 @@ if TYPE_CHECKING:
 
 
 class RBLNQwen2_5_VisionTransformerPretrainedModel(RBLNModel):
+    """
+    RBLN optimized Qwen2.5-VL vision transformer model.
+
+    This class provides hardware-accelerated inference for Qwen2.5-VL vision transformers
+    on RBLN devices, supporting image and video encoding for multimodal vision-language tasks
+    with window-based attention mechanisms.
+    """
+
     auto_model_class = None
 
     def __post_init__(self, **kwargs):

optimum/rbln/transformers/models/resnet/configuration_resnet.py

@@ -17,4 +17,9 @@ from ...configuration_generic import RBLNModelForImageClassificationConfig
 
 
 class RBLNResNetForImageClassificationConfig(RBLNModelForImageClassificationConfig):
-    ""
+    """
+    Configuration class for RBLNResNetForImageClassification.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized ResNet models for image classification tasks.
+    """

optimum/rbln/transformers/models/resnet/modeling_resnet.py

@@ -18,5 +18,9 @@ from ...modeling_generic import RBLNModelForImageClassification
 
 class RBLNResNetForImageClassification(RBLNModelForImageClassification):
     """
-    ResNet model for image classification tasks on RBLN NPU.
+    RBLN optimized ResNet model for image classification tasks.
+
+    This class provides hardware-accelerated inference for ResNet models
+    on RBLN devices, supporting image classification with convolutional neural networks
+    designed for computer vision tasks.
     """

optimum/rbln/transformers/models/roberta/configuration_roberta.py

@@ -16,8 +16,18 @@ from ...configuration_generic import RBLNModelForMaskedLMConfig, RBLNModelForSeq
 
 
 class RBLNRobertaForMaskedLMConfig(RBLNModelForMaskedLMConfig):
-    ""
+    """
+    Configuration class for RBLNRobertaForMaskedLM.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized RoBERTa models for masked language modeling tasks.
+    """
 
 
 class RBLNRobertaForSequenceClassificationConfig(RBLNModelForSequenceClassificationConfig):
-    ""
+    """
+    Configuration class for RBLNRobertaForSequenceClassification.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized RoBERTa models for sequence classification tasks.
+    """

optimum/rbln/transformers/models/roberta/modeling_roberta.py

@@ -16,8 +16,24 @@ from ...modeling_generic import RBLNModelForMaskedLM, RBLNModelForSequenceClassi
 
 
 class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
+    """
+    RBLN optimized RoBERTa model for masked language modeling tasks.
+
+    This class provides hardware-accelerated inference for RoBERTa models
+    on RBLN devices, supporting masked language modeling tasks such as
+    token prediction and text completion.
+    """
+
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
 class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
+    """
+    RBLN optimized RoBERTa model for sequence classification tasks.
+
+    This class provides hardware-accelerated inference for RoBERTa models
+    on RBLN devices, supporting text classification tasks such as sentiment analysis,
+    topic classification, and other sequence-level prediction tasks.
+    """
+
     rbln_model_input_names = ["input_ids", "attention_mask"]

optimum/rbln/transformers/models/seq2seq/modeling_seq2seq.py

@@ -161,16 +161,20 @@ class RBLNModelForSeq2SeqLM(RBLNModel, ABC):
             if "key_value_states" in name:
                 context.mark_static_address(tensor)
 
-        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,
         )
 
-        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,
         )

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

@@ -18,6 +18,13 @@ 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,

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

@@ -58,6 +58,13 @@ class _SiglipVisionModel(torch.nn.Module):
 
 
 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 = {

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/time_series_transformer/configuration_time_series_transformer.py

@@ -4,6 +4,13 @@ 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,

optimum/rbln/transformers/models/time_series_transformer/modeling_time_series_transformer.py

@@ -194,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,
         )

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

@@ -16,4 +16,9 @@ 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

@@ -16,4 +16,10 @@ 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/configuration_wav2vec2.py

@@ -16,4 +16,11 @@ 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/whisper/configuration_whisper.py

@@ -24,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,

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

@@ -230,15 +230,19 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
             if "key_value_states" in name:
                 context.mark_static_address(tensor)
 
-        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,
         )
-        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,
         )

optimum/rbln/utils/runtime_utils.py

@@ -13,12 +13,57 @@
 # limitations under the License.
 
 import threading
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional, Union
 
 import rebel
 import torch
 
 
+def tp_and_devices_are_ok(
+    tensor_parallel_size: Optional[int] = None,
+    device: Optional[Union[int, List[int]]] = None,
+    npu: Optional[str] = None,
+) -> Optional[str]:
+    if tensor_parallel_size is None:
+        tensor_parallel_size = 1
+
+    if rebel.device_count() < tensor_parallel_size:
+        return (
+            f"Tensor parallel size {tensor_parallel_size} is greater than "
+            f"the number of available devices {rebel.device_count()}."
+        )
+
+    if device is None:
+        device = list(range(tensor_parallel_size))
+    elif isinstance(device, int):
+        device = [device]
+    elif isinstance(device, list):
+        if any(not isinstance(d, int) for d in device):
+            return "Device must be a(n) (list of) integer(s)."
+        if len(device) != tensor_parallel_size:
+            return (
+                f"The number of devices ({len(device)}) does not match tensor parallel size ({tensor_parallel_size})."
+            )
+    else:
+        return f"Invalid device: {device}"
+
+    for device_id in device:
+        if device_id < 0:  # if any device is dummy device, skip it
+            return None
+        if rebel.get_npu_name(device_id) is None:
+            return (
+                f"Device {device_id} is not a valid NPU device. Please check your NPU status with 'rbln-stat' command."
+            )
+
+    if npu is not None:
+        for device_id in device:
+            npu_name = rebel.get_npu_name(device_id)
+            if npu_name != npu:
+                return f"Device {device_id} ({npu_name}) is not on the same NPU as {npu}."
+
+    return None
+
+
 class RBLNPytorchRuntime:
     mandatory_members = []
 

{optimum_rbln-0.8.1a6.dist-info → optimum_rbln-0.8.1a7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.1a6
+Version: 0.8.1a7
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai