






import logging
import os
import threading
import time


try:
       import warnings
except ImportError:
       __magic_name__   :  Optional[Any]        =							None

try:
       import msvcrt
except ImportError:
       __magic_name__   :  int        =							None

try:
       import fcntl
except ImportError:
       __magic_name__   :  Tuple        =							None


# Backward compatibility
# ------------------------------------------------
try:
       TimeoutError
except NameError:
       __magic_name__   :  Any        =							OSError


# Data
# ------------------------------------------------
__magic_name__   :  Union[str, Any]        =							[
    '''Timeout''',
    '''BaseFileLock''',
    '''WindowsFileLock''',
    '''UnixFileLock''',
    '''SoftFileLock''',
    '''FileLock''',
]

__magic_name__   :  Tuple        =							'''3.0.12'''


__magic_name__   :  Dict        =							None
def 		lowercase__  (				) ->			Optional[int]:
      """simple docstring"""

      global _logger
      UpperCamelCase					       =  _logger or logging.getLogger(__name__)
      return _logger





class 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  lock_file
           return None







     def __str__( self      :    List[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  f'The file lock \'{self.lock_file}\' could not be acquired.'
           return temp





class 					A__     :

     '''simple docstring'''






     def __init__( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  lock
           return None







     def __enter__( self      :    Tuple ):






           """simple docstring"""




           return self.lock







     def __exit__( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    List[Any] ):






           """simple docstring"""




           self.lock.release()
           return None





class 					A__     :

     '''simple docstring'''






     def __init__( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    int=-1   ,						_SCREAMING_SNAKE_CASE      :    Any=None ):






           """simple docstring"""




           UpperCamelCase					       =  max_filename_length if max_filename_length is not None else 255
           # Hash the filename if it's too long
           UpperCamelCase					       =  self.hash_filename_if_too_long(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
           # The path to the lock file.
           UpperCamelCase					       =  lock_file

           # The file descriptor for the *_lock_file* as it is returned by the
           # os.open() function.
           # This file lock is only NOT None, if the object currently holds the
           # lock.
           UpperCamelCase					       =  None

           # The default timeout value.
           UpperCamelCase					       =  timeout

           # We use this lock primarily for the lock counter.
           UpperCamelCase					       =  threading.Lock()

           # The lock counter is used for implementing the nested locking
           # mechanism. Whenever the lock is acquired, the counter is increased and
           # the lock is only released, when this value is 0 again.
           UpperCamelCase					       =  0
           return None







     @property
     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any] ):






           """simple docstring"""




           return self._lock_file







     @property
     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           return self._timeout







     @timeout.setter
     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  float(_SCREAMING_SNAKE_CASE )
           return None







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           raise NotImplementedError()







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           raise NotImplementedError()







     @property
     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           return self._lock_file_fd is not None







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str=None   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=0.0_5 ):






           """simple docstring"""




           if timeout is None:
                 UpperCamelCase					       =  self.timeout

           # Increment the number right at the beginning.
           # We can still undo it, if something fails.
           with self._thread_lock:
                 self._lock_counter += 1

           UpperCamelCase					       =  id(self )
           UpperCamelCase					       =  self._lock_file
           UpperCamelCase					       =  time.time()
           try:
                 while True:
                       with self._thread_lock:
                             if not self.is_locked:
                                   logger().debug(f'Attempting to acquire lock {lock_id} on {lock_filename}' )
                                   self._acquire()

                       if self.is_locked:
                             logger().debug(f'Lock {lock_id} acquired on {lock_filename}' )
                             break
                       elif timeout >= 0 and time.time() - start_time > timeout:
                             logger().debug(f'Timeout on acquiring lock {lock_id} on {lock_filename}' )
                             raise Timeout(self._lock_file )
                       else:
                             logger().debug(
                                 f'Lock {lock_id} not acquired on {lock_filename}, waiting {poll_intervall} seconds ...' )
                             time.sleep(_SCREAMING_SNAKE_CASE )
           except:  # noqa
                 # Something did go wrong, so decrement the counter.
                 with self._thread_lock:
                       UpperCamelCase					       =  max(0   ,						self._lock_counter - 1 )

                 raise
           return _Acquire_ReturnProxy(lock=self )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Dict=False ):






           """simple docstring"""




           with self._thread_lock:
                 if self.is_locked:
                       self._lock_counter -= 1

                       if self._lock_counter == 0 or force:
                             UpperCamelCase					       =  id(self )
                             UpperCamelCase					       =  self._lock_file

                             logger().debug(f'Attempting to release lock {lock_id} on {lock_filename}' )
                             self._release()
                             UpperCamelCase					       =  0
                             logger().debug(f'Lock {lock_id} released on {lock_filename}' )

           return None







     def __enter__( self      :    Union[str, Any] ):






           """simple docstring"""




           self.acquire()
           return self







     def __exit__( self      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Optional[int] ):






           """simple docstring"""




           self.release()
           return None







     def __del__( self      :    List[str] ):






           """simple docstring"""




           self.release(force=_SCREAMING_SNAKE_CASE )
           return None







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  os.path.basename(_SCREAMING_SNAKE_CASE )
           if len(_SCREAMING_SNAKE_CASE ) > max_length and max_length > 0:
                 UpperCamelCase					       =  os.path.dirname(_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  str(hash(_SCREAMING_SNAKE_CASE ) )
                 UpperCamelCase					       =  filename[: max_length - len(_SCREAMING_SNAKE_CASE ) - 8] + '...' + hashed_filename + '.lock'
                 return os.path.join(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
           else:
                 return path





class 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    int   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[str]=-1   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=None ):






           """simple docstring"""




           from .file_utils import relative_to_absolute_path

           super().__init__(_SCREAMING_SNAKE_CASE   ,						timeout=_SCREAMING_SNAKE_CASE   ,						max_filename_length=_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  '\\\\?\\' + relative_to_absolute_path(self.lock_file )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  os.O_RDWR | os.O_CREAT | os.O_TRUNC

           try:
                 UpperCamelCase					       =  os.open(self._lock_file   ,						_SCREAMING_SNAKE_CASE )
           except OSError:
                 pass
           else:
                 try:
                       msvcrt.locking(_SCREAMING_SNAKE_CASE   ,						msvcrt.LK_NBLCK   ,						1 )
                 except OSError:
                       os.close(_SCREAMING_SNAKE_CASE )
                 else:
                       UpperCamelCase					       =  fd
           return None







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self._lock_file_fd
           UpperCamelCase					       =  None
           msvcrt.locking(_SCREAMING_SNAKE_CASE   ,						msvcrt.LK_UNLCK   ,						1 )
           os.close(_SCREAMING_SNAKE_CASE )

           try:
                 os.remove(self._lock_file )
           # Probably another instance of the application
           # that acquired the file lock.
           except OSError:
                 pass
           return None





class 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=-1   ,						_SCREAMING_SNAKE_CASE      :    Tuple=None ):






           """simple docstring"""




           UpperCamelCase					       =  os.statvfs(os.path.dirname(_SCREAMING_SNAKE_CASE ) ).f_namemax
           super().__init__(_SCREAMING_SNAKE_CASE   ,						timeout=_SCREAMING_SNAKE_CASE   ,						max_filename_length=_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  os.O_RDWR | os.O_CREAT | os.O_TRUNC
           UpperCamelCase					       =  os.open(self._lock_file   ,						_SCREAMING_SNAKE_CASE )

           try:
                 fcntl.flock(_SCREAMING_SNAKE_CASE   ,						fcntl.LOCK_EX | fcntl.LOCK_NB )
           except OSError:
                 os.close(_SCREAMING_SNAKE_CASE )
           else:
                 UpperCamelCase					       =  fd
           return None







     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  self._lock_file_fd
           UpperCamelCase					       =  None
           fcntl.flock(_SCREAMING_SNAKE_CASE   ,						fcntl.LOCK_UN )
           os.close(_SCREAMING_SNAKE_CASE )
           return None





class 					A__     (     __snake_case ):

     '''simple docstring'''






     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_TRUNC
           try:
                 UpperCamelCase					       =  os.open(self._lock_file   ,						_SCREAMING_SNAKE_CASE )
           except OSError:
                 pass
           else:
                 UpperCamelCase					       =  fd
           return None







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           os.close(self._lock_file_fd )
           UpperCamelCase					       =  None

           try:
                 os.remove(self._lock_file )
           # The file is already deleted and that's what we want.
           except OSError:
                 pass
           return None

__magic_name__   :  Tuple        =							None

if msvcrt:
       __magic_name__   :  Optional[int]        =							WindowsFileLock
elif fcntl:
       __magic_name__   :  List[Any]        =							UnixFileLock
else:
       __magic_name__   :  Tuple        =							SoftFileLock

       if warnings is not None:
              warnings.warn('''only soft file lock is available''')











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


__magic_name__   :  Optional[int]        =							logging.get_logger(__name__)
def 		lowercase__  (				_UpperCamelCase) ->			Dict:
      """simple docstring"""

      UpperCamelCase					       =  r'\w+[.]\d+'
      UpperCamelCase					       =  re.findall(_UpperCamelCase       ,							_UpperCamelCase)
      for pat in pats:
            UpperCamelCase					       =  key.replace(_UpperCamelCase       ,							'_'.join(pat.split('.')))
      return key
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase) ->			Tuple:
      """simple docstring"""

      UpperCamelCase					       =  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)
      ):
            UpperCamelCase					       =  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:
            UpperCamelCase					       =  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:
            UpperCamelCase					       =  pt_tuple_key[:-1] + ('embedding',)
            return renamed_pt_tuple_key, pt_tensor

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

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

      # old PyTorch layer norm weight
      UpperCamelCase					       =  pt_tuple_key[:-1] + ('weight',)
      if pt_tuple_key[-1] == "gamma":
            return renamed_pt_tuple_key, pt_tensor

      # old PyTorch layer norm bias
      UpperCamelCase					       =  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 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase=42) ->			Optional[Any]:
      """simple docstring"""

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

      # Step 2: Since the model is stateless, get random Flax params
      UpperCamelCase					       =  flax_model.init_weights(PRNGKey(_UpperCamelCase))

      UpperCamelCase					       =  flatten_dict(_UpperCamelCase)
      UpperCamelCase					       =  {}

      # Need to change some parameters name to match Flax names
      for pt_key, pt_tensor in pt_state_dict.items():
            UpperCamelCase					       =  rename_key(_UpperCamelCase)
            UpperCamelCase					       =  tuple(renamed_pt_key.split('.'))

            # Correctly rename weight parameters
            UpperCamelCase       ,  UpperCamelCase					       =  rename_key_and_reshape_tensor(_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase)

            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
            UpperCamelCase					       =  jnp.asarray(_UpperCamelCase)

      return unflatten_dict(_UpperCamelCase)



