



# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.utils import logging


a__      =		logging.get_logger(__name__)


if is_torch_available():
  import torch
  from torch import nn
  from torch.utils.data import Dataset

  from transformers import Trainer







  class 						snake_case  (		SCREAMING_SNAKE_CASE_					):



     '''simple docstring'''






     def __init__(   self							:     Dict			,						lowerCAmelCase							:     int = 101)			->			Optional[Any]:

      """simple docstring"""







      _snake_case      :					Optional[Any]					=   length

     def __len__(   self							:     Optional[Any])			->			Optional[int]:

      """simple docstring"""







      return self.length


     def __getitem__(   self							:     Dict			,						lowerCAmelCase							:     Any)			->			int:

      """simple docstring"""







      return i







  class 						snake_case  :



     '''simple docstring'''






     def __call__(   self							:     Optional[int]			,						lowerCAmelCase							:     List[str])			->			Optional[int]:

      """simple docstring"""







      return {"input_ids": torch.tensor(lowerCAmelCase), "labels": torch.tensor(lowerCAmelCase)}







  class 						snake_case  (		nn.Module					):



     '''simple docstring'''






     def __init__(   self							:     Dict)			->			Union[str, Any]:

      """simple docstring"""







      super().__init__()
      # Add some (unused) params otherwise DDP will complain.
      _snake_case      :					Optional[int]					=   nn.Linear(120			,						80)


     def        UpperCamelCase_					(   self							:     Union[str, Any]			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     Any=None)			->			int:

      """simple docstring"""







      if labels is not None:
       return torch.tensor(0.0			,						device=input_ids.device), input_ids
      else:
       return input_ids



def  lowercase    (				)					->       Optional[Any]:
 _snake_case      :					Optional[int]					=   HfArgumentParser((TrainingArguments,)			)
 sys.argv += ["--output_dir", "./examples"]
 _snake_case      :					Optional[int]					=   parser.parse_args_into_dataclasses()[0]

 logger.warning(
     F'''Process rank: {training_args.local_rank}, device: {training_args.device}, '''
     F'''tpu_num_cores: {training_args.tpu_num_cores}'''			,    )

 # Essentially, what we want to verify in the distributed case is
 # that we get all samples back, in the right order.
 # (this is crucial for prediction for instance)
 for dataset_length in [1_001, 256, 15]:
  _snake_case      :					Optional[int]					=   DummyDataset(SCREAMING_SNAKE_CASE__			)

  def compute_metrics(SCREAMING_SNAKE_CASE__    :		EvalPrediction			) -> Dict:
   _snake_case      :					List[str]					=   list(range(len(SCREAMING_SNAKE_CASE__			)			)			)
   _snake_case      :					Optional[Any]					=   p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
   return {"success": success}

  _snake_case      :					Optional[Any]					=   Trainer(
      model=DummyModel()			,    args=SCREAMING_SNAKE_CASE__			,    data_collator=DummyDataCollator()			,    eval_dataset=SCREAMING_SNAKE_CASE__			,    compute_metrics=SCREAMING_SNAKE_CASE__			,    )
  _snake_case      :					Any					=   trainer.evaluate()
  logger.info(SCREAMING_SNAKE_CASE__			)
  if metrics["eval_success"] is not True:
   logger.error(SCREAMING_SNAKE_CASE__			)
   exit(1			)

  _snake_case      :					Union[str, Any]					=   trainer.predict(SCREAMING_SNAKE_CASE__			)
  logger.info(p.metrics			)
  if p.metrics["test_success"] is not True:
   logger.error(p.metrics			)
   exit(1			)

  _snake_case      :					str					=   2

  _snake_case      :					Dict					=   trainer.evaluate()
  logger.info(SCREAMING_SNAKE_CASE__			)
  if metrics["eval_success"] is not True:
   logger.error(SCREAMING_SNAKE_CASE__			)
   exit(1			)

  _snake_case      :					Tuple					=   trainer.predict(SCREAMING_SNAKE_CASE__			)
  logger.info(p.metrics			)
  if p.metrics["test_success"] is not True:
   logger.error(p.metrics			)
   exit(1			)

  _snake_case      :					Any					=   None

 logger.info("""🔥 All distributed tests successful"""			)



def  lowercase    (				SCREAMING_SNAKE_CASE__    :		int			)					->       List[str]:
 # For xla_spawn (TPUs)
 main()


if __name__ == "__main__":
  main()





def  lowercase    (				SCREAMING_SNAKE_CASE__    :		float			,    SCREAMING_SNAKE_CASE__    :		float			)					->       float:
 return a + b


if __name__ == "__main__":
  a__      =		5
  a__      =		6
  print(F'''The sum of {a} + {b} is {add(a, b)}''')
