


'''simple docstring'''
import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
     import transformer_engine.pytorch as te



def 				lowerCAmelCase__	(       UpperCAmelCase    ,       UpperCAmelCase=True    ,       UpperCAmelCase=True    ,       UpperCAmelCase=True     ):






       """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(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 % 16 != 0 for p in module.weight.shape     ):
                            return
                     snake_case__					:	Optional[int]          =					module.bias is not None
                     snake_case__					:	Any          =					te.Linear(
                         module.in_features    ,       module.out_features    ,       bias=lowerCAmelCase_    ,       params_dtype=module.weight.dtype     )
                     snake_case__					:	Dict          =					module.weight.data.clone()
                     if has_bias:
                            snake_case__					:	Any          =					module.bias.data.clone()

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

                     setattr(lowerCAmelCase_    ,       lowerCAmelCase_    ,       lowerCAmelCase_     )
              elif isinstance(lowerCAmelCase_    ,       te.Linear     ) and not to_transformer_engine and _convert_linear:
                     snake_case__					:	Dict          =					module.bias is not None
                     snake_case__					:	Union[str, Any]          =					nn.Linear(
                         module.in_features    ,       module.out_features    ,       bias=lowerCAmelCase_    ,       params_dtype=module.weight.dtype     )
                     snake_case__					:	Optional[Any]          =					module.weight.data.clone()
                     if has_bias:
                            snake_case__					:	Union[str, Any]          =					module.bias.data.clone()

                     setattr(lowerCAmelCase_    ,       lowerCAmelCase_    ,       lowerCAmelCase_     )
              elif isinstance(lowerCAmelCase_    ,       te.LayerNorm     ) and not to_transformer_engine and _convert_ln:
                     snake_case__					:	Union[str, Any]          =					nn.LayerNorm(module.normalized_shape[0]    ,       eps=module.eps    ,       params_dtype=module.weight.dtype     )
                     snake_case__					:	List[str]          =					module.weight.data.clone()
                     snake_case__					:	List[Any]          =					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 				lowerCAmelCase__	(       UpperCAmelCase     ):






       """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(lowerCAmelCase_    ,       (te.LayerNorm, te.Linear)     ):
                     return True
       return False








def       lowerCAmelCase			(		lowerCAmelCase_			, lowerCAmelCase_			)->			set:
				assert isinstance(lowerCAmelCase_			, lowerCAmelCase_			), f"""The input value of [set_a={set_a}] is not a set"""
				assert isinstance(lowerCAmelCase_			, lowerCAmelCase_			), f"""The input value of [set_b={set_b}] is not a set"""

				return {a + b for a in set_a for b in set_b}


if __name__ == "__main__":
						from doctest import testmod

						testmod()