
'''simple docstring'''







import torch

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



def 			_a		(	_SCREAMING_SNAKE_CASE		:   int  )     ->   Tuple:
						return (torch.arange(state.num_processes  ) + 1.0 + (state.num_processes * state.process_index)).to(state.device  )



def 			_a		(	_SCREAMING_SNAKE_CASE		:   Optional[Any]  )     ->   Any:
						_SCREAMING_SNAKE_CASE				= create_tensor(_SCREAMING_SNAKE_CASE  )
						_SCREAMING_SNAKE_CASE				= gather(_SCREAMING_SNAKE_CASE  )
						assert gathered_tensor.tolist() == list(range(1	,			state.num_processes**2 + 1  )  )



def 			_a		(	_SCREAMING_SNAKE_CASE		:   List[str]  )     ->   List[str]:
						_SCREAMING_SNAKE_CASE				= [state.process_index]
						_SCREAMING_SNAKE_CASE				= gather_object(_SCREAMING_SNAKE_CASE  )
						assert len(_SCREAMING_SNAKE_CASE  ) == state.num_processes, F'{gathered_obj}, {len(_SCREAMING_SNAKE_CASE  )} != {state.num_processes}'
						assert gathered_obj == list(range(state.num_processes  )  ), F'{gathered_obj} != {list(range(state.num_processes  )  )}'



def 			_a		(	_SCREAMING_SNAKE_CASE		:   List[str]  )     ->   List[str]:
						_SCREAMING_SNAKE_CASE				= create_tensor(_SCREAMING_SNAKE_CASE  )
						_SCREAMING_SNAKE_CASE				= broadcast(_SCREAMING_SNAKE_CASE  )
						assert broadcasted_tensor.shape == torch.Size([state.num_processes]  )
						assert broadcasted_tensor.tolist() == list(range(1	,			state.num_processes + 1  )  )



def 			_a		(	_SCREAMING_SNAKE_CASE		:   Any  )     ->   Union[str, Any]:
						# We need to pad the tensor with one more element if we are the main process
						# to ensure that we can pad
						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(_SCREAMING_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 			_a		(	_SCREAMING_SNAKE_CASE		:   Union[str, Any]  )     ->   Union[str, Any]:
						# For now runs on only two processes
						if state.num_processes != 2:
												return
						_SCREAMING_SNAKE_CASE				= create_tensor(_SCREAMING_SNAKE_CASE  )
						_SCREAMING_SNAKE_CASE				= reduce(_SCREAMING_SNAKE_CASE	,			"sum"  )
						_SCREAMING_SNAKE_CASE				= torch.tensor([4.0, 6]  ).to(state.device  )
						assert torch.allclose(_SCREAMING_SNAKE_CASE	,			_SCREAMING_SNAKE_CASE  ), F'{reduced_tensor} != {truth_tensor}'



def 			_a		(	_SCREAMING_SNAKE_CASE		:   Any  )     ->   Any:
						# For now runs on only two processes
						if state.num_processes != 2:
												return
						_SCREAMING_SNAKE_CASE				= create_tensor(_SCREAMING_SNAKE_CASE  )
						_SCREAMING_SNAKE_CASE				= reduce(_SCREAMING_SNAKE_CASE	,			"mean"  )
						_SCREAMING_SNAKE_CASE				= torch.tensor([2.0, 3]  ).to(state.device  )
						assert torch.allclose(_SCREAMING_SNAKE_CASE	,			_SCREAMING_SNAKE_CASE  ), F'{reduced_tensor} != {truth_tensor}'



def 			_a		(	_SCREAMING_SNAKE_CASE		:   int  )     ->   List[str]:
						# For xla_spawn (TPUs)
						main()



def 			_a		(	)     ->   Tuple:
						_SCREAMING_SNAKE_CASE				= PartialState()
						state.print(F'State: {state}'  )
						state.print("testing gather"  )
						test_gather(_SCREAMING_SNAKE_CASE  )
						state.print("testing gather_object"  )
						test_gather_object(_SCREAMING_SNAKE_CASE  )
						state.print("testing broadcast"  )
						test_broadcast(_SCREAMING_SNAKE_CASE  )
						state.print("testing pad_across_processes"  )
						test_pad_across_processes(_SCREAMING_SNAKE_CASE  )
						state.print("testing reduce_sum"  )
						test_reduce_sum(_SCREAMING_SNAKE_CASE  )
						state.print("testing reduce_mean"  )
						test_reduce_mean(_SCREAMING_SNAKE_CASE  )


if __name__ == "__main__":
						main()

'''simple docstring'''







from typing import TYPE_CHECKING

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


_snake_case				:   int     =				{
    """configuration_longformer""": [
        """LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """LongformerConfig""",
        """LongformerOnnxConfig""",
    ],
    """tokenization_longformer""": ["""LongformerTokenizer"""],
}

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

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_snake_case				:   int     =				[
						    """LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """LongformerForMaskedLM""",
						    """LongformerForMultipleChoice""",
						    """LongformerForQuestionAnswering""",
						    """LongformerForSequenceClassification""",
						    """LongformerForTokenClassification""",
						    """LongformerModel""",
						    """LongformerPreTrainedModel""",
						    """LongformerSelfAttention""",
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_snake_case				:   Tuple     =				[
						    """TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """TFLongformerForMaskedLM""",
						    """TFLongformerForMultipleChoice""",
						    """TFLongformerForQuestionAnswering""",
						    """TFLongformerForSequenceClassification""",
						    """TFLongformerForTokenClassification""",
						    """TFLongformerModel""",
						    """TFLongformerPreTrainedModel""",
						    """TFLongformerSelfAttention""",
						]


if TYPE_CHECKING:
						from .configuration_longformer import (
						    LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
						    LongformerConfig,
						    LongformerOnnxConfig,
						)
						from .tokenization_longformer import LongformerTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_longformer_fast import LongformerTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_longformer import (
												    LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
												    LongformerForMaskedLM,
												    LongformerForMultipleChoice,
												    LongformerForQuestionAnswering,
												    LongformerForSequenceClassification,
												    LongformerForTokenClassification,
												    LongformerModel,
												    LongformerPreTrainedModel,
												    LongformerSelfAttention,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_longformer import (
												    TF_LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TFLongformerForMaskedLM,
												    TFLongformerForMultipleChoice,
												    TFLongformerForQuestionAnswering,
												    TFLongformerForSequenceClassification,
												    TFLongformerForTokenClassification,
												    TFLongformerModel,
												    TFLongformerPreTrainedModel,
												    TFLongformerSelfAttention,
												)

else:
						import sys

						_snake_case				:   Dict     =				_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)