


import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
						import transformer_engine.pytorch as te


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

																		setattr(_lowerCAmelCase   ,	_lowerCAmelCase   ,	_lowerCAmelCase				)
												elif isinstance(_lowerCAmelCase   ,	nn.LayerNorm				) and to_transformer_engine and _convert_ln:
																		__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		te.LayerNorm(module.normalized_shape[0]   ,	eps=module.eps   ,	params_dtype=module.weight.dtype				)
																		__SCREAMING_SNAKE_CASE						:					Any		      =		module.weight.data.clone()
																		__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		module.bias.data.clone()

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

																		setattr(_lowerCAmelCase   ,	_lowerCAmelCase   ,	_lowerCAmelCase				)
												elif isinstance(_lowerCAmelCase   ,	te.LayerNorm				) and not to_transformer_engine and _convert_ln:
																		__SCREAMING_SNAKE_CASE						:					Tuple		      =		nn.LayerNorm(module.normalized_shape[0]   ,	eps=module.eps   ,	params_dtype=module.weight.dtype				)
																		__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		module.weight.data.clone()
																		__SCREAMING_SNAKE_CASE						:					Dict		      =		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 	_UpperCamelCase							(				lowercase__				):
						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




from .number_theory.prime_numbers import next_prime
class 	_lowercase				:






				'''simple docstring'''



				def __init__(	self      :Tuple      ,      lowerCAmelCase__      :int      ,      lowerCAmelCase__      :int | None = None      ,      lowerCAmelCase__      :float | None = None      ,      )		->      None:
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		size_table
										__SCREAMING_SNAKE_CASE						:					Dict		      =		[None] * self.size_table
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		0.75 if lim_charge is None else lim_charge
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		1 if charge_factor is None else charge_factor
										__SCREAMING_SNAKE_CASE						:					list		      =		[]
										__SCREAMING_SNAKE_CASE						:					dict		      =		{}


				def 		__magic_name__(	self      :List[Any]     )		->      Tuple:
										return self._keys


				def 		__magic_name__(	self      :Any     )		->      Dict:
										return sum(1 for slot in self.values if slot is not None     ) / (
										    self.size_table * self.charge_factor
										)


				def 		__magic_name__(	self      :Optional[Any]      ,      lowerCAmelCase__      :Union[str, Any]     )		->      Optional[Any]:
										return key % self.size_table


				def 		__magic_name__(	self      :int      ,      lowerCAmelCase__      :Union[str, Any]     )		->      List[Any]:
										print(f'''step {step_ord}'''     )
										print(list(range(len(self.values     )     )     )     )
										print(self.values     )


				def 		__magic_name__(	self      :List[str]      ,      lowerCAmelCase__      :Optional[Any]     )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					List[Any]		      =		1
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		values
										for value in values:
																self.insert_data(lowerCAmelCase__     )
																self._step_by_step(lowerCAmelCase__     )
																i += 1


				def 		__magic_name__(	self      :Tuple      ,      lowerCAmelCase__      :List[Any]      ,      lowerCAmelCase__      :Union[str, Any]     )		->      str:
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		data
										__SCREAMING_SNAKE_CASE						:					Any		      =		data


				def 		__magic_name__(	self      :Optional[int]      ,      lowerCAmelCase__      :Any      ,      lowerCAmelCase__      :List[str]=None     )		->      int:
										__SCREAMING_SNAKE_CASE						:					Tuple		      =		self.hash_function(key + 1     )

										while self.values[new_key] is not None and self.values[new_key] != key:
																if self.values.count(lowerCAmelCase__     ) > 0:
																						__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		self.hash_function(new_key + 1     )
																else:
																						__SCREAMING_SNAKE_CASE						:					int		      =		None
																						break

										return new_key


				def 		__magic_name__(	self      :Optional[int]     )		->      Union[str, Any]:
										__SCREAMING_SNAKE_CASE						:					str		      =		[value for value in self.values if value is not None]
										__SCREAMING_SNAKE_CASE						:					Optional[int]		      =		next_prime(self.size_table      ,      factor=2     )
										self._keys.clear()
										__SCREAMING_SNAKE_CASE						:					Optional[Any]		      =		[None] * self.size_table  # hell's pointers D: don't DRY ;/
										for value in survivor_values:
																self.insert_data(lowerCAmelCase__     )



				def 		__magic_name__(	self      :Tuple      ,      lowerCAmelCase__      :str     )		->      Dict:
										__SCREAMING_SNAKE_CASE						:					Union[str, Any]		      =		self.hash_function(lowerCAmelCase__     )

										if self.values[key] is None:
																self._set_value(lowerCAmelCase__      ,      lowerCAmelCase__     )

										elif self.values[key] == data:
																pass

										else:
																__SCREAMING_SNAKE_CASE						:					Tuple		      =		self._collision_resolution(lowerCAmelCase__      ,      lowerCAmelCase__     )
																if collision_resolution is not None:
																						self._set_value(lowerCAmelCase__      ,      lowerCAmelCase__     )
																else:
																						self.rehashing()
																						self.insert_data(lowerCAmelCase__     )
