


'''simple docstring'''





import unittest

from transformers import MraConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
    import torch

    from transformers import (
        MraForMaskedLM,
        MraForMultipleChoice,
        MraForQuestionAnswering,
        MraForSequenceClassification,
        MraForTokenClassification,
        MraModel,
    )
    from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST




class 				UpperCAmelCase			:


 def __init__(      self							:	str			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	List[str]=2			,       __lowerCamelCase							:	Dict=8			,       __lowerCamelCase							:	Union[str, Any]=True			,       __lowerCamelCase							:	Union[str, Any]=True			,       __lowerCamelCase							:	int=True			,       __lowerCamelCase							:	Optional[Any]=True			,       __lowerCamelCase							:	int=9_9			,       __lowerCamelCase							:	Optional[int]=1_6			,       __lowerCamelCase							:	Optional[int]=5			,       __lowerCamelCase							:	str=2			,       __lowerCamelCase							:	Union[str, Any]=3_6			,       __lowerCamelCase							:	Optional[int]="gelu"			,       __lowerCamelCase							:	Any=0.0			,       __lowerCamelCase							:	Dict=0.0			,       __lowerCamelCase							:	Optional[int]=5_1_2			,       __lowerCamelCase							:	Union[str, Any]=1_6			,       __lowerCamelCase							:	Optional[Any]=2			,       __lowerCamelCase							:	Optional[int]=0.02			,       __lowerCamelCase							:	List[Any]=3			,       __lowerCamelCase							:	Dict=4			,       __lowerCamelCase							:	Optional[int]=None			,       ):
    UpperCAmelCase__	:Tuple      =      parent
    UpperCAmelCase__	:Dict      =      batch_size
    UpperCAmelCase__	:str      =      seq_length
    UpperCAmelCase__	:List[Any]      =      is_training
    UpperCAmelCase__	:Tuple      =      use_input_mask
    UpperCAmelCase__	:Union[str, Any]      =      use_token_type_ids
    UpperCAmelCase__	:Optional[Any]      =      use_labels
    UpperCAmelCase__	:Optional[Any]      =      vocab_size
    UpperCAmelCase__	:Dict      =      hidden_size
    UpperCAmelCase__	:Union[str, Any]      =      num_hidden_layers
    UpperCAmelCase__	:Optional[int]      =      num_attention_heads
    UpperCAmelCase__	:Tuple      =      intermediate_size
    UpperCAmelCase__	:List[Any]      =      hidden_act
    UpperCAmelCase__	:Dict      =      hidden_dropout_prob
    UpperCAmelCase__	:int      =      attention_probs_dropout_prob
    UpperCAmelCase__	:Dict      =      max_position_embeddings
    UpperCAmelCase__	:List[Any]      =      type_vocab_size
    UpperCAmelCase__	:Tuple      =      type_sequence_label_size
    UpperCAmelCase__	:Tuple      =      initializer_range
    UpperCAmelCase__	:int      =      num_labels
    UpperCAmelCase__	:Any      =      num_choices
    UpperCAmelCase__	:str      =      scope


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    UpperCAmelCase__	:Any      =      ids_tensor([self.batch_size, self.seq_length]			,       self.vocab_size )

    UpperCAmelCase__	:Tuple      =      None
    if self.use_input_mask:
       UpperCAmelCase__	:List[str]      =      random_attention_mask([self.batch_size, self.seq_length] )

    UpperCAmelCase__	:str      =      None
    if self.use_token_type_ids:
       UpperCAmelCase__	:List[Any]      =      ids_tensor([self.batch_size, self.seq_length]			,       self.type_vocab_size )

    UpperCAmelCase__	:Optional[int]      =      None
    UpperCAmelCase__	:Any      =      None
    UpperCAmelCase__	:Optional[int]      =      None
    if self.use_labels:
       UpperCAmelCase__	:Dict      =      ids_tensor([self.batch_size]			,       self.type_sequence_label_size )
       UpperCAmelCase__	:str      =      ids_tensor([self.batch_size, self.seq_length]			,       self.num_labels )
       UpperCAmelCase__	:List[Any]      =      ids_tensor([self.batch_size]			,       self.num_choices )

    UpperCAmelCase__	:int      =      self.get_config()

    return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    return MraConfig(
        vocab_size=self.vocab_size			,       hidden_size=self.hidden_size			,       num_hidden_layers=self.num_hidden_layers			,       num_attention_heads=self.num_attention_heads			,       intermediate_size=self.intermediate_size			,       hidden_act=self.hidden_act			,       hidden_dropout_prob=self.hidden_dropout_prob			,       attention_probs_dropout_prob=self.attention_probs_dropout_prob			,       max_position_embeddings=self.max_position_embeddings			,       type_vocab_size=self.type_vocab_size			,       is_decoder=__lowerCamelCase			,       initializer_range=self.initializer_range			,       )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__	:Tuple      =      self.get_config()
    UpperCAmelCase__	:Optional[int]      =      3_0_0
    return config


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    (
        (
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						
    )	:Union[str, Any]      =      self.prepare_config_and_inputs()

    UpperCAmelCase__	:List[str]      =      True
    UpperCAmelCase__	:Dict      =      floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
    UpperCAmelCase__	:Optional[Any]      =      ids_tensor([self.batch_size, self.seq_length]			,       vocab_size=2 )

    return (
        config,
        input_ids,
        token_type_ids,
        input_mask,
        sequence_labels,
        token_labels,
        choice_labels,
        encoder_hidden_states,
        encoder_attention_mask,
    )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict			,       __lowerCamelCase							:	Any			,       __lowerCamelCase							:	Any			,       __lowerCamelCase							:	List[Any]			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	Union[str, Any]			,       __lowerCamelCase							:	Union[str, Any] ):
    UpperCAmelCase__	:Dict      =      MraModel(config=__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:Optional[Any]      =      model(__lowerCamelCase			,       attention_mask=__lowerCamelCase			,       token_type_ids=__lowerCamelCase )
    UpperCAmelCase__	:List[str]      =      model(__lowerCamelCase			,       token_type_ids=__lowerCamelCase )
    UpperCAmelCase__	:Optional[int]      =      model(__lowerCamelCase )
    self.parent.assertEqual(result.last_hidden_state.shape			,       (self.batch_size, self.seq_length, self.hidden_size) )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any			,       __lowerCamelCase							:	str			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Any			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	List[Any]			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Optional[int]			,       ):
    UpperCAmelCase__	:Dict      =      True
    UpperCAmelCase__	:int      =      MraModel(__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:Any      =      model(
        __lowerCamelCase			,       attention_mask=__lowerCamelCase			,       token_type_ids=__lowerCamelCase			,       encoder_hidden_states=__lowerCamelCase			,       encoder_attention_mask=__lowerCamelCase			,       )
    UpperCAmelCase__	:Tuple      =      model(
        __lowerCamelCase			,       attention_mask=__lowerCamelCase			,       token_type_ids=__lowerCamelCase			,       encoder_hidden_states=__lowerCamelCase			,       )
    UpperCAmelCase__	:Tuple      =      model(__lowerCamelCase			,       attention_mask=__lowerCamelCase			,       token_type_ids=__lowerCamelCase )
    self.parent.assertEqual(result.last_hidden_state.shape			,       (self.batch_size, self.seq_length, self.hidden_size) )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	Any			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	Optional[int]			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	int ):
    UpperCAmelCase__	:Tuple      =      MraForMaskedLM(config=__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:Dict      =      model(__lowerCamelCase			,       attention_mask=__lowerCamelCase			,       token_type_ids=__lowerCamelCase			,       labels=__lowerCamelCase )
    self.parent.assertEqual(result.logits.shape			,       (self.batch_size, self.seq_length, self.vocab_size) )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Optional[int]			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	int			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	List[Any]			,       __lowerCamelCase							:	int ):
    UpperCAmelCase__	:Union[str, Any]      =      MraForQuestionAnswering(config=__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:Union[str, Any]      =      model(
        __lowerCamelCase			,       attention_mask=__lowerCamelCase			,       token_type_ids=__lowerCamelCase			,       start_positions=__lowerCamelCase			,       end_positions=__lowerCamelCase			,       )
    self.parent.assertEqual(result.start_logits.shape			,       (self.batch_size, self.seq_length) )
    self.parent.assertEqual(result.end_logits.shape			,       (self.batch_size, self.seq_length) )


 def        __SCREAMING_SNAKE_CASE					(      self							:	str			,       __lowerCamelCase							:	int			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	List[Any] ):
    UpperCAmelCase__	:Optional[Any]      =      self.num_labels
    UpperCAmelCase__	:Optional[Any]      =      MraForSequenceClassification(__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:Union[str, Any]      =      model(__lowerCamelCase			,       attention_mask=__lowerCamelCase			,       token_type_ids=__lowerCamelCase			,       labels=__lowerCamelCase )
    self.parent.assertEqual(result.logits.shape			,       (self.batch_size, self.num_labels) )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any]			,       __lowerCamelCase							:	List[str]			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	Any			,       __lowerCamelCase							:	int			,       __lowerCamelCase							:	Union[str, Any]			,       __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	int ):
    UpperCAmelCase__	:List[str]      =      self.num_labels
    UpperCAmelCase__	:List[Any]      =      MraForTokenClassification(config=__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:List[Any]      =      model(__lowerCamelCase			,       attention_mask=__lowerCamelCase			,       token_type_ids=__lowerCamelCase			,       labels=__lowerCamelCase )
    self.parent.assertEqual(result.logits.shape			,       (self.batch_size, self.seq_length, self.num_labels) )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any]			,       __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	Optional[int]			,       __lowerCamelCase							:	Union[str, Any]			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Any			,       __lowerCamelCase							:	List[Any]			,       __lowerCamelCase							:	Optional[Any] ):
    UpperCAmelCase__	:str      =      self.num_choices
    UpperCAmelCase__	:Union[str, Any]      =      MraForMultipleChoice(config=__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:Optional[Any]      =      input_ids.unsqueeze(1 ).expand(-1			,       self.num_choices			,       -1 ).contiguous()
    UpperCAmelCase__	:Any      =      token_type_ids.unsqueeze(1 ).expand(-1			,       self.num_choices			,       -1 ).contiguous()
    UpperCAmelCase__	:Any      =      input_mask.unsqueeze(1 ).expand(-1			,       self.num_choices			,       -1 ).contiguous()
    UpperCAmelCase__	:Union[str, Any]      =      model(
        __lowerCamelCase			,       attention_mask=__lowerCamelCase			,       token_type_ids=__lowerCamelCase			,       labels=__lowerCamelCase			,       )
    self.parent.assertEqual(result.logits.shape			,       (self.batch_size, self.num_choices) )





 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    UpperCAmelCase__	:Any      =      self.prepare_config_and_inputs()
    (
        (
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						
    )	:List[Any]      =      config_and_inputs
    UpperCAmelCase__	:Union[str, Any]      =      {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
    return config, inputs_dict




@require_torch
class 				UpperCAmelCase			(			_snake_case						, unittest.TestCase						):
 UpperCAmelCase				=					(
     (
         MraModel,
         MraForMaskedLM,
         MraForMultipleChoice,
         MraForQuestionAnswering,
         MraForSequenceClassification,
         MraForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 UpperCAmelCase				=					False
 UpperCAmelCase				=					False
 UpperCAmelCase				=					False
 UpperCAmelCase				=					False

 UpperCAmelCase				=					()


 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict ):
    UpperCAmelCase__	:List[Any]      =      MraModelTester(self )
    UpperCAmelCase__	:List[str]      =      ConfigTester(self			,       config_class=__lowerCamelCase			,       hidden_size=3_7 )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    self.config_tester.run_common_tests()


 def        __SCREAMING_SNAKE_CASE					(      self							:	int ):
    UpperCAmelCase__	:Union[str, Any]      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    UpperCAmelCase__	:str      =      self.model_tester.prepare_config_and_inputs()
    for type in ["absolute", "relative_key", "relative_key_query"]:
       UpperCAmelCase__	:List[str]      =      type
       self.model_tester.create_and_check_model(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__	:List[Any]      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	str ):
    UpperCAmelCase__	:int      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_multiple_choice(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__	:int      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_question_answering(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    UpperCAmelCase__	:Union[str, Any]      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_sequence_classification(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	str ):
    UpperCAmelCase__	:List[str]      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase )


 @slow
 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       UpperCAmelCase__	:Any      =      MraModel.from_pretrained(__lowerCamelCase )
       self.assertIsNotNone(__lowerCamelCase )





 @unittest.skip(reason='''MRA does not output attentions''' )
 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    return




@require_torch
class 				UpperCAmelCase			(			unittest.TestCase						):


 @slow
 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict ):
    UpperCAmelCase__	:Optional[Any]      =      MraModel.from_pretrained('''uw-madison/mra-base-512-4''' )
    UpperCAmelCase__	:List[str]      =      torch.arange(2_5_6 ).unsqueeze(0 )

    with torch.no_grad():
       UpperCAmelCase__	:Optional[Any]      =      model(__lowerCamelCase )[0]

    UpperCAmelCase__	:Any      =      torch.Size((1, 2_5_6, 7_6_8) )
    self.assertEqual(output.shape			,       __lowerCamelCase )

    UpperCAmelCase__	:Optional[Any]      =      torch.tensor(
        [[[-0.01_40, 0.08_30, -0.03_81], [0.15_46, 0.14_02, 0.02_20], [0.11_62, 0.08_51, 0.01_65]]] )

    self.assertTrue(torch.allclose(output[:, :3, :3]			,       __lowerCamelCase			,       atol=1e-4 ) )


 @slow
 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any] ):
    UpperCAmelCase__	:Tuple      =      MraForMaskedLM.from_pretrained('''uw-madison/mra-base-512-4''' )
    UpperCAmelCase__	:Any      =      torch.arange(2_5_6 ).unsqueeze(0 )

    with torch.no_grad():
       UpperCAmelCase__	:Optional[int]      =      model(__lowerCamelCase )[0]

    UpperCAmelCase__	:int      =      5_0_2_6_5

    UpperCAmelCase__	:Union[str, Any]      =      torch.Size((1, 2_5_6, vocab_size) )
    self.assertEqual(output.shape			,       __lowerCamelCase )

    UpperCAmelCase__	:Optional[Any]      =      torch.tensor(
        [[[9.25_95, -3.60_38, 11.88_19], [9.38_69, -3.26_93, 11.09_56], [11.85_24, -3.49_38, 13.12_10]]] )

    self.assertTrue(torch.allclose(output[:, :3, :3]			,       __lowerCamelCase			,       atol=1e-4 ) )





 @slow
 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    UpperCAmelCase__	:Union[str, Any]      =      MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3''' )
    UpperCAmelCase__	:Any      =      torch.arange(4_0_9_6 ).unsqueeze(0 )

    with torch.no_grad():
       UpperCAmelCase__	:Optional[Any]      =      model(__lowerCamelCase )[0]

    UpperCAmelCase__	:List[str]      =      5_0_2_6_5

    UpperCAmelCase__	:Union[str, Any]      =      torch.Size((1, 4_0_9_6, vocab_size) )
    self.assertEqual(output.shape			,       __lowerCamelCase )

    UpperCAmelCase__	:Optional[int]      =      torch.tensor(
        [[[5.47_89, -2.35_64, 7.50_64], [7.90_67, -1.33_69, 9.96_68], [9.07_12, -1.81_06, 7.03_80]]] )

    self.assertTrue(torch.allclose(output[:, :3, :3]			,       __lowerCamelCase			,       atol=1e-4 ) )







'''simple docstring'''





import importlib.util
import json
import os
import warnings
from dataclasses import dataclass, field

import torch

from ..training_args import TrainingArguments
from ..utils import cached_property, is_sagemaker_dp_enabled, logging


__lowerCamelCase												=						logging.get_logger(__name__)


def        a__     (					):
   # Get the sagemaker specific mp parameters from smp_options variable.
   UpperCAmelCase__	:Any      =      os.getenv('''SM_HP_MP_PARAMETERS''',   '''{}'''						)
   try:
      # Parse it and check the field "partitions" is included, it is required for model parallel.
      UpperCAmelCase__	:Any      =      json.loads(UpperCamelCase_						)
      if "partitions" not in smp_options:
         return False
   except json.JSONDecodeError:
      return False

   # Get the sagemaker specific framework parameters from mpi_options variable.
   UpperCAmelCase__	:Optional[Any]      =      os.getenv('''SM_FRAMEWORK_PARAMS''',   '''{}'''						)
   try:
      # Parse it and check the field "sagemaker_distributed_dataparallel_enabled".
      UpperCAmelCase__	:str      =      json.loads(UpperCamelCase_						)
      if not mpi_options.get('''sagemaker_mpi_enabled''',   UpperCamelCase_						):
         return False
   except json.JSONDecodeError:
      return False
   # Lastly, check if the `smdistributed` module is present.
   return importlib.util.find_spec('''smdistributed'''						) is not None


if is_sagemaker_model_parallel_available():
    import smdistributed.modelparallel.torch as smp

    smp.init()




@dataclass
class 				UpperCAmelCase			(			_snake_case						):
 UpperCAmelCase				=					field(
     default=""						, metadata={"help": "Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer"}						, )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    super().__post_init__()
    warnings.warn(
        '''`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use '''
        '''`TrainingArguments` instead.'''			,       __lowerCamelCase			,       )


 @cached_property
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    logger.info('''PyTorch: setting up devices''' )
    if torch.distributed.is_available() and torch.distributed.is_initialized() and self.local_rank == -1:
       logger.warning(
           '''torch.distributed process group is initialized, but local_rank == -1. '''
           '''In order to use Torch DDP, launch your script with `python -m torch.distributed.launch''' )
    if self.no_cuda:
       UpperCAmelCase__	:Union[str, Any]      =      torch.device('''cpu''' )
       UpperCAmelCase__	:Dict      =      0
    elif is_sagemaker_model_parallel_available():
       UpperCAmelCase__	:Dict      =      smp.local_rank()
       UpperCAmelCase__	:List[str]      =      torch.device('''cuda'''			,       __lowerCamelCase )
       UpperCAmelCase__	:Optional[int]      =      1
    elif is_sagemaker_dp_enabled():
       import smdistributed.dataparallel.torch.torch_smddp  # noqa: F401

       torch.distributed.init_process_group(backend='''smddp'''			,       timeout=self.ddp_timeout_delta )
       UpperCAmelCase__	:int      =      int(os.getenv('''SMDATAPARALLEL_LOCAL_RANK''' ) )
       UpperCAmelCase__	:int      =      torch.device('''cuda'''			,       self.local_rank )
       UpperCAmelCase__	:Union[str, Any]      =      1
    elif self.local_rank == -1:
       # if n_gpu is > 1 we'll use nn.DataParallel.
       # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
       # Explicitly set CUDA to the first (index 0) CUDA device, otherwise `set_device` will
       # trigger an error that a device index is missing. Index 0 takes into account the
       # GPUs available in the environment, so `CUDA_VISIBLE_DEVICES=1,2` with `cuda:0`
       # will use the first GPU in that env, i.e. GPU#1
       UpperCAmelCase__	:str      =      torch.device('''cuda:0''' if torch.cuda.is_available() else '''cpu''' )
       # Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at
       # the default value.
       UpperCAmelCase__	:str      =      torch.cuda.device_count()
    else:
       # Here, we'll use torch.distributed.
       # Initializes the distributed backend which will take care of synchronizing nodes/GPUs
       if not torch.distributed.is_initialized():
          torch.distributed.init_process_group(backend='''nccl'''			,       timeout=self.ddp_timeout_delta )
       UpperCAmelCase__	:str      =      torch.device('''cuda'''			,       self.local_rank )
       UpperCAmelCase__	:Union[str, Any]      =      1

    if device.type == "cuda":
       torch.cuda.set_device(__lowerCamelCase )

    return device


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    if is_sagemaker_model_parallel_available():
       return smp.dp_size()

    return super().world_size


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict ):
    return not is_sagemaker_model_parallel_available()





 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	int ):
    return False



