





# 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				(      snake_case							)		->    Any:
		return (torch.arange(state.num_processes							) + 1.0 + (state.num_processes * state.process_index)).to(state.device							)



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



def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    int:
		__lowercase											=     [state.process_index]
		__lowercase											=     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  SCREAMING_SNAKE_CASE				(      snake_case							)		->    List[str]:
		__lowercase											=     create_tensor(snake_case							)
		__lowercase											=     broadcast(snake_case							)
		assert broadcasted_tensor.shape == torch.Size([state.num_processes]							)
		assert broadcasted_tensor.tolist() == list(range(1						,   state.num_processes + 1							)							)



def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    List[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:
				__lowercase											=     torch.arange(state.num_processes + 1							).to(state.device							)
		else:
				__lowercase											=     torch.arange(state.num_processes							).to(state.device							)
		__lowercase											=     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  SCREAMING_SNAKE_CASE				(      snake_case							)		->    int:
		# For now runs on only two processes
		if state.num_processes != 2:
				return
		__lowercase											=     create_tensor(snake_case							)
		__lowercase											=     reduce(snake_case						,   'sum'							)
		__lowercase											=     torch.tensor([4.0, 6]							).to(state.device							)
		assert torch.allclose(snake_case						,   snake_case							), F"{reduced_tensor} != {truth_tensor}"



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



def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    Tuple:
		# For xla_spawn (TPUs)
		main()



def  SCREAMING_SNAKE_CASE				(      )		->    int:
		__lowercase											=     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()











import os
from collections import namedtuple

import pytest

from datasets import ClassLabel, Features, Sequence, Value
from datasets.commands.test import TestCommand
from datasets.info import DatasetInfo, DatasetInfosDict


SCREAMING_SNAKE_CASE_   :      Optional[Any]			=   namedtuple(
    '''_TestCommandArgs''',
    [
        '''dataset''',
        '''name''',
        '''cache_dir''',
        '''data_dir''',
        '''all_configs''',
        '''save_infos''',
        '''ignore_verifications''',
        '''force_redownload''',
        '''clear_cache''',
    ],
    defaults=[None, None, None, False, False, False, False, False],
)



def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    Optional[int]:
		return (abs(source - target							) / target) < 0.01



@pytest.mark.integration
def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    Dict:
		__lowercase											=     _TestCommandArgs(dataset=snake_case						,   all_configs=snake_case						,   save_infos=snake_case							)
		__lowercase											=     TestCommand(*snake_case							)
		test_command.run()
		__lowercase											=     os.path.join(snake_case						,   'README.md'							)
		assert os.path.exists(snake_case							)
		__lowercase											=     DatasetInfosDict.from_directory(snake_case							)
		__lowercase											=     DatasetInfosDict(
		    {
		        'default': DatasetInfo(
		            features=Features(
		                {
		                    'tokens': Sequence(Value('string'							)							),
		                    'ner_tags': Sequence(
		                        ClassLabel(names=['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC']							)							),
		                    'langs': Sequence(Value('string'							)							),
		                    'spans': Sequence(Value('string'							)							),
		                }							)						,   splits=[
		                {
		                    'name': 'train',
		                    'num_bytes': 2_351_563,
		                    'num_examples': 10_000,
		                },
		                {
		                    'name': 'validation',
		                    'num_bytes': 238_418,
		                    'num_examples': 1_000,
		                },
		            ]						,   download_size=3_940_680						,   dataset_size=2_589_981						,   )
		    }							)
		assert dataset_infos.keys() == expected_dataset_infos.keys()
		for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
				__lowercase						,      __lowercase											=     getattr(dataset_infos['default']						,   snake_case							), getattr(expected_dataset_infos['default']						,   snake_case							)
				if key == "num_bytes":
						assert is_apercent_close(snake_case						,   snake_case							)
				elif key == "splits":
						assert list(snake_case							) == list(snake_case							)
						for split in result:
								assert result[split].name == expected[split].name
								assert result[split].num_examples == expected[split].num_examples
								assert is_apercent_close(result[split].num_bytes						,   expected[split].num_bytes							)
				else:
						result == expected




