


'''simple docstring'''







import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
  import transformer_engine.pytorch as te







def _UpperCamelCase			(	__A ,     __A=True ,     __A=True ,     __A=True       )				->						List[Any]:


      '''simple docstring'''



      if not is_fpa_available():
            raise ImportError("Using `convert_model` requires transformer_engine to be installed."       )
      for name, module in model.named_children():
            if isinstance(__A ,     nn.Linear       ) and to_transformer_engine and _convert_linear:
                  # Return early if the linear layer weights are not multiples of 16
                  if any(p % 16 != 0 for p in module.weight.shape       ):
                        return
                  UpperCamelCase__      =							module.bias is not None
                  UpperCamelCase__      =							te.Linear(
                      module.in_features ,     module.out_features ,     bias=__A ,     params_dtype=module.weight.dtype       )
                  UpperCamelCase__      =							module.weight.data.clone()
                  if has_bias:
                        UpperCamelCase__      =							module.bias.data.clone()

                  setattr(__A ,     __A ,     __A       )
            elif isinstance(__A ,     nn.LayerNorm       ) and to_transformer_engine and _convert_ln:
                  UpperCamelCase__      =							te.LayerNorm(module.normalized_shape[0] ,     eps=module.eps ,     params_dtype=module.weight.dtype       )
                  UpperCamelCase__      =							module.weight.data.clone()
                  UpperCamelCase__      =							module.bias.data.clone()

                  setattr(__A ,     __A ,     __A       )
            elif isinstance(__A ,     te.Linear       ) and not to_transformer_engine and _convert_linear:
                  UpperCamelCase__      =							module.bias is not None
                  UpperCamelCase__      =							nn.Linear(
                      module.in_features ,     module.out_features ,     bias=__A ,     params_dtype=module.weight.dtype       )
                  UpperCamelCase__      =							module.weight.data.clone()
                  if has_bias:
                        UpperCamelCase__      =							module.bias.data.clone()

                  setattr(__A ,     __A ,     __A       )
            elif isinstance(__A ,     te.LayerNorm       ) and not to_transformer_engine and _convert_ln:
                  UpperCamelCase__      =							nn.LayerNorm(module.normalized_shape[0] ,     eps=module.eps ,     params_dtype=module.weight.dtype       )
                  UpperCamelCase__      =							module.weight.data.clone()
                  UpperCamelCase__      =							module.bias.data.clone()

                  setattr(__A ,     __A ,     __A       )
            else:
                  convert_model(
                      __A ,     to_transformer_engine=__A ,     _convert_linear=_convert_linear ,     _convert_ln=_convert_ln ,     )







def _UpperCamelCase			(	__A       )				->						Optional[int]:


      '''simple docstring'''



      if not is_fpa_available():
            raise ImportError("Using `has_transformer_engine_layers` requires transformer_engine to be installed."       )
      for m in model.modules():
            if isinstance(__A ,     (te.LayerNorm, te.Linear)       ):
                  return True
      return False









'''simple docstring'''







from __future__ import annotations

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

a__    :   Dict								=			TypeVar('T')
a__    :   Optional[Any]								=			TypeVar('U')
class    lowercase_	(       Generic[T, U]  ):


    def __init__(	self    ,   a    ,   a       ):
          UpperCamelCase__      =							key
          UpperCamelCase__      =							val
          UpperCamelCase__      =							0
          UpperCamelCase__      =							None
          UpperCamelCase__      =							None


    def __repr__(	self       ):
          return "Node: key: {}, val: {}, freq: {}, has next: {}, has prev: {}".format(
              self.key    ,   self.val    ,   self.freq    ,   self.next is not None    ,   self.prev is not None       )

class    lowercase_	(       Generic[T, U]  ):


    def __init__(	self       ):
          UpperCamelCase__      =							DoubleLinkedListNode(a    ,   a       )
          UpperCamelCase__      =							DoubleLinkedListNode(a    ,   a       )
          UpperCamelCase__   ,      UpperCamelCase__      =							self.rear, self.head


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


    def   __a				(	self    ,   a       ):
          UpperCamelCase__      =							self.rear.prev

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

          UpperCamelCase__      =							node
          UpperCamelCase__      =							previous
          UpperCamelCase__      =							node
          UpperCamelCase__      =							self.rear
          node.freq += 1
          self._position_node(a       )


    def   __a				(	self    ,   a       ):

          while node.prev is not None and node.prev.freq > node.freq:
                # swap node with previous node
                UpperCamelCase__      =							node.prev

                UpperCamelCase__      =							previous_node.prev
                UpperCamelCase__      =							node.prev
                UpperCamelCase__      =							previous_node
                UpperCamelCase__      =							node


    def   __a				(	self    ,   a       ):

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

          UpperCamelCase__      =							node.next
          UpperCamelCase__      =							node.prev
          UpperCamelCase__      =							None
          UpperCamelCase__      =							None
          return node

class    lowercase_	(       Generic[T, U]  ):
    __UpperCAmelCase       =      {}


    def __init__(	self    ,   a       ):
          UpperCamelCase__      =							DoubleLinkedList()
          UpperCamelCase__      =							capacity
          UpperCamelCase__      =							0
          UpperCamelCase__      =							0
          UpperCamelCase__      =							0
          UpperCamelCase__      =							{}


    def __repr__(	self       ):
          return (
              f'''CacheInfo(hits={self.hits}, misses={self.miss}, '''
              f'''capacity={self.capacity}, current_size={self.num_keys})'''
          )


    def __contains__(	self    ,   a       ):
          return key in self.cache


    def   __a				(	self    ,   a       ):

          if key in self.cache:
                self.hits += 1
                UpperCamelCase__      =							self.cache[key]
                UpperCamelCase__      =							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(a       )
                return node.val
          self.miss += 1
          return None


    def   __a				(	self    ,   a    ,   a       ):

          if key not in self.cache:
                if self.num_keys >= self.capacity:
                      # delete first node when over capacity
                      UpperCamelCase__      =							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(a       ) is not None
                      # first_node guaranteed to be in list

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

          else:
                UpperCamelCase__      =							self.list.remove(self.cache[key]       )
                assert node is not None  # node guaranteed to be in list
                UpperCamelCase__      =							value
                self.list.add(a       )


    @classmethod
    def   __a				(	cls    ,   a = 1_28       ):

          def cache_decorator_inner(a       ) -> Callable[..., U]:
                def cache_decorator_wrapper(*a       ) -> U:
                      if func not in cls.decorator_function_to_instance_map:
                            UpperCamelCase__      =							LFUCache(a       )

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

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

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

                return cache_decorator_wrapper

          return cache_decorator_inner


if __name__ == "__main__":
  import doctest

  doctest.testmod()





