

# 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        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     return (torch.arange(state.num_processes       ) + 1.0 + (state.num_processes * state.process_index)).to(state.device       )
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					create_tensor(_A       )
     SCREAMING_SNAKE_CASE__					=					gather(_A       )
     assert gathered_tensor.tolist() == list(range(1   ,	state.num_processes**2 + 1       )       )
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					[state.process_index]
     SCREAMING_SNAKE_CASE__					=					gather_object(_A       )
     assert len(_A       ) == state.num_processes, F'''{gathered_obj}, {len(_A       )} != {state.num_processes}'''
     assert gathered_obj == list(range(state.num_processes       )       ), F'''{gathered_obj} != {list(range(state.num_processes       )       )}'''
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




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





     '''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(_A       )
     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        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     if state.num_processes != 2:
          return
     SCREAMING_SNAKE_CASE__					=					create_tensor(_A       )
     SCREAMING_SNAKE_CASE__					=					reduce(_A   ,	'''sum'''       )
     SCREAMING_SNAKE_CASE__					=					torch.tensor([4.0, 6]       ).to(state.device       )
     assert torch.allclose(_A   ,	_A       ), F'''{reduced_tensor} != {truth_tensor}'''
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     if state.num_processes != 2:
          return
     SCREAMING_SNAKE_CASE__					=					create_tensor(_A       )
     SCREAMING_SNAKE_CASE__					=					reduce(_A   ,	'''mean'''       )
     SCREAMING_SNAKE_CASE__					=					torch.tensor([2.0, 3]       ).to(state.device       )
     assert torch.allclose(_A   ,	_A       ), F'''{reduced_tensor} != {truth_tensor}'''
def        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     main()
def        UpperCAmelCase_     (		):





     '''simple docstring'''




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


if __name__ == "__main__":
      main()





from importlib import import_module

from .logging import get_logger


_SCREAMING_SNAKE_CASE  :						Optional[int]         =    get_logger(__name__)






class  UpperCAmelCase__	:






  """simple docstring"""


  def __init__(							self		:	Optional[int]			,     __lowerCamelCase		:	Optional[Any]			,     __lowerCamelCase		:	Any=None      )	->							int:
       SCREAMING_SNAKE_CASE__					=					attrs or []
       if module is not None:
            for key in module.__dict__:
                 if key in attrs or not key.startswith('''__'''      ):
                      setattr(self			,     __lowerCamelCase			,     getattr(__lowerCamelCase			,     __lowerCamelCase      )      )
       SCREAMING_SNAKE_CASE__					=					module._original_module if isinstance(__lowerCamelCase			,     _PatchedModuleObj      ) else module






class  UpperCAmelCase__	:






  """simple docstring"""


  a     =      []






  def __init__(							self		:	int			,     __lowerCamelCase		:	Union[str, Any]			,     __lowerCamelCase		:	str			,     __lowerCamelCase		:	Optional[int]			,     __lowerCamelCase		:	List[Any]=None      )	->							Any:
       SCREAMING_SNAKE_CASE__					=					obj
       SCREAMING_SNAKE_CASE__					=					target
       SCREAMING_SNAKE_CASE__					=					new
       SCREAMING_SNAKE_CASE__					=					target.split('''.'''      )[0]
       SCREAMING_SNAKE_CASE__					=					{}
       SCREAMING_SNAKE_CASE__					=					attrs or []






  def __enter__(							self		:	int      )	->							Tuple:
       *SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__					=					self.target.split('''.'''      )

       # Patch modules:
       # it's used to patch attributes of submodules like "os.path.join";
       # in this case we need to patch "os" and "os.path"

       for i in range(len(__lowerCamelCase      )      ):
            try:
                 SCREAMING_SNAKE_CASE__					=					import_module('''.'''.join(submodules[: i + 1]      )      )
            except ModuleNotFoundError:
                 continue
            # We iterate over all the globals in self.obj in case we find "os" or "os.path"
            for attr in self.obj.__dir__():
                 SCREAMING_SNAKE_CASE__					=					getattr(self.obj			,     __lowerCamelCase      )
                 # We don't check for the name of the global, but rather if its value *is* "os" or "os.path".
                 # This allows to patch renamed modules like "from os import path as ospath".
                 if obj_attr is submodule or (
                     (isinstance(__lowerCamelCase			,     _PatchedModuleObj      ) and obj_attr._original_module is submodule)
                 ):
                      SCREAMING_SNAKE_CASE__					=					obj_attr
                      # patch at top level
                      setattr(self.obj			,     __lowerCamelCase			,     _PatchedModuleObj(__lowerCamelCase			,     attrs=self.attrs      )      )
                      SCREAMING_SNAKE_CASE__					=					getattr(self.obj			,     __lowerCamelCase      )
                      # construct lower levels patches
                      for key in submodules[i + 1 :]:
                           setattr(__lowerCamelCase			,     __lowerCamelCase			,     _PatchedModuleObj(getattr(__lowerCamelCase			,     __lowerCamelCase			,     __lowerCamelCase      )			,     attrs=self.attrs      )      )
                           SCREAMING_SNAKE_CASE__					=					getattr(__lowerCamelCase			,     __lowerCamelCase      )
                      # finally set the target attribute
                      setattr(__lowerCamelCase			,     __lowerCamelCase			,     self.new      )

        # Patch attribute itself:
        # it's used for builtins like "open",
        # and also to patch "os.path.join" we may also need to patch "join"
        # itself if it was imported as "from os.path import join".

       if submodules:  # if it's an attribute of a submodule like "os.path.join"
            try:
                 SCREAMING_SNAKE_CASE__					=					getattr(import_module('''.'''.join(__lowerCamelCase      )      )			,     __lowerCamelCase      )
            except (AttributeError, ModuleNotFoundError):
                 return
            # We iterate over all the globals in self.obj in case we find "os.path.join"
            for attr in self.obj.__dir__():
                 # We don't check for the name of the global, but rather if its value *is* "os.path.join".
                 # This allows to patch renamed attributes like "from os.path import join as pjoin".
                 if getattr(self.obj			,     __lowerCamelCase      ) is attr_value:
                      SCREAMING_SNAKE_CASE__					=					getattr(self.obj			,     __lowerCamelCase      )
                      setattr(self.obj			,     __lowerCamelCase			,     self.new      )
       elif target_attr in globals()["__builtins__"]:  # if it'a s builtin like "open"
            SCREAMING_SNAKE_CASE__					=					globals()['''__builtins__'''][target_attr]
            setattr(self.obj			,     __lowerCamelCase			,     self.new      )
       else:
            raise RuntimeError(f'''Tried to patch attribute {target_attr} instead of a submodule.'''      )






  def __exit__(							self		:	Optional[Any]			,     *__lowerCamelCase		:	Tuple      )	->							List[str]:
       for attr in list(self.original      ):
            setattr(self.obj			,     __lowerCamelCase			,     self.original.pop(__lowerCamelCase      )      )






  def    lowercase_	(							self		:	List[Any]      )	->							Optional[Any]:
       self.__enter__()
       self._active_patches.append(self      )







  def    lowercase_	(							self		:	int      )	->							int:
       try:
            self._active_patches.remove(self      )
       except ValueError:
            # If the patch hasn't been started this will fail
            return None

       return self.__exit__()


