



'''simple docstring'''






import logging
import re

import pytorch_quantization
import pytorch_quantization.nn as quant_nn
import torch
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor


UpperCamelCase__:   Tuple								=			logging.getLogger(__name__)

UpperCamelCase__:   str								=			50  # max width of layer names
UpperCamelCase__:   List[str]								=			70  # max width of quantizer names


def  snake_case_     (					_lowerCAmelCase  :						List[Any]  )      ->   Any:
							UpperCAmelCase :       List[str] 							=		parser.add_argument_group('''quant_trainer arguments'''  )
							group.add_argument('''--wprec'''   ,      type=_lowerCAmelCase   ,      default=8   ,      help='''weight precision'''  )
							group.add_argument('''--aprec'''   ,      type=_lowerCAmelCase   ,      default=8   ,      help='''activation precision'''  )
							group.add_argument('''--quant-per-tensor'''   ,      action='''store_true'''   ,      help='''per tensor weight scaling'''  )
							group.add_argument('''--quant-disable'''   ,      action='''store_true'''   ,      help='''disable all quantizers'''  )
							group.add_argument('''--quant-disable-embeddings'''   ,      action='''store_true'''   ,      help='''disable all embeddings quantizers'''  )
							group.add_argument('''--quant-disable-keyword'''   ,      type=_lowerCAmelCase   ,      nargs='''+'''   ,      help='''disable quantizers by keyword'''  )
							group.add_argument('''--quant-disable-layer-module'''   ,      type=_lowerCAmelCase   ,      help='''disable quantizers by keyword under layer.'''  )
							group.add_argument('''--quant-enable-layer-module'''   ,      type=_lowerCAmelCase   ,      help='''enable quantizers by keyword under layer'''  )
							group.add_argument('''--calibrator'''   ,      default='''max'''   ,      help='''which quantization range calibrator to use'''  )
							group.add_argument('''--percentile'''   ,      default=_lowerCAmelCase   ,      type=_lowerCAmelCase   ,      help='''percentile for PercentileCalibrator'''  )
							group.add_argument('''--fuse-qkv'''   ,      action='''store_true'''   ,      help='''use the same scale factor for qkv'''  )
							group.add_argument('''--clip-gelu'''   ,      metavar='''N'''   ,      type=_lowerCAmelCase   ,      help='''clip gelu output maximum value to N'''  )
							group.add_argument(
							    '''--recalibrate-weights'''   ,      action='''store_true'''   ,      help=(
							        '''recalibrate weight amaxes by taking the max of the weights.'''
							        ''' amaxes will be computed with the current quantization granularity (axis).'''
							    )   ,      )


def  snake_case_     (					_lowerCAmelCase  :						Optional[Any]  )      ->   List[Any]:

							if args.calibrator == "max":
														UpperCAmelCase :       str 							=		'''max'''
							elif args.calibrator == "percentile":
														if args.percentile is None:
																					raise ValueError('''Specify --percentile when using percentile calibrator'''  )
														UpperCAmelCase :       Optional[int] 							=		'''histogram'''
							elif args.calibrator == "mse":
														UpperCAmelCase :       List[str] 							=		'''histogram'''
							else:
														raise ValueError(f"""Invalid calibrator {args.calibrator}"""  )

							UpperCAmelCase :       Optional[int] 							=		QuantDescriptor(num_bits=args.aprec   ,      calib_method=_lowerCAmelCase  )
							UpperCAmelCase :       List[str] 							=		QuantDescriptor(num_bits=args.wprec   ,      axis=(None if args.quant_per_tensor else (0,))  )
							quant_nn.QuantLinear.set_default_quant_desc_input(_lowerCAmelCase  )
							quant_nn.QuantLinear.set_default_quant_desc_weight(_lowerCAmelCase  )


def  snake_case_     (					_lowerCAmelCase  :						Optional[int]   ,      _lowerCAmelCase  :						Optional[Any]   ,      _lowerCAmelCase  :						Union[str, Any]=False   ,      _lowerCAmelCase  :						int=False  )      ->   List[str]:
							logger.info('''Configuring Model for Quantization'''  )
							logger.info(f"""using quantization package {pytorch_quantization.__file__}"""  )

							if not calib:
														if args.quant_disable_embeddings:
																					set_quantizer_by_name(_lowerCAmelCase   ,      ['''embeddings''']   ,      which='''weight'''   ,      _disabled=_lowerCAmelCase  )

														if args.quant_disable:
																					set_quantizer_by_name(_lowerCAmelCase   ,      ['''''']   ,      _disabled=_lowerCAmelCase  )

														if args.quant_disable_keyword:
																					set_quantizer_by_name(_lowerCAmelCase   ,      args.quant_disable_keyword   ,      _disabled=_lowerCAmelCase  )

														if args.quant_disable_layer_module:
																					set_quantizer_by_name(_lowerCAmelCase   ,      [R'''layer.\d+.''' + args.quant_disable_layer_module]   ,      _disabled=_lowerCAmelCase  )

														if args.quant_enable_layer_module:
																					set_quantizer_by_name(_lowerCAmelCase   ,      [R'''layer.\d+.''' + args.quant_enable_layer_module]   ,      _disabled=_lowerCAmelCase  )

														if args.recalibrate_weights:
																					recalibrate_weights(_lowerCAmelCase  )

														if args.fuse_qkv:
																					fuse_qkv(_lowerCAmelCase   ,      _lowerCAmelCase  )

							if args.clip_gelu:
														clip_gelu(_lowerCAmelCase   ,      args.clip_gelu  )

							# if args.local_rank in [-1, 0] and not calib:
							print_quant_summary(_lowerCAmelCase  )


def  snake_case_     (					_lowerCAmelCase  :						str  )      ->   List[Any]:
							logger.info('''Enabling Calibration'''  )
							for name, module in model.named_modules():
														if name.endswith('''_quantizer'''  ):
																					if module._calibrator is not None:
																												module.disable_quant()
																												module.enable_calib()
																					else:
																												module.disable()
																					logger.info(f"""{name:80}: {module}"""  )


def  snake_case_     (					_lowerCAmelCase  :						int   ,      _lowerCAmelCase  :						Union[str, Any]  )      ->   str:
							logger.info('''Loading calibrated amax'''  )
							for name, module in model.named_modules():
														if name.endswith('''_quantizer'''  ):
																					if module._calibrator is not None:
																												if isinstance(module._calibrator   ,      calib.MaxCalibrator  ):
																																			module.load_calib_amax()
																												else:
																																			module.load_calib_amax('''percentile'''   ,      percentile=args.percentile  )
																												module.enable_quant()
																												module.disable_calib()
																					else:
																												module.enable()
							model.cuda()
							print_quant_summary(_lowerCAmelCase  )


def  snake_case_     (					_lowerCAmelCase  :						List[str]   ,      _lowerCAmelCase  :						str  )      ->   Any:

							def fusea(_lowerCAmelCase  :						str   ,      _lowerCAmelCase  :						Dict   ,      _lowerCAmelCase  :						Optional[Any]  ):
														for mod in [qq, qk, qv]:
																					if not hasattr(_lowerCAmelCase   ,      '''_amax'''  ):
																												print('''          WARNING: NO AMAX BUFFER'''  )
																												return
														UpperCAmelCase :       Union[str, Any] 							=		qq._amax.detach().item()
														UpperCAmelCase :       Tuple 							=		qk._amax.detach().item()
														UpperCAmelCase :       Optional[int] 							=		qv._amax.detach().item()

														UpperCAmelCase :       str 							=		max(_lowerCAmelCase   ,      _lowerCAmelCase   ,      _lowerCAmelCase  )
														qq._amax.fill_(_lowerCAmelCase  )
														qk._amax.fill_(_lowerCAmelCase  )
														qv._amax.fill_(_lowerCAmelCase  )
														logger.info(f"""          q={q:5.2f} k={k:5.2f} v={v:5.2f} -> {amax:5.2f}"""  )

							for name, mod in model.named_modules():
														if name.endswith('''.attention.self'''  ):
																					logger.info(f"""FUSE_QKV: {name:{name_width}}"""  )
																					fusea(mod.matmul_q_input_quantizer   ,      mod.matmul_k_input_quantizer   ,      mod.matmul_v_input_quantizer  )
																					if args.quant_per_tensor:
																												fusea(mod.query._weight_quantizer   ,      mod.key._weight_quantizer   ,      mod.value._weight_quantizer  )


def  snake_case_     (					_lowerCAmelCase  :						Tuple   ,      _lowerCAmelCase  :						Union[str, Any]  )      ->   Optional[Any]:

							for name, mod in model.named_modules():
														if name.endswith('''.output.dense'''  ) and not name.endswith('''attention.output.dense'''  ):
																					UpperCAmelCase :       Union[str, Any] 							=		mod._input_quantizer._amax.data.detach().item()
																					mod._input_quantizer._amax.data.detach().clamp_(max=_lowerCAmelCase  )
																					UpperCAmelCase :       Dict 							=		mod._input_quantizer._amax.data.detach().item()
																					logger.info(f"""CLIP_GELU: {name:{name_width}} amax: {amax_init:5.2f} -> {amax:5.2f}"""  )


def  snake_case_     (					_lowerCAmelCase  :						Dict  )      ->   Any:

							for name, mod in model.named_modules():
														if hasattr(_lowerCAmelCase   ,      '''_weight_quantizer'''  ) and mod._weight_quantizer.axis is not None:
																					UpperCAmelCase :       Any 							=		mod.weight.shape[0]
																					UpperCAmelCase :       Dict 							=		mod._weight_quantizer._amax.detach()
																					UpperCAmelCase :       Any 							=		torch.ones(_lowerCAmelCase   ,      dtype=amax.dtype   ,      device=amax.device  ) * amax
																					print(f"""expanding {name} {amax} -> {mod._weight_quantizer._amax}"""  )


def  snake_case_     (					_lowerCAmelCase  :						Optional[Any]  )      ->   Union[str, Any]:

							for name, mod in model.named_modules():
														if hasattr(_lowerCAmelCase   ,      '''_weight_quantizer'''  ):
																					if not hasattr(mod.weight_quantizer   ,      '''_amax'''  ):
																												print('''RECALIB: {name:{name_width}} WARNING: NO AMAX BUFFER'''  )
																												continue

																					# determine which axes to reduce across
																					# e.g. a 4D tensor quantized per axis 0 should reduce over (1,2,3)
																					UpperCAmelCase :       str 							=		set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis  )
																					UpperCAmelCase :       Any 							=		set(range(len(mod.weight.size()  )  )  ) - axis_set
																					UpperCAmelCase :       List[str] 							=		pytorch_quantization.utils.reduce_amax(mod.weight   ,      axis=_lowerCAmelCase   ,      keepdims=_lowerCAmelCase  ).detach()
																					logger.info(f"""RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}"""  )
																					UpperCAmelCase :       Optional[Any] 							=		amax


def  snake_case_     (					_lowerCAmelCase  :						List[Any]   ,      _lowerCAmelCase  :						Union[str, Any]=25   ,      _lowerCAmelCase  :						Dict=180   ,      _lowerCAmelCase  :						str=None  )      ->   str:

							if ignore is None:
														UpperCAmelCase :       Optional[Any] 							=		[]
							elif not isinstance(_lowerCAmelCase   ,      _lowerCAmelCase  ):
														UpperCAmelCase :       Any 							=		[ignore]

							UpperCAmelCase :       Tuple 							=		0
							for name, mod in model.named_modules():
														if not hasattr(_lowerCAmelCase   ,      '''weight'''  ):
																					continue
														UpperCAmelCase :       Union[str, Any] 							=		max(_lowerCAmelCase   ,      len(_lowerCAmelCase  )  )

							for name, mod in model.named_modules():
														UpperCAmelCase :       Union[str, Any] 							=		getattr(_lowerCAmelCase   ,      '''_input_quantizer'''   ,      _lowerCAmelCase  )
														UpperCAmelCase :       Tuple 							=		getattr(_lowerCAmelCase   ,      '''_weight_quantizer'''   ,      _lowerCAmelCase  )
														if not hasattr(_lowerCAmelCase   ,      '''weight'''  ):
																					continue
														if type(_lowerCAmelCase  ) in ignore:
																					continue
														if [True for s in ignore if type(_lowerCAmelCase  ) is str and s in name]:
																					continue
														UpperCAmelCase :       str 							=		f"""Act:{input_q.extra_repr()}"""
														UpperCAmelCase :       int 							=		f"""Wgt:{weight_q.extra_repr()}"""
														UpperCAmelCase :       Tuple 							=		f"""{name:{name_width}} {act_str} {wgt_str}"""
														if len(_lowerCAmelCase  ) <= line_width:
																					logger.info(_lowerCAmelCase  )
														else:
																					logger.info(f"""{name:{name_width}} {act_str}"""  )
																					logger.info(f"""{"  ":{name_width}} {wgt_str}"""  )


def  snake_case_     (					_lowerCAmelCase  :						Dict  )      ->   List[Any]:
							UpperCAmelCase :       List[Any] 							=		0
							for name, mod in model.named_modules():
														if isinstance(_lowerCAmelCase   ,      pytorch_quantization.nn.TensorQuantizer  ):
																					print(f"""{name:80} {mod}"""  )
																					count += 1
							print(f"""{count} TensorQuantizers found in model"""  )


def  snake_case_     (					_lowerCAmelCase  :						Any   ,      _lowerCAmelCase  :						Optional[Any]   ,      _lowerCAmelCase  :						int   ,      _lowerCAmelCase  :						Tuple   ,      _lowerCAmelCase  :						Optional[Any]  )      ->   Tuple:
							UpperCAmelCase :       Any 							=		getattr(_lowerCAmelCase   ,      _lowerCAmelCase   ,      _lowerCAmelCase  )
							if quantizer_mod is not None:
														assert hasattr(_lowerCAmelCase   ,      _lowerCAmelCase  )
														setattr(_lowerCAmelCase   ,      _lowerCAmelCase   ,      _lowerCAmelCase  )
							else:
														logger.warning(f"""{name} has no {quantizer}"""  )


def  snake_case_     (					_lowerCAmelCase  :						int   ,      _lowerCAmelCase  :						Tuple   ,      _lowerCAmelCase  :						Union[str, Any]="both"   ,      **_lowerCAmelCase  :						Optional[Any]  )      ->   List[str]:
							UpperCAmelCase :       List[Any] 							=		f"""Warning: changing {which} quantizers of {name:{qname_width}}"""
							for k, v in kwargs.items():
														s += f""" {k}={v}"""
														if which in ["input", "both"]:
																					set_quantizer(_lowerCAmelCase   ,      _lowerCAmelCase   ,      '''_input_quantizer'''   ,      _lowerCAmelCase   ,      _lowerCAmelCase  )
														if which in ["weight", "both"]:
																					set_quantizer(_lowerCAmelCase   ,      _lowerCAmelCase   ,      '''_weight_quantizer'''   ,      _lowerCAmelCase   ,      _lowerCAmelCase  )
							logger.info(_lowerCAmelCase  )




def  snake_case_     (					_lowerCAmelCase  :						Union[str, Any]   ,      _lowerCAmelCase  :						Union[str, Any]   ,      **_lowerCAmelCase  :						Any  )      ->   str:

							for name, mod in model.named_modules():
														if hasattr(_lowerCAmelCase   ,      '''_input_quantizer'''  ) or hasattr(_lowerCAmelCase   ,      '''_weight_quantizer'''  ):
																					for n in names:
																												if re.search(_lowerCAmelCase   ,      _lowerCAmelCase  ):
																																			set_quantizers(_lowerCAmelCase   ,      _lowerCAmelCase   ,      **_lowerCAmelCase  )
														elif name.endswith('''_quantizer'''  ):
																					for n in names:
																												if re.search(_lowerCAmelCase   ,      _lowerCAmelCase  ):
																																			UpperCAmelCase :       List[str] 							=		f"""Warning: changing {name:{name_width}}"""
																																			for k, v in kwargs.items():
																																										s += f""" {k}={v}"""
																																										setattr(_lowerCAmelCase   ,      _lowerCAmelCase   ,      _lowerCAmelCase  )
																																			logger.info(_lowerCAmelCase  )










'''simple docstring'''






import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
						import jax.numpy as jnp

if is_tf_available():
						import tensorflow as tf

if is_torch_available():
						import torch

class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	def      A	(						self			:     Any		)			->		str:
								UpperCAmelCase :       Any 							=		{
								    '''task_specific_params''': {
								        '''summarization''': {'''length_penalty''': 1.0, '''max_length''': 128, '''min_length''': 12, '''num_beams''': 4},
								        '''summarization_cnn''': {'''length_penalty''': 2.0, '''max_length''': 142, '''min_length''': 56, '''num_beams''': 4},
								        '''summarization_xsum''': {'''length_penalty''': 1.0, '''max_length''': 62, '''min_length''': 11, '''num_beams''': 6},
								    }
								}
								UpperCAmelCase :       int 							=		{
								    '''task_specific_params.summarization.length_penalty''': 1.0,
								    '''task_specific_params.summarization.max_length''': 128,
								    '''task_specific_params.summarization.min_length''': 12,
								    '''task_specific_params.summarization.num_beams''': 4,
								    '''task_specific_params.summarization_cnn.length_penalty''': 2.0,
								    '''task_specific_params.summarization_cnn.max_length''': 142,
								    '''task_specific_params.summarization_cnn.min_length''': 56,
								    '''task_specific_params.summarization_cnn.num_beams''': 4,
								    '''task_specific_params.summarization_xsum.length_penalty''': 1.0,
								    '''task_specific_params.summarization_xsum.max_length''': 62,
								    '''task_specific_params.summarization_xsum.min_length''': 11,
								    '''task_specific_params.summarization_xsum.num_beams''': 6,
								}

								self.assertEqual(flatten_dict(__snake_case		)		,			__snake_case		)







	def      A	(						self			:     int		)			->		str:
								UpperCAmelCase :       Any 							=		np.random.randn(3		,			4		)
								self.assertTrue(np.allclose(transpose(__snake_case		)		,			x.transpose()		)		)

								UpperCAmelCase :       str 							=		np.random.randn(3		,			4		,			5		)
								self.assertTrue(np.allclose(transpose(__snake_case		,			axes=(1, 2, 0)		)		,			x.transpose((1, 2, 0)		)		)		)







	@require_torch
	def      A	(						self			:     str		)			->		Union[str, Any]:
								UpperCAmelCase :       Any 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       List[Any] 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		)		,			transpose(__snake_case		).numpy()		)		)

								UpperCAmelCase :       Tuple 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       Any 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		,			axes=(1, 2, 0)		)		,			transpose(__snake_case		,			axes=(1, 2, 0)		).numpy()		)		)







	@require_tf
	def      A	(						self			:     List[str]		)			->		Optional[Any]:
								UpperCAmelCase :       int 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Optional[int] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		)		,			transpose(__snake_case		).numpy()		)		)

								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       str 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		,			axes=(1, 2, 0)		)		,			transpose(__snake_case		,			axes=(1, 2, 0)		).numpy()		)		)







	@require_flax
	def      A	(						self			:     Tuple		)			->		Any:
								UpperCAmelCase :       List[Any] 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       List[str] 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		)		,			np.asarray(transpose(__snake_case		)		)		)		)

								UpperCAmelCase :       Dict 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       int 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(transpose(__snake_case		,			axes=(1, 2, 0)		)		,			np.asarray(transpose(__snake_case		,			axes=(1, 2, 0)		)		)		)		)







	def      A	(						self			:     Optional[Any]		)			->		Any:
								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(3		,			4		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(4, 3)		)		,			np.reshape(__snake_case		,			(4, 3)		)		)		)

								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(3		,			4		,			5		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(12, 5)		)		,			np.reshape(__snake_case		,			(12, 5)		)		)		)







	@require_torch
	def      A	(						self			:     Union[str, Any]		)			->		int:
								UpperCAmelCase :       Dict 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Optional[Any] 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(4, 3)		)		,			reshape(__snake_case		,			(4, 3)		).numpy()		)		)

								UpperCAmelCase :       List[str] 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       List[Any] 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(12, 5)		)		,			reshape(__snake_case		,			(12, 5)		).numpy()		)		)







	@require_tf
	def      A	(						self			:     int		)			->		List[str]:
								UpperCAmelCase :       List[Any] 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       List[str] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(4, 3)		)		,			reshape(__snake_case		,			(4, 3)		).numpy()		)		)

								UpperCAmelCase :       List[Any] 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       Optional[Any] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(12, 5)		)		,			reshape(__snake_case		,			(12, 5)		).numpy()		)		)







	@require_flax
	def      A	(						self			:     Any		)			->		Dict:
								UpperCAmelCase :       Tuple 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Union[str, Any] 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(4, 3)		)		,			np.asarray(reshape(__snake_case		,			(4, 3)		)		)		)		)

								UpperCAmelCase :       Any 							=		np.random.randn(3		,			4		,			5		)
								UpperCAmelCase :       Optional[Any] 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(reshape(__snake_case		,			(12, 5)		)		,			np.asarray(reshape(__snake_case		,			(12, 5)		)		)		)		)







	def      A	(						self			:     List[Any]		)			->		List[Any]:
								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(1		,			3		,			4		)
								self.assertTrue(np.allclose(squeeze(__snake_case		)		,			np.squeeze(__snake_case		)		)		)

								UpperCAmelCase :       str 							=		np.random.randn(1		,			4		,			1		,			5		)
								self.assertTrue(np.allclose(squeeze(__snake_case		,			axis=2		)		,			np.squeeze(__snake_case		,			axis=2		)		)		)







	@require_torch
	def      A	(						self			:     Optional[int]		)			->		Union[str, Any]:
								UpperCAmelCase :       Optional[int] 							=		np.random.randn(1		,			3		,			4		)
								UpperCAmelCase :       List[str] 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		)		,			squeeze(__snake_case		).numpy()		)		)

								UpperCAmelCase :       Any 							=		np.random.randn(1		,			4		,			1		,			5		)
								UpperCAmelCase :       str 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		,			axis=2		)		,			squeeze(__snake_case		,			axis=2		).numpy()		)		)







	@require_tf
	def      A	(						self			:     Optional[Any]		)			->		Dict:
								UpperCAmelCase :       int 							=		np.random.randn(1		,			3		,			4		)
								UpperCAmelCase :       Optional[int] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		)		,			squeeze(__snake_case		).numpy()		)		)

								UpperCAmelCase :       List[str] 							=		np.random.randn(1		,			4		,			1		,			5		)
								UpperCAmelCase :       Optional[int] 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		,			axis=2		)		,			squeeze(__snake_case		,			axis=2		).numpy()		)		)







	@require_flax
	def      A	(						self			:     List[Any]		)			->		Dict:
								UpperCAmelCase :       Optional[int] 							=		np.random.randn(1		,			3		,			4		)
								UpperCAmelCase :       int 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		)		,			np.asarray(squeeze(__snake_case		)		)		)		)

								UpperCAmelCase :       str 							=		np.random.randn(1		,			4		,			1		,			5		)
								UpperCAmelCase :       int 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(squeeze(__snake_case		,			axis=2		)		,			np.asarray(squeeze(__snake_case		,			axis=2		)		)		)		)







	def      A	(						self			:     Optional[Any]		)			->		int:
								UpperCAmelCase :       Optional[Any] 							=		np.random.randn(3		,			4		)
								self.assertTrue(np.allclose(expand_dims(__snake_case		,			axis=1		)		,			np.expand_dims(__snake_case		,			axis=1		)		)		)







	@require_torch
	def      A	(						self			:     List[str]		)			->		Tuple:
								UpperCAmelCase :       Tuple 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Tuple 							=		torch.tensor(__snake_case		)
								self.assertTrue(np.allclose(expand_dims(__snake_case		,			axis=1		)		,			expand_dims(__snake_case		,			axis=1		).numpy()		)		)







	@require_tf
	def      A	(						self			:     List[str]		)			->		Union[str, Any]:
								UpperCAmelCase :       Union[str, Any] 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       Any 							=		tf.constant(__snake_case		)
								self.assertTrue(np.allclose(expand_dims(__snake_case		,			axis=1		)		,			expand_dims(__snake_case		,			axis=1		).numpy()		)		)







	@require_flax
	def      A	(						self			:     Any		)			->		List[Any]:
								UpperCAmelCase :       List[str] 							=		np.random.randn(3		,			4		)
								UpperCAmelCase :       str 							=		jnp.array(__snake_case		)
								self.assertTrue(np.allclose(expand_dims(__snake_case		,			axis=1		)		,			np.asarray(expand_dims(__snake_case		,			axis=1		)		)		)		)





