




# 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			(							lowercase	:       Union[str, Any]     ):


      '''simple docstring'''



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

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Any     ):


      '''simple docstring'''



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

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       str     ):


      '''simple docstring'''



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

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Union[str, Any]     ):


      '''simple docstring'''



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

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Union[str, Any]     ):


      '''simple docstring'''



      if state.is_main_process:
            lowerCamelCase_		     =    torch.arange(state.num_processes + 1     ).to(state.device     )
      else:
            lowerCamelCase_		     =    torch.arange(state.num_processes     ).to(state.device     )
      lowerCamelCase_		     =    pad_across_processes(lowercase     )
      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			(							lowercase	:       Optional[int]     ):


      '''simple docstring'''



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

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Union[str, Any]     ):


      '''simple docstring'''



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

def 			_SCREAMING_SNAKE_CASE			(							lowercase	:       Optional[int]     ):


      '''simple docstring'''



      main()

def 			_SCREAMING_SNAKE_CASE			(							):


      '''simple docstring'''



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


if __name__ == "__main__":
 main()









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowerCamelCase				:		str     =   {
    "configuration_nezha": ["NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP", "NezhaConfig"],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase				:		List[str]     =   [
     "NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST",
     "NezhaForNextSentencePrediction",
     "NezhaForMaskedLM",
     "NezhaForPreTraining",
     "NezhaForMultipleChoice",
     "NezhaForQuestionAnswering",
     "NezhaForSequenceClassification",
     "NezhaForTokenClassification",
     "NezhaModel",
     "NezhaPreTrainedModel",
 ]


if TYPE_CHECKING:
 from .configuration_nezha import NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP, NezhaConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_nezha import (
      NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST,
      NezhaForMaskedLM,
      NezhaForMultipleChoice,
      NezhaForNextSentencePrediction,
      NezhaForPreTraining,
      NezhaForQuestionAnswering,
      NezhaForSequenceClassification,
      NezhaForTokenClassification,
      NezhaModel,
      NezhaPreTrainedModel,
  )


else:
 import sys

 lowerCamelCase				:		str     =   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



