





from __future__ import annotations

from enum import Enum, unique
from typing import TypeVar

# Create a generic variable that can be 'Enum', or any subclass.
_A				:							Dict							= TypeVar("""T""", bound="""Enum""")



@unique
class 				__snake_case  (     __SCREAMING_SNAKE_CASE						):


    '''simple docstring'''







    lowerCamelCase__						:		Optional[Any]							=  8_0
    lowerCamelCase__						:		Optional[int]							=  7_0
    lowerCamelCase__						:		List[str]							=  6_0
    lowerCamelCase__						:		int							=  5_0
    lowerCamelCase__						:		int							=  4_0
    lowerCamelCase__						:		List[Any]							=  3_0
    lowerCamelCase__						:		Optional[Any]							=  2_0
    lowerCamelCase__						:		List[Any]							=  1_0





@unique
class 				__snake_case  (     __SCREAMING_SNAKE_CASE						):


    '''simple docstring'''







    lowerCamelCase__						:		Optional[int]							=  2_4
    lowerCamelCase__						:		Dict							=  2_1
    lowerCamelCase__						:		int							=  1_8
    lowerCamelCase__						:		Dict							=  1_5
    lowerCamelCase__						:		List[Any]							=  1_2
    lowerCamelCase__						:		Dict							=  9
    lowerCamelCase__						:		Optional[int]							=  6
    lowerCamelCase__						:		Any							=  3
    lowerCamelCase__						:		Any							=  2
    lowerCamelCase__						:		Optional[Any]							=  1
    lowerCamelCase__						:		Union[str, Any]							=  -1
    lowerCamelCase__						:		Any							=  -2
    lowerCamelCase__						:		Any							=  -3
    lowerCamelCase__						:		Any							=  -6
    lowerCamelCase__						:		Optional[Any]							=  -9
    lowerCamelCase__						:		int							=  -1_2
    lowerCamelCase__						:		str							=  -1_5
    lowerCamelCase__						:		List[Any]							=  -1_8
    lowerCamelCase__						:		Optional[int]							=  -2_1
    lowerCamelCase__						:		str							=  -2_4



    @classmethod
    def    lowercase_				(					cls    ):





        '''simple docstring'''

        return {unit.name: unit.value for unit in cls if unit.value > 0}




    @classmethod
    def    lowercase_				(					cls    ):





        '''simple docstring'''

        return {unit.name: unit.value for unit in cls if unit.value < 0}





def        __snake_case       (						lowerCAmelCase_  )							->	str:
    SCREAMING_SNAKE_CASE__       							= SIUnit.get_positive() if value > 0 else SIUnit.get_negative()
    for name_prefix, value_prefix in prefixes.items():
        SCREAMING_SNAKE_CASE__       							= value / (1_0**value_prefix)
        if numerical_part > 1:
            return f'''{numerical_part!s} {name_prefix}'''
    return str(lowerCAmelCase_  )


def        __snake_case       (						lowerCAmelCase_  )							->	str:
    for prefix in BinaryUnit:
        SCREAMING_SNAKE_CASE__       							= value / (2**prefix.value)
        if numerical_part > 1:
            return f'''{numerical_part!s} {prefix.name}'''
    return str(lowerCAmelCase_  )


if __name__ == "__main__":
  import doctest

  doctest.testmod()







import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
  import transformer_engine.pytorch as te


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True  )							->	Tuple:
    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(lowerCAmelCase_   ,       nn.Linear  ) and to_transformer_engine and _convert_linear:
            # Return early if the linear layer weights are not multiples of 16
            if any(p % 1_6 != 0 for p in module.weight.shape  ):
                return
            SCREAMING_SNAKE_CASE__       							= module.bias is not None
            SCREAMING_SNAKE_CASE__       							= te.Linear(
                module.in_features   ,       module.out_features   ,       bias=lowerCAmelCase_   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            if has_bias:
                SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

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

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

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

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



def        __snake_case       (						lowerCAmelCase_  )							->	Union[str, Any]:
    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(lowerCAmelCase_   ,       (te.LayerNorm, te.Linear)  ):
            return True
    return False
