



import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
	from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
	import torch_xla.core.xla_model as xm



def    snake_case_				(    _SCREAMING_SNAKE_CASE   ):
	if is_torch_version("<"  ,	"2.0.0"   ) or not hasattr(_SCREAMING_SNAKE_CASE  ,	"_dynamo"   ):
		return False
	return isinstance(_SCREAMING_SNAKE_CASE  ,	torch._dynamo.eval_frame.OptimizedModule   )



def    snake_case_				(    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE = True   ):
	__lowercase		       =			(torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

	__lowercase		       =			is_compiled_module(_SCREAMING_SNAKE_CASE   )
	if is_compiled:
		__lowercase		       =			model
		__lowercase		       =			model._orig_mod

	if is_deepspeed_available():
		options += (DeepSpeedEngine,)

	while isinstance(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   ):
		__lowercase		       =			model.module

	if not keep_fpaa_wrapper:
		__lowercase		       =			getattr(_SCREAMING_SNAKE_CASE  ,	"forward"   )
		__lowercase		       =			model.__dict__.pop("_original_forward"  ,	_SCREAMING_SNAKE_CASE   )
		if original_forward is not None:
			while hasattr(_SCREAMING_SNAKE_CASE  ,	"__wrapped__"   ):
				__lowercase		       =			forward.__wrapped__
				if forward == original_forward:
					break
			__lowercase		       =			forward
		if getattr(_SCREAMING_SNAKE_CASE  ,	"_converted_to_transformer_engine"  ,	_SCREAMING_SNAKE_CASE   ):
			convert_model(_SCREAMING_SNAKE_CASE  ,	to_transformer_engine=_SCREAMING_SNAKE_CASE   )

	if is_compiled:
		__lowercase		       =			model
		__lowercase		       =			compiled_model

	return model



def    snake_case_				(    ):
	PartialState().wait_for_everyone()



def    snake_case_				(    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   ):
	if PartialState().distributed_type == DistributedType.TPU:
		xm.save(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
	elif PartialState().local_process_index == 0:
		torch.save(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )



@contextmanager
def    snake_case_				(    **_SCREAMING_SNAKE_CASE   ):
	for key, value in kwargs.items():
		__lowercase		       =			str(_SCREAMING_SNAKE_CASE   )

	yield

	for key in kwargs:
		if key.upper() in os.environ:
			del os.environ[key.upper()]



def    snake_case_				(    _SCREAMING_SNAKE_CASE   ):
	if not hasattr(_SCREAMING_SNAKE_CASE  ,	"__qualname__"   ) and not hasattr(_SCREAMING_SNAKE_CASE  ,	"__name__"   ):
		__lowercase		       =			getattr(_SCREAMING_SNAKE_CASE  ,	"__class__"  ,	_SCREAMING_SNAKE_CASE   )
	if hasattr(_SCREAMING_SNAKE_CASE  ,	"__qualname__"   ):
		return obj.__qualname__
	if hasattr(_SCREAMING_SNAKE_CASE  ,	"__name__"   ):
		return obj.__name__
	return str(_SCREAMING_SNAKE_CASE   )



def    snake_case_				(    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   ):
	for key, value in source.items():
		if isinstance(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   ):
			__lowercase		       =			destination.setdefault(_SCREAMING_SNAKE_CASE  ,	{}   )
			merge_dicts(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
		else:
			__lowercase		       =			value

	return destination



def    snake_case_				(    _SCREAMING_SNAKE_CASE = None   ):
	if port is None:
		__lowercase		       =			2_9_5_0_0
	with socket.socket(socket.AF_INET  ,	socket.SOCK_STREAM   ) as s:
		return s.connect_ex(("localhost", port)   ) == 0








from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

snake_case__					:     Union[str, Any]    		=   TypeVar("""T""")
snake_case__					:     Optional[int]    		=   TypeVar("""U""")

class   _A   ( Generic[T, U]	):


						'''simple docstring'''





						def __init__(	self	:							Optional[int]							,      lowerCamelCase	:							T | None							,      lowerCamelCase	:							U | None	):






							'''simple docstring'''



							__lowercase		       =			key
							__lowercase		       =			val
							__lowercase		       =			None
							__lowercase		       =			None



						def __repr__(	self	:							Any	):






							'''simple docstring'''



							return (
							    f"""Node: key: {self.key}, val: {self.val}, """
							    f"""has next: {bool(self.next	)}, has prev: {bool(self.prev	)}"""
							)

class   _A   ( Generic[T, U]	):


						'''simple docstring'''





						def __init__(	self	:							Dict	):






							'''simple docstring'''



							__lowercase		       =			DoubleLinkedListNode(lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			DoubleLinkedListNode(lowerCamelCase							,      lowerCamelCase	)
							__lowercase       ,    __lowercase		       =			self.rear, self.head



						def __repr__(	self	:							Optional[Any]	):






							'''simple docstring'''



							__lowercase		       =			["DoubleLinkedList"]
							__lowercase		       =			self.head
							while node.next is not None:
								rep.append(str(lowerCamelCase	)	)
								__lowercase		       =			node.next
							rep.append(str(self.rear	)	)
							return ",\n    ".join(lowerCamelCase	)



						def     _snake_case	(	self	:							Union[str, Any]							,      lowerCamelCase	:							DoubleLinkedListNode[T, U]	):






							'''simple docstring'''



							__lowercase		       =			self.rear.prev

							# All nodes other than self.head are guaranteed to have non-None previous
							assert previous is not None

							__lowercase		       =			node
							__lowercase		       =			previous
							__lowercase		       =			node
							__lowercase		       =			self.rear



						def     _snake_case	(	self	:							Optional[int]							,      lowerCamelCase	:							DoubleLinkedListNode[T, U]	):






							'''simple docstring'''



							if node.prev is None or node.next is None:
								return None

							__lowercase		       =			node.next
							__lowercase		       =			node.prev
							__lowercase		       =			None
							__lowercase		       =			None
							return node

class   _A   ( Generic[T, U]	):


						'''simple docstring'''





						_snake_case			:     dict[Callable[[T], U], LRUCache[T, U]]      							=      {}



						def __init__(	self	:							List[Any]							,      lowerCamelCase	:							int	):






							'''simple docstring'''



							__lowercase		       =			DoubleLinkedList()
							__lowercase		       =			capacity
							__lowercase		       =			0
							__lowercase		       =			0
							__lowercase		       =			0
							__lowercase		       =			{}



						def __repr__(	self	:							Optional[Any]	):






							'''simple docstring'''



							return (
							    f"""CacheInfo(hits={self.hits}, misses={self.miss}, """
							    f"""capacity={self.capacity}, current size={self.num_keys})"""
							)



						def __contains__(	self	:							Dict							,      lowerCamelCase	:							T	):






							'''simple docstring'''



							return key in self.cache



						def     _snake_case	(	self	:							List[Any]							,      lowerCamelCase	:							T	):






							'''simple docstring'''



							if key in self.cache:
								self.hits += 1
								__lowercase		       =			self.cache[key]
								__lowercase		       =			self.list.remove(self.cache[key]	)
								assert node == value_node

								# node is guaranteed not None because it is in self.cache
								assert node is not None
								self.list.add(lowerCamelCase	)
								return node.val
							self.miss += 1
							return None



						def     _snake_case	(	self	:							Union[str, Any]							,      lowerCamelCase	:							T							,      lowerCamelCase	:							U	):






							'''simple docstring'''



							if key not in self.cache:
								if self.num_keys >= self.capacity:
									# delete first node (oldest) when over capacity
									__lowercase		       =			self.list.head.next

									# guaranteed to have a non-None first node when num_keys > 0
									# explain to type checker via assertions
									assert first_node is not None
									assert first_node.key is not None
									assert (
									    self.list.remove(lowerCamelCase	) is not None
									)  # node guaranteed to be in list assert node.key is not None

									del self.cache[first_node.key]
									self.num_keys -= 1
								__lowercase		       =			DoubleLinkedListNode(lowerCamelCase							,      lowerCamelCase	)
								self.list.add(self.cache[key]	)
								self.num_keys += 1

							else:
								# bump node to the end of the list, update value
								__lowercase		       =			self.list.remove(self.cache[key]	)
								assert node is not None  # node guaranteed to be in list
								__lowercase		       =			value
								self.list.add(lowerCamelCase	)



						@classmethod
						def     _snake_case	(	cls	:							Union[str, Any]							,      lowerCamelCase	:							int = 128	):






							'''simple docstring'''



							def cache_decorator_inner(lowerCamelCase	:							Callable[[T], U]	) -> Callable[..., U]:
								def cache_decorator_wrapper(*lowerCamelCase	:							T	) -> U:
									if func not in cls.decorator_function_to_instance_map:
										__lowercase		       =			LRUCache(lowerCamelCase	)

									__lowercase		       =			cls.decorator_function_to_instance_map[func].get(args[0]	)
									if result is None:
										__lowercase		       =			func(*lowerCamelCase	)
										cls.decorator_function_to_instance_map[func].put(args[0]							,      lowerCamelCase	)
									return result

								def cache_info() -> LRUCache[T, U]:
									return cls.decorator_function_to_instance_map[func]

								setattr(lowerCamelCase							,      "cache_info"							,      lowerCamelCase	)  # noqa: B010

								return cache_decorator_wrapper

							return cache_decorator_inner


if __name__ == "__main__":
	import doctest

	doctest.testmod()



