
import argparse

import requests
import torch
from PIL import Image

from transformers import CLIPProcessor, GroupViTConfig, GroupViTModel






def      lowerCamelCase__       (       __lowerCAmelCase						:						Dict       ):




			"""simple docstring"""


			if "img_encoder.pos_embed" in name:
						lowerCAmelCase_      			=			name.replace("img_encoder.pos_embed"					,							"vision_model.embeddings.position_embeddings"       )
			if "img_encoder.patch_embed.proj" in name:
						lowerCAmelCase_      			=			name.replace("img_encoder.patch_embed.proj"					,							"vision_model.embeddings.patch_embeddings.projection"       )
			if "img_encoder.patch_embed.norm" in name:
						lowerCAmelCase_      			=			name.replace("img_encoder.patch_embed.norm"					,							"vision_model.embeddings.layernorm"       )
			if "img_encoder.layers" in name:
						lowerCAmelCase_      			=			name.replace("img_encoder.layers"					,							"vision_model.encoder.stages"       )
			if "blocks" in name and "res" not in name:
						lowerCAmelCase_      			=			name.replace("blocks"					,							"layers"       )
			if "attn" in name and "pre_assign" not in name:
						lowerCAmelCase_      			=			name.replace("attn"					,							"self_attn"       )
			if "proj" in name and "self_attn" in name and "text" not in name:
						lowerCAmelCase_      			=			name.replace("proj"					,							"out_proj"       )
			if "pre_assign_attn.attn.proj" in name:
						lowerCAmelCase_      			=			name.replace("pre_assign_attn.attn.proj"					,							"pre_assign_attn.attn.out_proj"       )
			if "norm1" in name:
						lowerCAmelCase_      			=			name.replace("norm1"					,							"layer_norm1"       )
			if "norm2" in name and "pre_assign" not in name:
						lowerCAmelCase_      			=			name.replace("norm2"					,							"layer_norm2"       )
			if "img_encoder.norm" in name:
						lowerCAmelCase_      			=			name.replace("img_encoder.norm"					,							"vision_model.layernorm"       )
			# text encoder
			if "text_encoder.token_embedding" in name:
						lowerCAmelCase_      			=			name.replace("text_encoder.token_embedding"					,							"text_model.embeddings.token_embedding"       )
			if "text_encoder.positional_embedding" in name:
						lowerCAmelCase_      			=			name.replace("text_encoder.positional_embedding"					,							"text_model.embeddings.position_embedding.weight"       )
			if "text_encoder.transformer.resblocks." in name:
						lowerCAmelCase_      			=			name.replace("text_encoder.transformer.resblocks."					,							"text_model.encoder.layers."       )
			if "ln_1" in name:
						lowerCAmelCase_      			=			name.replace("ln_1"					,							"layer_norm1"       )
			if "ln_2" in name:
						lowerCAmelCase_      			=			name.replace("ln_2"					,							"layer_norm2"       )
			if "c_fc" in name:
						lowerCAmelCase_      			=			name.replace("c_fc"					,							"fc1"       )
			if "c_proj" in name:
						lowerCAmelCase_      			=			name.replace("c_proj"					,							"fc2"       )
			if "text_encoder" in name:
						lowerCAmelCase_      			=			name.replace("text_encoder"					,							"text_model"       )
			if "ln_final" in name:
						lowerCAmelCase_      			=			name.replace("ln_final"					,							"final_layer_norm"       )
			# projection layers
			if "img_projector.linear_hidden." in name:
						lowerCAmelCase_      			=			name.replace("img_projector.linear_hidden."					,							"visual_projection."       )
			if "img_projector.linear_out." in name:
						lowerCAmelCase_      			=			name.replace("img_projector.linear_out."					,							"visual_projection.3."       )
			if "text_projector.linear_hidden" in name:
						lowerCAmelCase_      			=			name.replace("text_projector.linear_hidden"					,							"text_projection"       )
			if "text_projector.linear_out" in name:
						lowerCAmelCase_      			=			name.replace("text_projector.linear_out"					,							"text_projection.3"       )

			return name






def      lowerCamelCase__       (       __lowerCAmelCase						:						Dict					,							__lowerCAmelCase						:						str       ):




			"""simple docstring"""


			for key in orig_state_dict.copy().keys():
						lowerCAmelCase_      			=			orig_state_dict.pop(__lowerCAmelCase       )

						if "qkv" in key:
									# weights and biases of the key, value and query projections of vision encoder's attention layers require special treatment:
									# we need to split them up into separate matrices/vectors
									lowerCAmelCase_      			=			key.split("."       )
									lowerCAmelCase_			,						lowerCAmelCase_      			=			int(key_split[2]       ), int(key_split[4]       )
									lowerCAmelCase_      			=			config.vision_config.hidden_size
									if "weight" in key:
												lowerCAmelCase_      			=			val[:dim, :]
												lowerCAmelCase_      			=			val[dim : dim * 2, :]
												lowerCAmelCase_      			=			val[-dim:, :]
									else:
												lowerCAmelCase_      			=			val[:dim]
												lowerCAmelCase_      			=			val[dim : dim * 2]
												lowerCAmelCase_      			=			val[-dim:]
						elif "in_proj" in key:
									# weights and biases of the key, value and query projections of text encoder's attention layers require special treatment:
									# we need to split them up into separate matrices/vectors
									lowerCAmelCase_      			=			key.split("."       )
									lowerCAmelCase_      			=			int(key_split[3]       )
									lowerCAmelCase_      			=			config.text_config.hidden_size
									if "weight" in key:
												lowerCAmelCase_      			=			val[:dim, :]
												lowerCAmelCase_      			=			val[
												    dim : dim * 2, :
												]
												lowerCAmelCase_      			=			val[-dim:, :]
									else:
												lowerCAmelCase_      			=			val[:dim]
												lowerCAmelCase_      			=			val[dim : dim * 2]
												lowerCAmelCase_      			=			val[-dim:]
						else:
									lowerCAmelCase_      			=			rename_key(__lowerCAmelCase       )
									# squeeze if necessary
									if (
									    "text_projection.0" in new_name
									    or "text_projection.3" in new_name
									    or "visual_projection.0" in new_name
									    or "visual_projection.3" in new_name
									):
												lowerCAmelCase_      			=			val.squeeze_()
									else:
												lowerCAmelCase_      			=			val

			return orig_state_dict






def      lowerCamelCase__       (       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			"http://images.cocodataset.org/val2017/000000039769.jpg"
			lowerCAmelCase_      			=			Image.open(requests.get(__lowerCAmelCase					,							stream=__lowerCAmelCase       ).raw       )
			return im






@torch.no_grad()
def      lowerCamelCase__       (       __lowerCAmelCase						:						Optional[Any]					,							__lowerCAmelCase						:						Optional[Any]					,							__lowerCAmelCase						:						Optional[Any]="groupvit-gcc-yfcc"					,							__lowerCAmelCase						:						int=False       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			GroupViTConfig()
			lowerCAmelCase_      			=			GroupViTModel(__lowerCAmelCase       ).eval()

			lowerCAmelCase_      			=			torch.load(__lowerCAmelCase					,							map_location="cpu"       )["model"]
			lowerCAmelCase_      			=			convert_state_dict(__lowerCAmelCase					,							__lowerCAmelCase       )
			lowerCAmelCase_			,						lowerCAmelCase_      			=			model.load_state_dict(__lowerCAmelCase					,							strict=__lowerCAmelCase       )
			assert missing_keys == ["text_model.embeddings.position_ids"]
			assert (unexpected_keys == ["multi_label_logit_scale"]) or (len(__lowerCAmelCase       ) == 0)

			# verify result
			lowerCAmelCase_      			=			CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32"       )
			lowerCAmelCase_      			=			prepare_img()
			lowerCAmelCase_      			=			processor(text=["a photo of a cat", "a photo of a dog"]					,							images=__lowerCAmelCase					,							padding=__lowerCAmelCase					,							return_tensors="pt"       )

			with torch.no_grad():
						lowerCAmelCase_      			=			model(**__lowerCAmelCase       )

			if model_name == "groupvit-gcc-yfcc":
						lowerCAmelCase_      			=			torch.tensor([[13.3_523, 6.3_629]]       )
			elif model_name == "groupvit-gcc-redcaps":
						lowerCAmelCase_      			=			torch.tensor([[16.1_873, 8.6_230]]       )
			else:
						raise ValueError(F"""Model name {model_name} not supported."""       )
			assert torch.allclose(outputs.logits_per_image					,							__lowerCAmelCase					,							atol=1e-3       )

			processor.save_pretrained(__lowerCAmelCase       )
			model.save_pretrained(__lowerCAmelCase       )
			print("Successfully saved processor and model to"					,							__lowerCAmelCase       )

			if push_to_hub:
						print("Pushing to the hub..."       )
						processor.push_to_hub(__lowerCAmelCase					,							organization="nielsr"       )
						model.push_to_hub(__lowerCAmelCase					,							organization="nielsr"       )


if __name__ == "__main__":
			_A     =  argparse.ArgumentParser()
			parser.add_argument(
			    "--pytorch_dump_folder_path", default=None, type=str, help="Path to dump the processor and PyTorch model."
			)
			parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to GroupViT checkpoint")
			parser.add_argument(
			    "--model_name",
			    default="groupvit-gccy-fcc",
			    type=str,
			    help="Name of the model. Expecting either 'groupvit-gcc-yfcc' or 'groupvit-gcc-redcaps'",
			)
			parser.add_argument(
			    "--push_to_hub",
			    action="store_true",
			    help="Whether or not to push the converted model and processor to the 🤗 hub using the provided `model_name`.",
			)
			_A     =  parser.parse_args()

			convert_groupvit_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)





import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin





class 			_lowerCAmelCase      (			__a      ,   __a ):




				@register_to_config
				def __init__(      self       ,       _UpperCamelCase = 128       ,       _UpperCamelCase = 256       ,       _UpperCamelCase = 2000.0       ,       _UpperCamelCase = 768       ,       _UpperCamelCase = 12       ,       _UpperCamelCase = 12       ,       _UpperCamelCase = 64       ,       _UpperCamelCase = 2_048       ,       _UpperCamelCase = 0.1       ,       )       ->   str:
							super().__init__()

							lowerCAmelCase_      			=			nn.Sequential(
							    nn.Linear(_UpperCamelCase       ,       d_model * 4       ,       bias=_UpperCamelCase  )       ,       nn.SiLU()       ,       nn.Linear(d_model * 4       ,       d_model * 4       ,       bias=_UpperCamelCase  )       ,       nn.SiLU()       ,       )

							lowerCAmelCase_      			=			nn.Embedding(_UpperCamelCase       ,       _UpperCamelCase  )
							lowerCAmelCase_      			=			False

							lowerCAmelCase_      			=			nn.Linear(_UpperCamelCase       ,       _UpperCamelCase       ,       bias=_UpperCamelCase  )

							lowerCAmelCase_      			=			nn.Dropout(p=_UpperCamelCase  )

							lowerCAmelCase_      			=			nn.ModuleList()
							for lyr_num in range(_UpperCamelCase  ):
										# FiLM conditional T5 decoder
										lowerCAmelCase_      			=			DecoderLayer(d_model=_UpperCamelCase       ,       d_kv=_UpperCamelCase       ,       num_heads=_UpperCamelCase       ,       d_ff=_UpperCamelCase       ,       dropout_rate=_UpperCamelCase  )
										self.decoders.append(_UpperCamelCase  )

							lowerCAmelCase_      			=			TaLayerNorm(_UpperCamelCase  )

							lowerCAmelCase_      			=			nn.Dropout(p=_UpperCamelCase  )
							lowerCAmelCase_      			=			nn.Linear(_UpperCamelCase       ,       _UpperCamelCase       ,       bias=_UpperCamelCase  )




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   int:
							lowerCAmelCase_      			=			torch.mul(query_input.unsqueeze(-1  )       ,       key_input.unsqueeze(-2  )  )
							return mask.unsqueeze(-3  )




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   Any:
							lowerCAmelCase_			,						lowerCAmelCase_			,						lowerCAmelCase_      			=			decoder_input_tokens.shape
							assert decoder_noise_time.shape == (batch,)

							# decoder_noise_time is in [0, 1), so rescale to expected timing range.
							lowerCAmelCase_      			=			get_timestep_embedding(
							    decoder_noise_time * self.config.max_decoder_noise_time       ,       embedding_dim=self.config.d_model       ,       max_period=self.config.max_decoder_noise_time       ,       ).to(dtype=self.dtype  )

							lowerCAmelCase_      			=			self.conditioning_emb(_UpperCamelCase  ).unsqueeze(1  )

							assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

							lowerCAmelCase_      			=			decoder_input_tokens.shape[1]

							# If we want to use relative positions for audio context, we can just offset
							# this sequence by the length of encodings_and_masks.
							lowerCAmelCase_      			=			torch.broadcast_to(
							    torch.arange(_UpperCamelCase       ,       device=decoder_input_tokens.device  )       ,       (batch, seq_length)       ,       )

							lowerCAmelCase_      			=			self.position_encoding(_UpperCamelCase  )

							lowerCAmelCase_      			=			self.continuous_inputs_projection(_UpperCamelCase  )
							inputs += position_encodings
							lowerCAmelCase_      			=			self.dropout(_UpperCamelCase  )

							# decoder: No padding present.
							lowerCAmelCase_      			=			torch.ones(
							    decoder_input_tokens.shape[:2]       ,       device=decoder_input_tokens.device       ,       dtype=inputs.dtype  )

							# Translate encoding masks to encoder-decoder masks.
							lowerCAmelCase_      			=			[(x, self.encoder_decoder_mask(_UpperCamelCase       ,       _UpperCamelCase  )) for x, y in encodings_and_masks]

							# cross attend style: concat encodings
							lowerCAmelCase_      			=			torch.cat([x[0] for x in encodings_and_encdec_masks]       ,       dim=1  )
							lowerCAmelCase_      			=			torch.cat([x[1] for x in encodings_and_encdec_masks]       ,       dim=-1  )

							for lyr in self.decoders:
										lowerCAmelCase_      			=			lyr(
										    _UpperCamelCase       ,       conditioning_emb=_UpperCamelCase       ,       encoder_hidden_states=_UpperCamelCase       ,       encoder_attention_mask=_UpperCamelCase       ,       )[0]

							lowerCAmelCase_      			=			self.decoder_norm(_UpperCamelCase  )
							lowerCAmelCase_      			=			self.post_dropout(_UpperCamelCase  )

							lowerCAmelCase_      			=			self.spec_out(_UpperCamelCase  )
							return spec_out





class 			_lowerCAmelCase      (			nn.Module ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase=1e-6  )       ->   Dict:
							super().__init__()
							lowerCAmelCase_      			=			nn.ModuleList()

							# cond self attention: layer 0
							self.layer.append(
							    TaLayerSelfAttentionCond(d_model=_UpperCamelCase       ,       d_kv=_UpperCamelCase       ,       num_heads=_UpperCamelCase       ,       dropout_rate=_UpperCamelCase  )  )

							# cross attention: layer 1
							self.layer.append(
							    TaLayerCrossAttention(
							        d_model=_UpperCamelCase       ,       d_kv=_UpperCamelCase       ,       num_heads=_UpperCamelCase       ,       dropout_rate=_UpperCamelCase       ,       layer_norm_epsilon=_UpperCamelCase       ,       )  )

							# Film Cond MLP + dropout: last layer
							self.layer.append(
							    TaLayerFFCond(d_model=_UpperCamelCase       ,       d_ff=_UpperCamelCase       ,       dropout_rate=_UpperCamelCase       ,       layer_norm_epsilon=_UpperCamelCase  )  )




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase=None       ,       _UpperCamelCase=None       ,       _UpperCamelCase=None       ,       _UpperCamelCase=None       ,       _UpperCamelCase=None       ,       )       ->   Any:
							lowerCAmelCase_      			=			self.layer[0](
							    _UpperCamelCase       ,       conditioning_emb=_UpperCamelCase       ,       attention_mask=_UpperCamelCase       ,       )

							if encoder_hidden_states is not None:
										lowerCAmelCase_      			=			torch.where(encoder_attention_mask > 0       ,       0       ,       -1e1_0  ).to(
										    encoder_hidden_states.dtype  )

										lowerCAmelCase_      			=			self.layer[1](
										    _UpperCamelCase       ,       key_value_states=_UpperCamelCase       ,       attention_mask=_UpperCamelCase       ,       )

							# Apply Film Conditional Feed Forward layer
							lowerCAmelCase_      			=			self.layer[-1](_UpperCamelCase       ,       _UpperCamelCase  )

							return (hidden_states,)





class 			_lowerCAmelCase      (			nn.Module ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   Optional[int]:
							super().__init__()
							lowerCAmelCase_      			=			TaLayerNorm(_UpperCamelCase  )
							lowerCAmelCase_      			=			TaFiLMLayer(in_features=d_model * 4       ,       out_features=_UpperCamelCase  )
							lowerCAmelCase_      			=			Attention(query_dim=_UpperCamelCase       ,       heads=_UpperCamelCase       ,       dim_head=_UpperCamelCase       ,       out_bias=_UpperCamelCase       ,       scale_qk=_UpperCamelCase  )
							lowerCAmelCase_      			=			nn.Dropout(_UpperCamelCase  )




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase=None       ,       _UpperCamelCase=None       ,       )       ->   Dict:
							# pre_self_attention_layer_norm
							lowerCAmelCase_      			=			self.layer_norm(_UpperCamelCase  )

							if conditioning_emb is not None:
										lowerCAmelCase_      			=			self.FiLMLayer(_UpperCamelCase       ,       _UpperCamelCase  )

							# Self-attention block
							lowerCAmelCase_      			=			self.attention(_UpperCamelCase  )

							lowerCAmelCase_      			=			hidden_states + self.dropout(_UpperCamelCase  )

							return hidden_states





class 			_lowerCAmelCase      (			nn.Module ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   Dict:
							super().__init__()
							lowerCAmelCase_      			=			Attention(query_dim=_UpperCamelCase       ,       heads=_UpperCamelCase       ,       dim_head=_UpperCamelCase       ,       out_bias=_UpperCamelCase       ,       scale_qk=_UpperCamelCase  )
							lowerCAmelCase_      			=			TaLayerNorm(_UpperCamelCase       ,       eps=_UpperCamelCase  )
							lowerCAmelCase_      			=			nn.Dropout(_UpperCamelCase  )




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase=None       ,       _UpperCamelCase=None       ,       )       ->   int:
							lowerCAmelCase_      			=			self.layer_norm(_UpperCamelCase  )
							lowerCAmelCase_      			=			self.attention(
							    _UpperCamelCase       ,       encoder_hidden_states=_UpperCamelCase       ,       attention_mask=attention_mask.squeeze(1  )       ,       )
							lowerCAmelCase_      			=			hidden_states + self.dropout(_UpperCamelCase  )
							return layer_output





class 			_lowerCAmelCase      (			nn.Module ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   Tuple:
							super().__init__()
							lowerCAmelCase_      			=			TaDenseGatedActDense(d_model=_UpperCamelCase       ,       d_ff=_UpperCamelCase       ,       dropout_rate=_UpperCamelCase  )
							lowerCAmelCase_      			=			TaFiLMLayer(in_features=d_model * 4       ,       out_features=_UpperCamelCase  )
							lowerCAmelCase_      			=			TaLayerNorm(_UpperCamelCase       ,       eps=_UpperCamelCase  )
							lowerCAmelCase_      			=			nn.Dropout(_UpperCamelCase  )




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase=None  )       ->   Tuple:
							lowerCAmelCase_      			=			self.layer_norm(_UpperCamelCase  )
							if conditioning_emb is not None:
										lowerCAmelCase_      			=			self.film(_UpperCamelCase       ,       _UpperCamelCase  )

							lowerCAmelCase_      			=			self.DenseReluDense(_UpperCamelCase  )
							lowerCAmelCase_      			=			hidden_states + self.dropout(_UpperCamelCase  )
							return hidden_states





class 			_lowerCAmelCase      (			nn.Module ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   Union[str, Any]:
							super().__init__()
							lowerCAmelCase_      			=			nn.Linear(_UpperCamelCase       ,       _UpperCamelCase       ,       bias=_UpperCamelCase  )
							lowerCAmelCase_      			=			nn.Linear(_UpperCamelCase       ,       _UpperCamelCase       ,       bias=_UpperCamelCase  )
							lowerCAmelCase_      			=			nn.Linear(_UpperCamelCase       ,       _UpperCamelCase       ,       bias=_UpperCamelCase  )
							lowerCAmelCase_      			=			nn.Dropout(_UpperCamelCase  )
							lowerCAmelCase_      			=			NewGELUActivation()




				def   __a		(      self       ,       _UpperCamelCase  )       ->   int:
							lowerCAmelCase_      			=			self.act(self.wi_a(_UpperCamelCase  )  )
							lowerCAmelCase_      			=			self.wi_a(_UpperCamelCase  )
							lowerCAmelCase_      			=			hidden_gelu * hidden_linear
							lowerCAmelCase_      			=			self.dropout(_UpperCamelCase  )

							lowerCAmelCase_      			=			self.wo(_UpperCamelCase  )
							return hidden_states





class 			_lowerCAmelCase      (			nn.Module ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase=1e-6  )       ->   int:
							super().__init__()
							lowerCAmelCase_      			=			nn.Parameter(torch.ones(_UpperCamelCase  )  )
							lowerCAmelCase_      			=			eps




				def   __a		(      self       ,       _UpperCamelCase  )       ->   Union[str, Any]:
							# T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean
							# Square Layer Normalization https://arxiv.org/abs/1910.07467 thus variance is calculated
							# w/o mean and there is no bias. Additionally we want to make sure that the accumulation for
							# half-precision inputs is done in fp32

							lowerCAmelCase_      			=			hidden_states.to(torch.floataa  ).pow(2  ).mean(-1       ,       keepdim=_UpperCamelCase  )
							lowerCAmelCase_      			=			hidden_states * torch.rsqrt(variance + self.variance_epsilon  )

							# convert into half-precision if necessary
							if self.weight.dtype in [torch.floataa, torch.bfloataa]:
										lowerCAmelCase_      			=			hidden_states.to(self.weight.dtype  )

							return self.weight * hidden_states





class 			_lowerCAmelCase      (			nn.Module ):




				def   __a		(      self       ,       _UpperCamelCase  )       ->   torch.Tensor:
							return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi  ) * (input + 0.044715 * torch.pow(_UpperCamelCase       ,       3.0  ))  ))





class 			_lowerCAmelCase      (			nn.Module ):




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   str:
							super().__init__()
							lowerCAmelCase_      			=			nn.Linear(_UpperCamelCase       ,       out_features * 2       ,       bias=_UpperCamelCase  )




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase  )       ->   int:
							lowerCAmelCase_      			=			self.scale_bias(_UpperCamelCase  )
							lowerCAmelCase_			,						lowerCAmelCase_      			=			torch.chunk(_UpperCamelCase       ,       2       ,       -1  )
							lowerCAmelCase_      			=			x * (1 + scale) + shift
							return x



