
"""simple docstring"""







import inspect
import os
from argparse import ArgumentParser, Namespace
from importlib import import_module

import huggingface_hub
import numpy as np
from packaging import version

from .. import (
    FEATURE_EXTRACTOR_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoImageProcessor,
    AutoProcessor,
    AutoTokenizer,
    is_datasets_available,
    is_tf_available,
    is_torch_available,
)
from ..utils import TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, logging
from . import BaseTransformersCLICommand


if is_tf_available():
							import tensorflow as tf

							tf.config.experimental.enable_tensor_float_aa_execution(False)

if is_torch_available():
							import torch

if is_datasets_available():
							from datasets import load_dataset


lowerCamelCase__     =						5E-5  # larger error tolerance than in our internal tests, to avoid flaky user-facing errors
def        __lowerCAmelCase							(_UpperCamelCase						):
	return PTtoTFCommand(
	    args.model_name			,    args.local_dir			,    args.max_error			,    args.new_weights			,    args.no_pr			,    args.push			,    args.extra_commit_description			,    args.override_model_class			,    )




class 					A__			(      _lowerCamelCase):

		@staticmethod
		def 				__lowerCamelCase				( _SCREAMING_SNAKE_CASE					):
			__lowerCAmelCase     :		Tuple		     = parser.add_parser(
			    'pt-to-tf'			,						help=(
			        'CLI tool to run convert a transformers model from a PyTorch checkpoint to a TensorFlow checkpoint.'
			        ' Can also be used to validate existing weights without opening PRs, with --no-pr.'
			    )			,						)
			train_parser.add_argument(
			    '--model-name'			,						type=_SCREAMING_SNAKE_CASE			,						required=_SCREAMING_SNAKE_CASE			,						help='The model name, including owner/organization, as seen on the hub.'			,						)
			train_parser.add_argument(
			    '--local-dir'			,						type=_SCREAMING_SNAKE_CASE			,						default=''			,						help='Optional local directory of the model repository. Defaults to /tmp/{model_name}'			,						)
			train_parser.add_argument(
			    '--max-error'			,						type=_SCREAMING_SNAKE_CASE			,						default=_SCREAMING_SNAKE_CASE			,						help=(
			        f"Maximum error tolerance. Defaults to {MAX_ERROR}. This flag should be avoided, use at your own risk."
			    )			,						)
			train_parser.add_argument(
			    '--new-weights'			,						action='store_true'			,						help='Optional flag to create new TensorFlow weights, even if they already exist.'			,						)
			train_parser.add_argument(
			    '--no-pr'			,						action='store_true'			,						help='Optional flag to NOT open a PR with converted weights.'					)
			train_parser.add_argument(
			    '--push'			,						action='store_true'			,						help='Optional flag to push the weights directly to `main` (requires permissions)'			,						)
			train_parser.add_argument(
			    '--extra-commit-description'			,						type=_SCREAMING_SNAKE_CASE			,						default=''			,						help='Optional additional commit description to use when opening a PR (e.g. to tag the owner).'			,						)
			train_parser.add_argument(
			    '--override-model-class'			,						type=_SCREAMING_SNAKE_CASE			,						default=_SCREAMING_SNAKE_CASE			,						help='If you think you know better than the auto-detector, you can specify the model class here. '
			    'Can be either an AutoModel class or a specific model class like BertForSequenceClassification.'			,						)
			train_parser.set_defaults(func=_SCREAMING_SNAKE_CASE					)

		@staticmethod
		def 				__lowerCamelCase				( _SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE					):
			__lowerCAmelCase     :		int		     = set(pt_outputs.keys()					)
			__lowerCAmelCase     :		Union[str, Any]		     = set(tf_outputs.keys()					)
			if pt_out_attrs != tf_out_attrs:
				raise ValueError(
				    f"The model outputs have different attributes, aborting. (Pytorch: {pt_out_attrs}, TensorFlow:"
				    f" {tf_out_attrs})"					)

			# 2. For each output attribute, computes the difference
			def _find_pt_tf_differences(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE=""					):
				# If the current attribute is a tensor, it is a leaf and we make the comparison. Otherwise, we will dig in
				# recursivelly, keeping the name of the attribute.
				if isinstance(_SCREAMING_SNAKE_CASE			,						torch.Tensor					):
					__lowerCAmelCase     :		Tuple		     = np.max(np.abs(pt_out.numpy() - tf_out.numpy()					)					)
					__lowerCAmelCase     :		Union[str, Any]		     = tensor_difference
				else:
					__lowerCAmelCase     :		Dict		     = attr_name
					for i, pt_item in enumerate(_SCREAMING_SNAKE_CASE					):
						# If it is a named attribute, we keep the name. Otherwise, just its index.
						if isinstance(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE					):
							__lowerCAmelCase     :		Union[str, Any]		     = root_name + pt_item
							__lowerCAmelCase     :		Tuple		     = tf_out[pt_item]
							__lowerCAmelCase     :		Any		     = pt_out[pt_item]
						else:
							__lowerCAmelCase     :		Dict		     = root_name + f"[{i}]"
							__lowerCAmelCase     :		str		     = tf_out[i]
						__lowerCAmelCase     :		Union[str, Any]		     = _find_pt_tf_differences(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE					)

				return differences

			return _find_pt_tf_differences(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						{}					)

		def __init__( self			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						*_SCREAMING_SNAKE_CASE			,						):
			__lowerCAmelCase     :		str		     = logging.get_logger('transformers-cli/pt_to_tf'					)
			__lowerCAmelCase     :		Union[str, Any]		     = model_name
			__lowerCAmelCase     :		Optional[Any]		     = local_dir if local_dir else os.path.join('/tmp'			,						_SCREAMING_SNAKE_CASE					)
			__lowerCAmelCase     :		Optional[Any]		     = max_error
			__lowerCAmelCase     :		int		     = new_weights
			__lowerCAmelCase     :		List[str]		     = no_pr
			__lowerCAmelCase     :		List[str]		     = push
			__lowerCAmelCase     :		Union[str, Any]		     = extra_commit_description
			__lowerCAmelCase     :		List[Any]		     = override_model_class

		def 				__lowerCamelCase				( self			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE					):

			def _get_audio_input():
				__lowerCAmelCase     :		Any		     = load_dataset('hf-internal-testing/librispeech_asr_dummy'			,						'clean'			,						split='validation'					)
				__lowerCAmelCase     :		Dict		     = ds.sort('id'					).select(range(2					)					)[:2]['audio']
				__lowerCAmelCase     :		Any		     = [x['array'] for x in speech_samples]
				return raw_samples

			__lowerCAmelCase     :		str		     = type(pt_model.config					)
			if model_config_class in PROCESSOR_MAPPING:
				__lowerCAmelCase     :		List[str]		     = AutoProcessor.from_pretrained(self._local_dir					)
				if model_config_class in TOKENIZER_MAPPING and processor.tokenizer.pad_token is None:
					__lowerCAmelCase     :		Tuple		     = processor.tokenizer.eos_token
			elif model_config_class in IMAGE_PROCESSOR_MAPPING:
				__lowerCAmelCase     :		int		     = AutoImageProcessor.from_pretrained(self._local_dir					)
			elif model_config_class in FEATURE_EXTRACTOR_MAPPING:
				__lowerCAmelCase     :		str		     = AutoFeatureExtractor.from_pretrained(self._local_dir					)
			elif model_config_class in TOKENIZER_MAPPING:
				__lowerCAmelCase     :		List[str]		     = AutoTokenizer.from_pretrained(self._local_dir					)
				if processor.pad_token is None:
					__lowerCAmelCase     :		str		     = processor.eos_token
			else:
				raise ValueError(f"Unknown data processing type (model config type: {model_config_class})"					)

			__lowerCAmelCase     :		Dict		     = set(inspect.signature(pt_model.forward					).parameters.keys()					)
			__lowerCAmelCase     :		Tuple		     = {}
			if "input_ids" in model_forward_signature:
				processor_inputs.update(
				    {
				        'text': ['Hi there!', 'I am a batch with more than one row and different input lengths.'],
				        'padding': True,
				        'truncation': True,
				    }					)
			if "pixel_values" in model_forward_signature:
				__lowerCAmelCase     :		Union[str, Any]		     = load_dataset('cifar10'			,						'plain_text'			,						split='test'					)[:2]['img']
				processor_inputs.update({'images': sample_images}					)
			if "input_features" in model_forward_signature:
				__lowerCAmelCase     :		Optional[int]		     = inspect.signature(processor.feature_extractor					).parameters
				# Pad to the largest input length by default but take feature extractor default
				# padding value if it exists e.g. "max_length" and is not False or None
				if "padding" in feature_extractor_signature:
					__lowerCAmelCase     :		Dict		     = feature_extractor_signature['padding'].default
					if default_strategy is not False and default_strategy is not None:
						__lowerCAmelCase     :		Any		     = default_strategy
					else:
						__lowerCAmelCase     :		Dict		     = True
				else:
					__lowerCAmelCase     :		Optional[int]		     = True
				processor_inputs.update({'audio': _get_audio_input(), 'padding': padding_strategy}					)
			if "input_values" in model_forward_signature:  # Wav2Vec2 audio input
				processor_inputs.update({'audio': _get_audio_input(), 'padding': True}					)
			__lowerCAmelCase     :		List[Any]		     = processor(**_SCREAMING_SNAKE_CASE			,						return_tensors='pt'					)
			__lowerCAmelCase     :		str		     = processor(**_SCREAMING_SNAKE_CASE			,						return_tensors='tf'					)

			# Extra input requirements, in addition to the input modality
			if (
			    config.is_encoder_decoder
			    or (hasattr(_SCREAMING_SNAKE_CASE			,						'encoder'					) and hasattr(_SCREAMING_SNAKE_CASE			,						'decoder'					))
			    or "decoder_input_ids" in tf_dummy_inputs
			):
				__lowerCAmelCase     :		Optional[int]		     = np.asarray([[1], [1]]			,						dtype=_SCREAMING_SNAKE_CASE					) * (pt_model.config.decoder_start_token_id or 0)
				pt_input.update({'decoder_input_ids': torch.tensor(_SCREAMING_SNAKE_CASE					)}					)
				tf_input.update({'decoder_input_ids': tf.convert_to_tensor(_SCREAMING_SNAKE_CASE					)}					)

			return pt_input, tf_input







		def 				__lowerCamelCase				( self					):
			# hub version 0.9.0 introduced the possibility of programmatically opening PRs with normal write tokens.
			if version.parse(huggingface_hub.__version__					) < version.parse('0.9.0'					):
				raise ImportError(
				    'The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub'
				    ' installation.'					)
			else:
				from huggingface_hub import Repository, create_commit
				from huggingface_hub._commit_api import CommitOperationAdd

			# Fetch remote data
			__lowerCAmelCase     :		Union[str, Any]		     = Repository(local_dir=self._local_dir			,						clone_from=self._model_name					)

			# Load config and get the appropriate architecture -- the latter is needed to convert the head's weights
			__lowerCAmelCase     :		int		     = AutoConfig.from_pretrained(self._local_dir					)
			__lowerCAmelCase     :		List[str]		     = config.architectures
			if self._override_model_class is not None:
				if self._override_model_class.startswith('TF'					):
					__lowerCAmelCase     :		Dict		     = [self._override_model_class[2:]]
				else:
					__lowerCAmelCase     :		Dict		     = [self._override_model_class]
				try:
					__lowerCAmelCase     :		Dict		     = getattr(import_module('transformers'					)			,						architectures[0]					)
				except AttributeError:
					raise ValueError(f"Model class {self._override_model_class} not found in transformers."					)
				try:
					__lowerCAmelCase     :		str		     = getattr(import_module('transformers'					)			,						'TF' + architectures[0]					)
				except AttributeError:
					raise ValueError(f"TF model class TF{self._override_model_class} not found in transformers."					)
			elif architectures is None:  # No architecture defined -- use auto classes
				__lowerCAmelCase     :		Optional[Any]		     = getattr(import_module('transformers'					)			,						'AutoModel'					)
				__lowerCAmelCase     :		str		     = getattr(import_module('transformers'					)			,						'TFAutoModel'					)
				self._logger.warning('No detected architecture, using AutoModel/TFAutoModel'					)
			else:  # Architecture defined -- use it
				if len(_SCREAMING_SNAKE_CASE					) > 1:
					raise ValueError(f"More than one architecture was found, aborting. (architectures = {architectures})"					)
				self._logger.warning(f"Detected architecture: {architectures[0]}"					)
				__lowerCAmelCase     :		List[Any]		     = getattr(import_module('transformers'					)			,						architectures[0]					)
				try:
					__lowerCAmelCase     :		Optional[Any]		     = getattr(import_module('transformers'					)			,						'TF' + architectures[0]					)
				except AttributeError:
					raise AttributeError(f"The TensorFlow equivalent of {architectures[0]} doesn't exist in transformers."					)

        # Check the TF dummy inputs to see what keys we need in the forward pass
			__lowerCAmelCase     :		int		     = tf_class.from_config(_SCREAMING_SNAKE_CASE					)
			__lowerCAmelCase     :		int		     = tf_from_pt_model.dummy_inputs

			del tf_from_pt_model  # Try to keep only one model in memory at a time

			# Load the model and get some basic inputs
			__lowerCAmelCase     :		str		     = pt_class.from_pretrained(self._local_dir					)
			pt_model.eval()

			__lowerCAmelCase     ,  __lowerCAmelCase     :		Tuple		     = self.get_inputs(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE					)

			with torch.no_grad():
				__lowerCAmelCase     :		int		     = pt_model(**_SCREAMING_SNAKE_CASE			,						output_hidden_states=_SCREAMING_SNAKE_CASE					)
			del pt_model  # will no longer be used, and may have a large memory footprint

			__lowerCAmelCase     :		Any		     = tf_class.from_pretrained(self._local_dir			,						from_pt=_SCREAMING_SNAKE_CASE					)
			__lowerCAmelCase     :		Dict		     = tf_from_pt_model(**_SCREAMING_SNAKE_CASE			,						output_hidden_states=_SCREAMING_SNAKE_CASE			,						training=_SCREAMING_SNAKE_CASE					)

			# Confirms that cross loading PT weights into TF worked.
			__lowerCAmelCase     :		str		     = self.find_pt_tf_differences(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE					)
			__lowerCAmelCase     :		str		     = {k: v for k, v in crossload_differences.items() if 'hidden' not in k}
			__lowerCAmelCase     :		Tuple		     = {k: v for k, v in crossload_differences.items() if 'hidden' in k}
			if len(_SCREAMING_SNAKE_CASE					) == 0 and architectures is not None:
				raise ValueError(
				    f"Something went wrong -- the config file has architectures ({architectures}), but no model head"
				    ' output was found. All outputs start with \'hidden\''					)
			__lowerCAmelCase     :		Optional[Any]		     = max(output_differences.values()					) if output_differences else 0.0
			__lowerCAmelCase     :		str		     = max(hidden_differences.values()					)
			if max_crossload_output_diff > self._max_error or max_crossload_hidden_diff > self._max_error:
				raise ValueError(
				    'The cross-loaded TensorFlow model has different outputs, something went wrong!\n'
				    + f"\nList of maximum output differences above the threshold ({self._max_error}):\n"
				    + '\n'.join([f"{k}: {v:.3e}" for k, v in output_differences.items() if v > self._max_error]					)
				    + f"\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n"
				    + '\n'.join([f"{k}: {v:.3e}" for k, v in hidden_differences.items() if v > self._max_error]					)					)

			# Save the weights in a TF format (if needed) and confirms that the results are still good
			__lowerCAmelCase     :		int		     = os.path.join(self._local_dir			,						_SCREAMING_SNAKE_CASE					)
			__lowerCAmelCase     :		int		     = os.path.join(self._local_dir			,						_SCREAMING_SNAKE_CASE					)
			if (not os.path.exists(_SCREAMING_SNAKE_CASE					) and not os.path.exists(_SCREAMING_SNAKE_CASE					)) or self._new_weights:
				tf_from_pt_model.save_pretrained(self._local_dir					)
			del tf_from_pt_model  # will no longer be used, and may have a large memory footprint

			__lowerCAmelCase     :		Optional[int]		     = tf_class.from_pretrained(self._local_dir					)
			__lowerCAmelCase     :		Union[str, Any]		     = tf_model(**_SCREAMING_SNAKE_CASE			,						output_hidden_states=_SCREAMING_SNAKE_CASE					)

			__lowerCAmelCase     :		str		     = self.find_pt_tf_differences(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE					)
			__lowerCAmelCase     :		Union[str, Any]		     = {k: v for k, v in conversion_differences.items() if 'hidden' not in k}
			__lowerCAmelCase     :		Tuple		     = {k: v for k, v in conversion_differences.items() if 'hidden' in k}
			if len(_SCREAMING_SNAKE_CASE					) == 0 and architectures is not None:
				raise ValueError(
				    f"Something went wrong -- the config file has architectures ({architectures}), but no model head"
				    ' output was found. All outputs start with \'hidden\''					)
			__lowerCAmelCase     :		List[str]		     = max(output_differences.values()					) if output_differences else 0.0
			__lowerCAmelCase     :		int		     = max(hidden_differences.values()					)
			if max_conversion_output_diff > self._max_error or max_conversion_hidden_diff > self._max_error:
				raise ValueError(
				    'The converted TensorFlow model has different outputs, something went wrong!\n'
				    + f"\nList of maximum output differences above the threshold ({self._max_error}):\n"
				    + '\n'.join([f"{k}: {v:.3e}" for k, v in output_differences.items() if v > self._max_error]					)
				    + f"\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n"
				    + '\n'.join([f"{k}: {v:.3e}" for k, v in hidden_differences.items() if v > self._max_error]					)					)

			__lowerCAmelCase     :		Tuple		     = 'Update TF weights' if self._new_weights else 'Add TF weights'
			if self._push:
				repo.git_add(auto_lfs_track=_SCREAMING_SNAKE_CASE					)
				repo.git_commit(_SCREAMING_SNAKE_CASE					)
				repo.git_push(blocking=_SCREAMING_SNAKE_CASE					)  # this prints a progress bar with the upload
				self._logger.warning(f"TF weights pushed into {self._model_name}"					)
			elif not self._no_pr:
				self._logger.warning('Uploading the weights into a new PR...'					)
				__lowerCAmelCase     :		int		     = (
				    'Model converted by the [`transformers`\' `pt_to_tf`'
				    ' CLI](https://github.com/huggingface/transformers/blob/main/src/transformers/commands/pt_to_tf.py). '
				    'All converted model outputs and hidden layers were validated against its PyTorch counterpart.\n\n'
				    f"Maximum crossload output difference={max_crossload_output_diff:.3e}; "
				    f"Maximum crossload hidden layer difference={max_crossload_hidden_diff:.3e};\n"
				    f"Maximum conversion output difference={max_conversion_output_diff:.3e}; "
				    f"Maximum conversion hidden layer difference={max_conversion_hidden_diff:.3e};\n"
				)
				if self._max_error > MAX_ERROR:
					commit_descrition += (
					    f"\n\nCAUTION: The maximum admissible error was manually increased to {self._max_error}!"
					)
				if self._extra_commit_description:
					commit_descrition += "\n\n" + self._extra_commit_description

				# sharded model -> adds all related files (index and .h5 shards)
				if os.path.exists(_SCREAMING_SNAKE_CASE					):
					__lowerCAmelCase     :		List[str]		     = [
					    CommitOperationAdd(path_in_repo=_SCREAMING_SNAKE_CASE			,						path_or_fileobj=_SCREAMING_SNAKE_CASE					)
					]
					for shard_path in tf.io.gfile.glob(self._local_dir + '/tf_model-*.h5'					):
						operations += [
						    CommitOperationAdd(path_in_repo=os.path.basename(_SCREAMING_SNAKE_CASE					)			,						path_or_fileobj=_SCREAMING_SNAKE_CASE					)
						]
				else:
					__lowerCAmelCase     :		Union[str, Any]		     = [CommitOperationAdd(path_in_repo=_SCREAMING_SNAKE_CASE			,						path_or_fileobj=_SCREAMING_SNAKE_CASE					)]

				__lowerCAmelCase     :		Optional[Any]		     = create_commit(
				    repo_id=self._model_name			,						operations=_SCREAMING_SNAKE_CASE			,						commit_message=_SCREAMING_SNAKE_CASE			,						commit_description=_SCREAMING_SNAKE_CASE			,						repo_type='model'			,						create_pr=_SCREAMING_SNAKE_CASE			,						).pr_url
				self._logger.warning(f"PR open in {hub_pr_url}"					)

"""simple docstring"""







import os
import sys
from pathlib import Path

import pytest

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node

from .utils import execute_subprocess_async, get_torch_dist_unique_port, require_torch
def        __lowerCAmelCase							():
	__lowerCAmelCase     :		Optional[int]		     = Dataset.from_dict({'i': range(17						)}						)
	__lowerCAmelCase     :		List[Any]		     = len(_UpperCamelCase						)
	__lowerCAmelCase     :		Tuple		     = 3
	__lowerCAmelCase     :		List[Any]		     = [
	    split_dataset_by_node(_UpperCamelCase			,    rank=_UpperCamelCase			,    world_size=_UpperCamelCase						) for rank in range(_UpperCamelCase						)
	]
	assert sum(len(_UpperCamelCase						) for ds in datasets_per_rank						) == full_size
	assert len({tuple(x.values()						) for ds in datasets_per_rank for x in ds}						) == full_size
def        __lowerCAmelCase							():
	def gen():
		return ({"i": i} for i in range(17						))

	__lowerCAmelCase     :		List[Any]		     = 3
	__lowerCAmelCase     :		Optional[Any]		     = IterableDataset.from_generator(_UpperCamelCase						)
	__lowerCAmelCase     :		List[Any]		     = len(list(_UpperCamelCase						)						)
	__lowerCAmelCase     :		str		     = [
	    split_dataset_by_node(_UpperCamelCase			,    rank=_UpperCamelCase			,    world_size=_UpperCamelCase						) for rank in range(_UpperCamelCase						)
	]
	assert sum(len(list(_UpperCamelCase						)						) for ds in datasets_per_rank						) == full_size
	assert len({tuple(x.values()						) for ds in datasets_per_rank for x in ds}						) == full_size
@pytest.mark.parametrize('shards_per_node'			,    [1, 2, 3]						)
def        __lowerCAmelCase							(_UpperCamelCase						):
	def gen(_UpperCamelCase						):
		for shard in shards:
			yield from ({"i": i, "shard": shard} for i in range(17						))

	__lowerCAmelCase     :		List[Any]		     = 3
	__lowerCAmelCase     :		Optional[int]		     = shards_per_node * world_size
	__lowerCAmelCase     :		Dict		     = {'shards': [F"shard_{shard_idx}.txt" for shard_idx in range(_UpperCamelCase						)]}
	__lowerCAmelCase     :		Dict		     = IterableDataset.from_generator(_UpperCamelCase			,    gen_kwargs=_UpperCamelCase						)
	__lowerCAmelCase     :		str		     = len(list(_UpperCamelCase						)						)
	assert full_ds.n_shards == world_size * shards_per_node
	__lowerCAmelCase     :		Union[str, Any]		     = [
	    split_dataset_by_node(_UpperCamelCase			,    rank=_UpperCamelCase			,    world_size=_UpperCamelCase						) for rank in range(_UpperCamelCase						)
	]
	assert [ds.n_shards for ds in datasets_per_rank] == [shards_per_node] * world_size
	assert sum(len(list(_UpperCamelCase						)						) for ds in datasets_per_rank						) == full_size
	assert len({tuple(x.values()						) for ds in datasets_per_rank for x in ds}						) == full_size
def        __lowerCAmelCase							():
	def gen():
		return ({"i": i} for i in range(17						))

	__lowerCAmelCase     :		str		     = 2
	__lowerCAmelCase     :		Any		     = IterableDataset.from_generator(_UpperCamelCase						)
	__lowerCAmelCase     :		Optional[int]		     = len(list(_UpperCamelCase						)						)

	__lowerCAmelCase     :		Dict		     = split_dataset_by_node(_UpperCamelCase			,    rank=0			,    world_size=_UpperCamelCase						).shuffle(seed=42						)
	assert len(list(_UpperCamelCase						)						) == 1 + full_size // world_size
	with pytest.raises(_UpperCamelCase						):
		split_dataset_by_node(_UpperCamelCase			,    rank=0			,    world_size=_UpperCamelCase						).shuffle()

	__lowerCAmelCase     :		List[str]		     = split_dataset_by_node(full_ds.shuffle(seed=42						)			,    rank=0			,    world_size=_UpperCamelCase						)
	assert len(list(_UpperCamelCase						)						) == 1 + full_size // world_size
	with pytest.raises(_UpperCamelCase						):
		split_dataset_by_node(full_ds.shuffle()			,    rank=0			,    world_size=_UpperCamelCase						)
@pytest.mark.parametrize('streaming'			,    [False, True]						)
@require_torch
@pytest.mark.skipif(os.name == 'nt'			,    reason='execute_subprocess_async doesn\'t support windows'						)
@pytest.mark.integration
def        __lowerCAmelCase							(_UpperCamelCase						):
	__lowerCAmelCase     :		Tuple		     = 2
	__lowerCAmelCase     :		Optional[Any]		     = get_torch_dist_unique_port()
	__lowerCAmelCase     :		List[Any]		     = Path(_UpperCamelCase						).resolve().parent / 'distributed_scripts' / 'launch_torch_distributed.py'
	__lowerCAmelCase     :		int		     = F"\n        -m torch.distributed.launch\n        --nproc_per_node={nproc_per_node}\n        --master_port={master_port}\n        {test_script}\n    ".split()
	__lowerCAmelCase     :		Tuple		     = F"\n        --streaming={streaming}\n    ".split()
	__lowerCAmelCase     :		Optional[Any]		     = [sys.executable] + distributed_args + args
	execute_subprocess_async(_UpperCamelCase			,    env=os.environ.copy()						)
@pytest.mark.parametrize(
    'nproc_per_node, num_workers'			,    [
        (2, 2),  # each node has 2 shards and each worker has 1 shards
        (3, 2),  # each node uses all the shards but skips examples, and each worker has 2 shards
    ]			,    )
@require_torch
@pytest.mark.skipif(os.name == 'nt'			,    reason='execute_subprocess_async doesn\'t support windows'						)
@pytest.mark.integration
def        __lowerCAmelCase							(_UpperCamelCase			,    _UpperCamelCase						):
	__lowerCAmelCase     :		Optional[int]		     = True
	__lowerCAmelCase     :		Optional[Any]		     = get_torch_dist_unique_port()
	__lowerCAmelCase     :		Tuple		     = Path(_UpperCamelCase						).resolve().parent / 'distributed_scripts' / 'launch_torch_distributed.py'
	__lowerCAmelCase     :		Optional[Any]		     = F"\n        -m torch.distributed.launch\n        --nproc_per_node={nproc_per_node}\n        --master_port={master_port}\n        {test_script}\n    ".split()
	__lowerCAmelCase     :		List[Any]		     = F"\n        --streaming={streaming}\n        --num_workers={num_workers}\n    ".split()
	__lowerCAmelCase     :		List[Any]		     = [sys.executable] + distributed_args + args
	execute_subprocess_async(_UpperCamelCase			,    env=os.environ.copy()						)