





import pytest

from datasets.splits import SplitDict, SplitInfo
from datasets.utils.py_utils import asdict


@pytest.mark.parametrize(
    'split_dict'   ,    [
        SplitDict(),
        SplitDict({'train': SplitInfo(name='train'   ,    num_bytes=1_337   ,    num_examples=42   ,    dataset_name='my_dataset'			)}			),
        SplitDict({'train': SplitInfo(name='train'   ,    num_bytes=1_337   ,    num_examples=42			)}			),
        SplitDict({'train': SplitInfo()}			),
    ]   ,    )
def       SCREAMING_SNAKE_CASE__  (       _lowercase     : Any			)   ->							List[Any]:




		'''simple docstring'''
		lowercase__					:      Union[str, Any]      =		split_dict._to_yaml_list()
		assert len(__UpperCamelCase			) == len(__UpperCamelCase			)
		lowercase__					:      Any      =		SplitDict._from_yaml_list(__UpperCamelCase			)
		for split_name, split_info in split_dict.items():
				# dataset_name field is deprecated, and is therefore not part of the YAML dump
				lowercase__					:      List[str]      =		None
				# the split name of split_dict takes over the name of the split info object
				lowercase__					:      Optional[Any]      =		split_name
		assert split_dict == reloaded





@pytest.mark.parametrize(
    'split_info'   ,    [SplitInfo(), SplitInfo(dataset_name=__UpperCamelCase			), SplitInfo(dataset_name='my_dataset'			)]			)
def       SCREAMING_SNAKE_CASE__  (       _lowercase     : str			)   ->							int:




		'''simple docstring'''
		# For backward compatibility, we need asdict(split_dict) to return split info dictrionaries with the "dataset_name"
		# field even if it's deprecated. This way old versionso of `datasets` can still reload dataset_infos.json files
		lowercase__					:      List[Any]      =		asdict(SplitDict({'train': split_info}			)			)
		assert "dataset_name" in split_dict_asdict["train"]
		assert split_dict_asdict["train"]["dataset_name"] == split_info.dataset_name







import operator as op


lowerCamelCase		       =      """scaler.pt"""
lowerCamelCase		       =      """pytorch_model"""
lowerCamelCase		       =      """random_states"""
lowerCamelCase		       =      """optimizer"""
lowerCamelCase		       =      """scheduler"""
lowerCamelCase		       =      """pytorch_model.bin"""
lowerCamelCase		       =      """pytorch_model.bin.index.json"""
lowerCamelCase		       =      """model.safetensors"""
lowerCamelCase		       =      """model.safetensors.index.json"""
lowerCamelCase		       =      """1.10.2"""
lowerCamelCase		       =      """py38"""
lowerCamelCase		       =      """4.17.0"""
lowerCamelCase		       =      ["""ml.p3.16xlarge""", """ml.p3dn.24xlarge""", """ml.p4dn.24xlarge"""]
lowerCamelCase		       =      ["""FULL_SHARD""", """SHARD_GRAD_OP""", """NO_SHARD""", """HYBRID_SHARD""", """HYBRID_SHARD_ZERO2"""]
lowerCamelCase		       =      ["""TRANSFORMER_BASED_WRAP""", """SIZE_BASED_WRAP""", """NO_WRAP"""]
lowerCamelCase		       =      ["""BACKWARD_PRE""", """BACKWARD_POST""", """NO_PREFETCH"""]
lowerCamelCase		       =      ["""FULL_STATE_DICT""", """LOCAL_STATE_DICT""", """SHARDED_STATE_DICT"""]
lowerCamelCase		       =      """2.0.1"""
lowerCamelCase		       =      ["""pdsh""", """standard""", """openmpi""", """mvapich"""]
lowerCamelCase		       =      ["""default""", """reduce-overhead""", """max-autotune"""]

lowerCamelCase		       =      {""">""": op.gt, """>=""": op.ge, """==""": op.eq, """!=""": op.ne, """<=""": op.le, """<""": op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
lowerCamelCase		       =      [
    """nnodes""",
    """nproc_per_node""",
    """rdzv_backend""",
    """rdzv_endpoint""",
    """rdzv_id""",
    """rdzv_conf""",
    """standalone""",
    """max_restarts""",
    """monitor_interval""",
    """start_method""",
    """role""",
    """module""",
    """m""",
    """no_python""",
    """run_path""",
    """log_dir""",
    """r""",
    """redirects""",
    """t""",
    """tee""",
    """node_rank""",
    """master_addr""",
    """master_port""",
]

lowerCamelCase		       =      ["""DEEPSPEED""", """MULTI_GPU""", """FSDP""", """MEGATRON_LM"""]
lowerCamelCase		       =      ["""DEEPSPEED""", """MULTI_XPU""", """FSDP"""]


