
'''simple docstring'''
import asyncio
import os
import shutil
import subprocess
import sys
import tempfile
import unittest
from distutils.util import strtobool
from functools import partial
from pathlib import Path
from typing import List, Union
from unittest import mock

import torch

from ..state import AcceleratorState, PartialState
from ..utils import (
    gather,
    is_bnb_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_mps_available,
    is_safetensors_available,
    is_tensorboard_available,
    is_torch_version,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE=False       ):
    """simple docstring"""






    try:
        snake_case_ :    Optional[Any]     			=						os.environ[key]
    except KeyError:
        # KEY isn't set, default to `default`.
        snake_case_ :    List[str]     			=						default
    else:
        # KEY is set, convert it to True or False.
        try:
            snake_case_ :    Dict     			=						strtobool(__SCREAMING_SNAKE_CASE       )
        except ValueError:
            # More values are supported, but let's keep the message simple.
            raise ValueError(f'If set, {key} must be yes or no.'       )
    return _value


a_											=						parse_flag_from_env("RUN_SLOW", default=False)

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skip("Test was skipped"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(_run_slow_tests,		"test is slow"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(not torch.cuda.is_available(),		"test requires only a CPU"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(torch.cuda.is_available(),		"test requires a GPU"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_xpu_available(),		"test requires a XPU"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_mps_available(),		"test requires a `mps` backend support in `torch`"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(
        is_transformers_available() and is_datasets_available(),		"test requires the Hugging Face suite"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_bnb_available(),		"test requires the bitsandbytes library"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_tpu_available(),		"test requires TPU"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(torch.cuda.device_count() == 1,		"test requires a GPU"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(torch.xpu.device_count() == 1,		"test requires a XPU"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(torch.cuda.device_count() > 1,		"test requires multiple GPUs"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(torch.xpu.device_count() > 1,		"test requires multiple XPUs"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_safetensors_available(),		"test requires safetensors"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_deepspeed_available(),		"test requires DeepSpeed"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_torch_version(">=",		"1.12.0"       ),		"test requires torch version >= 1.12.0"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE=None,		__SCREAMING_SNAKE_CASE=None       ):
    """simple docstring"""






    if test_case is None:
        return partial(__SCREAMING_SNAKE_CASE,		version=__SCREAMING_SNAKE_CASE       )
    return unittest.skipUnless(is_torch_version(">=",		__SCREAMING_SNAKE_CASE       ),		f'test requires torch version >= {version}'       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_tensorboard_available(),		"test requires Tensorboard"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_wandb_available(),		"test requires wandb"       )(__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(is_comet_ml_available(),		"test requires comet_ml"       )(__SCREAMING_SNAKE_CASE       )


a_											=						(
    any([is_wandb_available(), is_tensorboard_available()]) and not is_comet_ml_available()
)

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    return unittest.skipUnless(
        _atleast_one_tracker_available,		"test requires at least one tracker to be available and for `comet_ml` to not be installed",		)(__SCREAMING_SNAKE_CASE       )






class    UpperCAmelCase_				(	unittest.TestCase					):
 UpperCAmelCase_		    =						True





 @classmethod
 def        snake_case__		(    cls):
     snake_case_ :    Optional[Any]     			=						tempfile.mkdtemp()





 @classmethod
 def        snake_case__		(    cls):
     if os.path.exists(cls.tmpdir):
         shutil.rmtree(cls.tmpdir)





 def        snake_case__		(    self):
     if self.clear_on_setup:
         for path in Path(self.tmpdir).glob("**/*"):
             if path.is_file():
                 path.unlink()
             elif path.is_dir():
                 shutil.rmtree(lowercase_)






class    UpperCAmelCase_				(	unittest.TestCase					):





 def        snake_case__		(    self):
     super().tearDown()
     # Reset the state of the AcceleratorState singleton.
     AcceleratorState._reset_state()
     PartialState._reset_state()






class    UpperCAmelCase_				(	unittest.TestCase					):





 def        snake_case__		(    self ,       lowercase_):
     snake_case_ :    Dict     			=						mocks if isinstance(lowercase_ ,       (tuple, list)) else [mocks]
     for m in self.mocks:
         m.start()
         self.addCleanup(m.stop)


def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    snake_case_ :    Any     			=						AcceleratorState()
    snake_case_ :    Optional[int]     			=						tensor[None].clone().to(state.device       )
    snake_case_ :    Any     			=						gather(__SCREAMING_SNAKE_CASE       ).cpu()
    snake_case_ :    List[str]     			=						tensor[0].cpu()
    for i in range(tensors.shape[0]       ):
        if not torch.equal(tensors[i],		__SCREAMING_SNAKE_CASE       ):
            return False
    return True






class    UpperCAmelCase_				:





 def __init__(    self ,       lowercase_ ,       lowercase_ ,       lowercase_):
     snake_case_ :    Optional[Any]     			=						returncode
     snake_case_ :    List[Any]     			=						stdout
     snake_case_ :    Tuple     			=						stderr


async def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    while True:
        snake_case_ :    List[str]     			=						await stream.readline()
        if line:
            callback(__SCREAMING_SNAKE_CASE       )
        else:
            break

async def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE=None,		__SCREAMING_SNAKE_CASE=None,		__SCREAMING_SNAKE_CASE=None,		__SCREAMING_SNAKE_CASE=False,		__SCREAMING_SNAKE_CASE=False       ):
    """simple docstring"""






    if echo:
        print("\nRunning: ",		" ".join(__SCREAMING_SNAKE_CASE       )       )

    snake_case_ :    List[Any]     			=						await asyncio.create_subprocess_exec(
        cmd[0],		*cmd[1:],		stdin=__SCREAMING_SNAKE_CASE,		stdout=asyncio.subprocess.PIPE,		stderr=asyncio.subprocess.PIPE,		env=__SCREAMING_SNAKE_CASE,		)

    # note: there is a warning for a possible deadlock when using `wait` with huge amounts of data in the pipe
    # https://docs.python.org/3/library/asyncio-subprocess.html#asyncio.asyncio.subprocess.Process.wait
    #
    # If it starts hanging, will need to switch to the following code. The problem is that no data
    # will be seen until it's done and if it hangs for example there will be no debug info.
    # out, err = await p.communicate()
    # return _RunOutput(p.returncode, out, err)

    snake_case_ :    Any     			=						[]
    snake_case_ :    Any     			=						[]

    def tee(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE=""       ):
        snake_case_ :    int     			=						line.decode("utf-8"       ).rstrip()
        sink.append(__SCREAMING_SNAKE_CASE       )
        if not quiet:
            print(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		file=__SCREAMING_SNAKE_CASE       )

    # XXX: the timeout doesn't seem to make any difference here
    await asyncio.wait(
        [
            asyncio.create_task(_read_stream(p.stdout,		lambda __SCREAMING_SNAKE_CASE       : tee(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		sys.stdout,		label="stdout:"       )       )       ),
            asyncio.create_task(_read_stream(p.stderr,		lambda __SCREAMING_SNAKE_CASE       : tee(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		sys.stderr,		label="stderr:"       )       )       ),
        ],		timeout=__SCREAMING_SNAKE_CASE,		)
    return _RunOutput(await p.wait(),		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE=None,		__SCREAMING_SNAKE_CASE=None,		__SCREAMING_SNAKE_CASE=1_8_0,		__SCREAMING_SNAKE_CASE=False,		__SCREAMING_SNAKE_CASE=True       ):
    """simple docstring"""






    snake_case_ :    Optional[int]     			=						asyncio.get_event_loop()
    snake_case_ :    Union[str, Any]     			=						loop.run_until_complete(
        _stream_subprocess(__SCREAMING_SNAKE_CASE,		env=__SCREAMING_SNAKE_CASE,		stdin=__SCREAMING_SNAKE_CASE,		timeout=__SCREAMING_SNAKE_CASE,		quiet=__SCREAMING_SNAKE_CASE,		echo=__SCREAMING_SNAKE_CASE       )       )

    snake_case_ :    Optional[Any]     			=						" ".join(__SCREAMING_SNAKE_CASE       )
    if result.returncode > 0:
        snake_case_ :    List[Any]     			=						"\n".join(result.stderr       )
        raise RuntimeError(
            f'\'{cmd_str}\' failed with returncode {result.returncode}\n\n'
            f'The combined stderr from workers follows:\n{stderr}'       )

    return result






class    UpperCAmelCase_				(	snake_case__					):
 pass


def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE=False       ):
    """simple docstring"""






    try:
        snake_case_ :    Dict     			=						subprocess.check_output(__SCREAMING_SNAKE_CASE,		stderr=subprocess.STDOUT       )
        if return_stdout:
            if hasattr(__SCREAMING_SNAKE_CASE,		"decode"       ):
                snake_case_ :    int     			=						output.decode("utf-8"       )
            return output
    except subprocess.CalledProcessError as e:
        raise SubprocessCallException(
            f'Command `{" ".join(__SCREAMING_SNAKE_CASE       )}` failed with the following error:\n\n{e.output.decode()}'       ) from e






'''simple docstring'''
import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


a_											=						logging.get_logger(__name__)

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    snake_case_ :    List[Any]     			=						r"\w+[.]\d+"
    snake_case_ :    Tuple     			=						re.findall(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )
    for pat in pats:
        snake_case_ :    Union[str, Any]     			=						key.replace(__SCREAMING_SNAKE_CASE,		"_".join(pat.split("."       )       )       )
    return key

def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       ):
    """simple docstring"""






    snake_case_ :    Tuple     			=						pt_tuple_key[:-1] + ("scale",)
    if (
        any("norm" in str_ for str_ in pt_tuple_key       )
        and (pt_tuple_key[-1] == "bias")
        and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
        and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
    ):
        snake_case_ :    int     			=						pt_tuple_key[:-1] + ("scale",)
        return renamed_pt_tuple_key, pt_tensor
    elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
        snake_case_ :    Dict     			=						pt_tuple_key[:-1] + ("scale",)
        return renamed_pt_tuple_key, pt_tensor

    # embedding
    if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
        snake_case_ :    List[Any]     			=						pt_tuple_key[:-1] + ("embedding",)
        return renamed_pt_tuple_key, pt_tensor

    # conv layer
    snake_case_ :    List[str]     			=						pt_tuple_key[:-1] + ("kernel",)
    if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
        snake_case_ :    str     			=						pt_tensor.transpose(2,		3,		1,		0       )
        return renamed_pt_tuple_key, pt_tensor

    # linear layer
    snake_case_ :    str     			=						pt_tuple_key[:-1] + ("kernel",)
    if pt_tuple_key[-1] == "weight":
        snake_case_ :    Any     			=						pt_tensor.T
        return renamed_pt_tuple_key, pt_tensor

    # old PyTorch layer norm weight
    snake_case_ :    Union[str, Any]     			=						pt_tuple_key[:-1] + ("weight",)
    if pt_tuple_key[-1] == "gamma":
        return renamed_pt_tuple_key, pt_tensor

    # old PyTorch layer norm bias
    snake_case_ :    Union[str, Any]     			=						pt_tuple_key[:-1] + ("bias",)
    if pt_tuple_key[-1] == "beta":
        return renamed_pt_tuple_key, pt_tensor

    return pt_tuple_key, pt_tensor


def 						UpperCamelCase_     (	__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE=4_2       ):
    """simple docstring"""






    snake_case_ :    int     			=						{k: v.numpy() for k, v in pt_state_dict.items()}

    # Step 2: Since the model is stateless, get random Flax params
    snake_case_ :    Optional[int]     			=						flax_model.init_weights(PRNGKey(__SCREAMING_SNAKE_CASE       )       )

    snake_case_ :    Union[str, Any]     			=						flatten_dict(__SCREAMING_SNAKE_CASE       )
    snake_case_ :    str     			=						{}

    # Need to change some parameters name to match Flax names
    for pt_key, pt_tensor in pt_state_dict.items():
        snake_case_ :    List[str]     			=						rename_key(__SCREAMING_SNAKE_CASE       )
        snake_case_ :    Tuple     			=						tuple(renamed_pt_key.split("."       )       )

        # Correctly rename weight parameters
        snake_case_			, snake_case_ :    Dict     			=						rename_key_and_reshape_tensor(__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       )

        if flax_key in random_flax_state_dict:
            if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                raise ValueError(
                    f'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
                    f'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'       )

        # also add unexpected weight so that warning is thrown
        snake_case_ :    Any     			=						jnp.asarray(__SCREAMING_SNAKE_CASE       )

    return unflatten_dict(__SCREAMING_SNAKE_CASE       )




