# Copyright 2023 The HuggingFace Team. 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.

import torch

from accelerate import PartialState
from accelerate.utils.operations import broadcast, gather, gather_object, pad_across_processes, reduce



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Union[str, Any]     )							-> Tuple:







  """simple docstring"""





  return (torch.arange(state.num_processes     ) + 1.0 + (state.num_processes * state.process_index)).to(state.device     )



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Any     )							-> Optional[int]:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				create_tensor(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				gather(__UpperCamelCase     )
  assert gathered_tensor.tolist() == list(range(1			,					state.num_processes**2 + 1     )     )



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Optional[Any]     )							-> Tuple:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				[state.process_index]
  SCREAMING_SNAKE_CASE__      							=				gather_object(__UpperCamelCase     )
  assert len(__UpperCamelCase     ) == state.num_processes, f"""{gathered_obj}, {len(__UpperCamelCase     )} != {state.num_processes}"""
  assert gathered_obj == list(range(state.num_processes     )     ), f"""{gathered_obj} != {list(range(state.num_processes     )     )}"""



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Optional[Any]     )							-> str:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				create_tensor(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				broadcast(__UpperCamelCase     )
  assert broadcasted_tensor.shape == torch.Size([state.num_processes]     )
  assert broadcasted_tensor.tolist() == list(range(1			,					state.num_processes + 1     )     )



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	List[str]     )							-> Dict:







  """simple docstring"""





  if state.is_main_process:
    SCREAMING_SNAKE_CASE__      							=				torch.arange(state.num_processes + 1     ).to(state.device     )
  else:
    SCREAMING_SNAKE_CASE__      							=				torch.arange(state.num_processes     ).to(state.device     )
  SCREAMING_SNAKE_CASE__      							=				pad_across_processes(__UpperCamelCase     )
  assert padded_tensor.shape == torch.Size([state.num_processes + 1]     )
  if not state.is_main_process:
    assert padded_tensor.tolist() == list(range(0			,					state.num_processes     )     ) + [0]



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	List[str]     )							-> Any:







  """simple docstring"""





  if state.num_processes != 2:
    return
  SCREAMING_SNAKE_CASE__      							=				create_tensor(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				reduce(__UpperCamelCase			,					"""sum"""     )
  SCREAMING_SNAKE_CASE__      							=				torch.tensor([4.0, 6]     ).to(state.device     )
  assert torch.allclose(__UpperCamelCase			,					__UpperCamelCase     ), f"""{reduced_tensor} != {truth_tensor}"""



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	str     )							-> Dict:







  """simple docstring"""





  if state.num_processes != 2:
    return
  SCREAMING_SNAKE_CASE__      							=				create_tensor(__UpperCamelCase     )
  SCREAMING_SNAKE_CASE__      							=				reduce(__UpperCamelCase			,					"""mean"""     )
  SCREAMING_SNAKE_CASE__      							=				torch.tensor([2.0, 3]     ).to(state.device     )
  assert torch.allclose(__UpperCamelCase			,					__UpperCamelCase     ), f"""{reduced_tensor} != {truth_tensor}"""



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Tuple     )							-> str:







  """simple docstring"""





  main()



def      __SCREAMING_SNAKE_CASE  ( )							-> Tuple:







  """simple docstring"""





  SCREAMING_SNAKE_CASE__      							=				PartialState()
  state.print(f"""State: {state}"""     )
  state.print("""testing gather"""     )
  test_gather(__UpperCamelCase     )
  state.print("""testing gather_object"""     )
  test_gather_object(__UpperCamelCase     )
  state.print("""testing broadcast"""     )
  test_broadcast(__UpperCamelCase     )
  state.print("""testing pad_across_processes"""     )
  test_pad_across_processes(__UpperCamelCase     )
  state.print("""testing reduce_sum"""     )
  test_reduce_sum(__UpperCamelCase     )
  state.print("""testing reduce_mean"""     )
  test_reduce_mean(__UpperCamelCase     )


if __name__ == "__main__":
   main()

import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
   from transformers import AutoModelForSeqaSeqLM, AutoTokenizer







@require_torch
@require_sentencepiece
@require_tokenizers
class __snake_case    (						unittest.TestCase		):




   @slow
   def      __a						(   self      :    List[str]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				AutoModelForSeqaSeqLM.from_pretrained("""google/mt5-small"""	,     return_dict=_lowercase				).to(_lowercase				)
     SCREAMING_SNAKE_CASE__      							=				AutoTokenizer.from_pretrained("""google/mt5-small"""				)

     SCREAMING_SNAKE_CASE__      							=				tokenizer("""Hello there"""	,     return_tensors="""pt"""				).input_ids
     SCREAMING_SNAKE_CASE__      							=				tokenizer("""Hi I am"""	,     return_tensors="""pt"""				).input_ids

     SCREAMING_SNAKE_CASE__      							=				model(input_ids.to(_lowercase				)	,     labels=labels.to(_lowercase				)				).loss
     SCREAMING_SNAKE_CASE__      							=				-(labels.shape[-1] * loss.item())

     SCREAMING_SNAKE_CASE__      							=				-84.91_27
     self.assertTrue(abs(mtf_score - EXPECTED_SCORE				) < 1E-4				)
