



import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
	import transformer_engine.pytorch as te



def 				_lowerCAmelCase      (_lowerCAmelCase				,		_lowerCAmelCase=True				,		_lowerCAmelCase=True				,		_lowerCAmelCase=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
															UpperCamelCase_											=				module.bias is not None
															UpperCamelCase_											=				te.Linear(
															    module.in_features				,		module.out_features				,		bias=_lowerCAmelCase				,		params_dtype=module.weight.dtype)
															UpperCamelCase_											=				module.weight.data.clone()
															if has_bias:
																				UpperCamelCase_											=				module.bias.data.clone()

															setattr(_lowerCAmelCase				,		_lowerCAmelCase				,		_lowerCAmelCase)
										elif isinstance(_lowerCAmelCase				,		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(_lowerCAmelCase				,		_lowerCAmelCase				,		_lowerCAmelCase)
										elif isinstance(_lowerCAmelCase				,		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=_lowerCAmelCase				,		params_dtype=module.weight.dtype)
															UpperCamelCase_											=				module.weight.data.clone()
															if has_bias:
																				UpperCamelCase_											=				module.bias.data.clone()

															setattr(_lowerCAmelCase				,		_lowerCAmelCase				,		_lowerCAmelCase)
										elif isinstance(_lowerCAmelCase				,		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(_lowerCAmelCase				,		_lowerCAmelCase				,		_lowerCAmelCase)
										else:
															convert_model(
															    _lowerCAmelCase				,		to_transformer_engine=_lowerCAmelCase				,		_convert_linear=_convert_linear				,		_convert_ln=_convert_ln				,		)






def 				_lowerCAmelCase      (_lowerCAmelCase):
					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 transformers import HfArgumentParser, PyTorchBenchmark, PyTorchBenchmarkArguments



def 				_lowerCAmelCase      ():
					UpperCamelCase_											=				HfArgumentParser(_lowerCAmelCase)
					try:
										UpperCamelCase_											=				parser.parse_args_into_dataclasses()[0]
					except ValueError as e:
										UpperCamelCase_											=				"Arg --no_{0} is no longer used, please use --no-{0} instead."
										UpperCamelCase_											=				" ".join(str(_lowerCAmelCase).split(" ")[:-1])
										UpperCamelCase_											=				""
										UpperCamelCase_											=				eval(str(_lowerCAmelCase).split(" ")[-1])
										UpperCamelCase_											=				[]
										for arg in depreciated_args:
															# arg[2:] removes '--'
															if arg[2:] in PyTorchBenchmarkArguments.deprecated_args:
																				# arg[5:] removes '--no_'
																				full_error_msg += arg_error_msg.format(arg[5:])
															else:
																				wrong_args.append(_lowerCAmelCase)
										if len(_lowerCAmelCase) > 0:
															UpperCamelCase_											=				full_error_msg + begin_error_msg + str(_lowerCAmelCase)
										raise ValueError(_lowerCAmelCase)

					UpperCamelCase_											=				PyTorchBenchmark(args=_lowerCAmelCase)
					benchmark.run()


if __name__ == "__main__":
	main()
