

# 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_(							_snake_case   :  Tuple	):




			"""simple docstring"""






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






def        UpperCamelCase_(							_snake_case   :  Union[str, Any]	):




			"""simple docstring"""






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






def        UpperCamelCase_(							_snake_case   :  List[str]	):




			"""simple docstring"""






			__a        =[state.process_index]
			__a        =gather_object(_snake_case	)
			assert len(_snake_case	) == state.num_processes, F'{gathered_obj}, {len(_snake_case	)} != {state.num_processes}'
			assert gathered_obj == list(range(state.num_processes	)	), F'{gathered_obj} != {list(range(state.num_processes	)	)}'






def        UpperCamelCase_(							_snake_case   :  Union[str, Any]	):




			"""simple docstring"""






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






def        UpperCamelCase_(							_snake_case   :  Optional[int]	):




			"""simple docstring"""






			if state.is_main_process:
						__a        =torch.arange(state.num_processes + 1	).to(state.device	)
			else:
						__a        =torch.arange(state.num_processes	).to(state.device	)
			__a        =pad_across_processes(_snake_case	)
			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_(							_snake_case   :  int	):




			"""simple docstring"""






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






def        UpperCamelCase_(							_snake_case   :  Dict	):




			"""simple docstring"""






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






def        UpperCamelCase_(							_snake_case   :  str	):




			"""simple docstring"""






			main()






def        UpperCamelCase_(							):




			"""simple docstring"""






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


if __name__ == "__main__":
						main()









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


_lowerCAmelCase					:			Tuple  				=			{
    "configuration_llama": ["LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP", "LlamaConfig"],
}

try:
						if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCAmelCase					:			Optional[int]  				=			["LlamaTokenizer"]

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCAmelCase					:			Optional[int]  				=			["LlamaTokenizerFast"]

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCAmelCase					:			List[Any]  				=			[
						    "LlamaForCausalLM",
						    "LlamaModel",
						    "LlamaPreTrainedModel",
						    "LlamaForSequenceClassification",
						]


if TYPE_CHECKING:
						from .configuration_llama import LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, LlamaConfig

						try:
												if not is_sentencepiece_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_llama import LlamaTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_llama_fast import LlamaTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_llama import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaPreTrainedModel


else:
						import sys

						_lowerCAmelCase					:			int  				=			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






