




# 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 					lowerCAmelCase_			(	__a		)					->		Dict:







  """simple docstring"""





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





def 					lowerCAmelCase_			(	__a		)					->		Optional[int]:







  """simple docstring"""





  lowerCamelCase__:				Optional[int]        =create_tensor(__a		)
  lowerCamelCase__:				Union[str, Any]        =gather(__a		)
  assert gathered_tensor.tolist() == list(range(1     ,   state.num_processes**2 + 1		)		)





def 					lowerCAmelCase_			(	__a		)					->		List[str]:







  """simple docstring"""





  lowerCamelCase__:				Optional[int]        =[state.process_index]
  lowerCamelCase__:				int        =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 					lowerCAmelCase_			(	__a		)					->		Union[str, Any]:







  """simple docstring"""





  lowerCamelCase__:				Any        =create_tensor(__a		)
  lowerCamelCase__:				Optional[Any]        =broadcast(__a		)
  assert broadcasted_tensor.shape == torch.Size([state.num_processes]		)
  assert broadcasted_tensor.tolist() == list(range(1     ,   state.num_processes + 1		)		)





def 					lowerCAmelCase_			(	__a		)					->		Any:







  """simple docstring"""





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







  """simple docstring"""





  if state.num_processes != 2:
    return
  lowerCamelCase__:				Union[str, Any]        =create_tensor(__a		)
  lowerCamelCase__:				Dict        =reduce(__a     ,   "sum"		)
  lowerCamelCase__:				List[Any]        =torch.tensor([4.0, 6]		).to(state.device		)
  assert torch.allclose(__a     ,   __a		), F"""{reduced_tensor} != {truth_tensor}"""





def 					lowerCAmelCase_			(	__a		)					->		Any:







  """simple docstring"""





  if state.num_processes != 2:
    return
  lowerCamelCase__:				int        =create_tensor(__a		)
  lowerCamelCase__:				List[str]        =reduce(__a     ,   "mean"		)
  lowerCamelCase__:				Tuple        =torch.tensor([2.0, 3]		).to(state.device		)
  assert torch.allclose(__a     ,   __a		), F"""{reduced_tensor} != {truth_tensor}"""





def 					lowerCAmelCase_			(	__a		)					->		Optional[int]:







  """simple docstring"""





  main()





def 					lowerCAmelCase_			(	)					->		Any:







  """simple docstring"""





  lowerCamelCase__:				int        =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()








import inspect
import os
import unittest

import torch

import accelerate
from accelerate import debug_launcher
from accelerate.test_utils import (
    execute_subprocess_async,
    require_cpu,
    require_huggingface_suite,
    require_multi_gpu,
    require_single_gpu,
)
from accelerate.utils import patch_environment



@require_huggingface_suite
class        _SCREAMING_SNAKE_CASE						(     unittest.TestCase   ):







 '''simple docstring'''

 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[int])	->Tuple:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =inspect.getfile(accelerate.test_utils)
   lowerCamelCase__:				str        =os.path.sep.join(
       mod_file.split(os.path.sep)[:-1] + ["scripts", "external_deps", "test_metrics.py"])

   from accelerate.test_utils.scripts.external_deps import test_metrics  # noqa: F401

   lowerCamelCase__:				str        =test_metrics

 @require_cpu
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Any:



   '''simple docstring'''



   debug_launcher(self.test_metrics.main	,			num_processes=1)

 @require_cpu
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->int:



   '''simple docstring'''



   debug_launcher(self.test_metrics.main)

 @require_single_gpu
 def 							SCREAMING_SNAKE_CASE_ (self    :       Any)	->int:



   '''simple docstring'''



   self.test_metrics.main()







 @require_multi_gpu
 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple)	->Dict:



   '''simple docstring'''



   print(F"""Found {torch.cuda.device_count()} devices.""")
   lowerCamelCase__:				Optional[Any]        =["torchrun", F"""--nproc_per_node={torch.cuda.device_count()}""", self.test_file_path]
   with patch_environment(omp_num_threads=1):
     execute_subprocess_async(UpperCAmelCase_	,			env=os.environ.copy())

