
import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
		import transformer_engine.pytorch as te




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  Union[str, Any]		,      SCREAMING_SNAKE_CASE_  :  Optional[int]=True		,      SCREAMING_SNAKE_CASE_  :  List[str]=True		,      SCREAMING_SNAKE_CASE_  :  int=True					)     ->    Dict:
						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(SCREAMING_SNAKE_CASE_		,      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
																		_lowercase									=				module.bias is not None
																		_lowercase									=				te.Linear(
																		    module.in_features		,      module.out_features		,      bias=SCREAMING_SNAKE_CASE_		,      params_dtype=module.weight.dtype					)
																		_lowercase									=				module.weight.data.clone()
																		if has_bias:
																								_lowercase									=				module.bias.data.clone()

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

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

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

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




def   UpperCamelCase__     (    SCREAMING_SNAKE_CASE_  :  str					)     ->    Optional[int]:
						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(SCREAMING_SNAKE_CASE_		,      (te.LayerNorm, te.Linear)					):
																		return True
						return False

import abc
import copy
import dataclasses
from dataclasses import dataclass
from typing import ClassVar, Dict, Type, TypeVar

from ..features import Features


A		:   Optional[int]         =    TypeVar('''T''', bound='''TaskTemplate''')



@dataclass(frozen=_a						)
class       a_   ( abc.ABC						):
			# `task` is not a ClassVar since we want it to be part of the `asdict` output for JSON serialization
			a  : str
			a  : ClassVar[Features]
			a  : ClassVar[Features]






			def       UpperCamelCase_							(			self							, __UpperCamelCase     ):
									return copy.deepcopy(self     )






			@property
			def       UpperCamelCase_							(			self     ):
									return Features(**self.input_schema							, **self.label_schema     )






			@property
			@abc.abstractmethod
			def       UpperCamelCase_							(			self     ):
									raise NotImplementedError






			@classmethod
			def       UpperCamelCase_							(			cls							, __UpperCamelCase     ):
									_lowercase									=				{f.name for f in dataclasses.fields(cls     )}
									return cls(**{k: v for k, v in template_dict.items() if k in field_names}     )