


'''simple docstring'''






import jax
import jax.numpy as jnp
from bigbird_flax import FlaxBigBirdForNaturalQuestions
from datasets import load_from_disk

from transformers import BigBirdTokenizerFast


_A      =     {0: """null""", 1: """short""", 2: """long""", 3: """yes""", 4: """no"""}
_A      =     set("""""".join(["""‘""", """’""", """´""", """`""", """.""", """,""", """-""", """\""""]))


def A_					(				__SCREAMING_SNAKE_CASE :					List[Any]	,   __SCREAMING_SNAKE_CASE :					Any=0	,   __SCREAMING_SNAKE_CASE :					List[Any]=None       )     ->  List[str]:
					return [" ".join(x.split(''' '''       )[begin:end]       ) for x in answers if len(x.split(''' '''       )       ) > 1]


def A_					(				__SCREAMING_SNAKE_CASE :					str	,   __SCREAMING_SNAKE_CASE :					Tuple=False       )     ->  Union[str, Any]:
					if make_sub_answers:
										# if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
										# *e.g.* if the correct answer contains a prefix such as "the", or "a"
										__SCREAMING_SNAKE_CASE					:							str			   = (
										    given_answers + get_sub_answers(__SCREAMING_SNAKE_CASE	,   begin=1       ) + get_sub_answers(__SCREAMING_SNAKE_CASE	,   end=-1       )
										)
					__SCREAMING_SNAKE_CASE					:							Any			   = []
					for answer in given_answers:
										__SCREAMING_SNAKE_CASE					:							Tuple			   = answer.replace('''_'''	,   ''' '''       ).lower()
										__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = ''''''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ''' ''' for c in alias       )
										answers.append(''' '''.join(alias.split()       ).strip()       )
					return set(__SCREAMING_SNAKE_CASE       )


def A_					(				__SCREAMING_SNAKE_CASE :					Optional[Any]	,   __SCREAMING_SNAKE_CASE :					Tuple	,   __SCREAMING_SNAKE_CASE :					int=1	,   __SCREAMING_SNAKE_CASE :					Any=1_00       )     ->  int:
					__SCREAMING_SNAKE_CASE					:							Optional[int]			   = jax.lax.top_k(__SCREAMING_SNAKE_CASE	,   __SCREAMING_SNAKE_CASE       )
					__SCREAMING_SNAKE_CASE					:							str			   = jax.lax.top_k(__SCREAMING_SNAKE_CASE	,   __SCREAMING_SNAKE_CASE       )

					__SCREAMING_SNAKE_CASE					:							Tuple			   = best_end_idx[:, None] - best_start_idx[None, :]
					__SCREAMING_SNAKE_CASE					:							Optional[int]			   = jnp.logical_or(widths < 0	,   widths > max_size       )
					__SCREAMING_SNAKE_CASE					:							Tuple			   = (best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
					__SCREAMING_SNAKE_CASE					:							Dict			   = jnp.argmax(__SCREAMING_SNAKE_CASE       ).item()

					return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]


def A_					(				__SCREAMING_SNAKE_CASE :					Union[str, Any]       )     ->  Any:
					__SCREAMING_SNAKE_CASE					:							int			   = sample['''question''']['''text''']
					__SCREAMING_SNAKE_CASE					:							Any			   = sample['''document''']['''tokens''']['''token''']
					__SCREAMING_SNAKE_CASE					:							List[Any]			   = sample['''document''']['''tokens''']['''is_html''']
					__SCREAMING_SNAKE_CASE					:							List[Any]			   = sample['''annotations''']['''long_answer''']
					__SCREAMING_SNAKE_CASE					:							int			   = sample['''annotations''']['''short_answers''']

					__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = ''' '''.join([context[i] for i in range(len(__SCREAMING_SNAKE_CASE       )       ) if not is_html[i]]       )

					# 0 - No ; 1 - Yes
					for answer in sample["annotations"]["yes_no_answer"]:
										if answer == 0 or answer == 1:
															return {
															    "question": question,
															    "context": context_string,
															    "short": [],
															    "long": [],
															    "category": "no" if answer == 0 else "yes",
															}

					__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = []
					for s in short_answers:
										short_targets.extend(s['''text''']       )
					__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = list(set(__SCREAMING_SNAKE_CASE       )       )

					__SCREAMING_SNAKE_CASE					:							Tuple			   = []
					for s in long_answers:
										if s["start_token"] == -1:
															continue
										__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = context[s['''start_token'''] : s['''end_token''']]
										__SCREAMING_SNAKE_CASE					:							List[str]			   = is_html[s['''start_token'''] : s['''end_token''']]
										__SCREAMING_SNAKE_CASE					:							str			   = ''' '''.join([answer[i] for i in range(len(__SCREAMING_SNAKE_CASE       )       ) if not html[i]]       )
										if new_answer not in long_targets:
															long_targets.append(__SCREAMING_SNAKE_CASE       )

					__SCREAMING_SNAKE_CASE					:							Any			   = '''long_short''' if len(short_targets + long_targets       ) > 0 else '''null'''

					return {
					    "question": question,
					    "context": context_string,
					    "short": short_targets,
					    "long": long_targets,
					    "category": category,
					}


def A_					(				)     ->  str:
					__SCREAMING_SNAKE_CASE					:							List[Any]			   = load_from_disk('''natural-questions-validation'''       )
					__SCREAMING_SNAKE_CASE					:							int			   = dataset.map(__SCREAMING_SNAKE_CASE       ).remove_columns(['''annotations''', '''document''', '''id''']       )
					print(__SCREAMING_SNAKE_CASE       )

					__SCREAMING_SNAKE_CASE					:							Dict			   = dataset.filter(lambda __SCREAMING_SNAKE_CASE       : (len(x['''question''']       ) + len(x['''context''']       )) < 4 * 40_96       )
					__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = short_validation_dataset.filter(lambda __SCREAMING_SNAKE_CASE       : x["category"] != "null"       )
					short_validation_dataset

					__SCREAMING_SNAKE_CASE					:							List[str]			   = '''vasudevgupta/flax-bigbird-natural-questions'''
					__SCREAMING_SNAKE_CASE					:							Optional[int]			   = FlaxBigBirdForNaturalQuestions.from_pretrained(__SCREAMING_SNAKE_CASE       )
					__SCREAMING_SNAKE_CASE					:							int			   = BigBirdTokenizerFast.from_pretrained(__SCREAMING_SNAKE_CASE       )

					@jax.jit
					def forward(*__SCREAMING_SNAKE_CASE :					List[Any]	,   **__SCREAMING_SNAKE_CASE :					Union[str, Any]       ):
										__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = model(*__SCREAMING_SNAKE_CASE	,   **__SCREAMING_SNAKE_CASE       )
										return start_logits, end_logits, jnp.argmax(__SCREAMING_SNAKE_CASE	,   axis=-1       )

					def evaluate(__SCREAMING_SNAKE_CASE :					Tuple       ):
										# encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
										__SCREAMING_SNAKE_CASE					:							Optional[int]			   = tokenizer(
										    example['''question''']	,   example['''context''']	,   return_tensors='''np'''	,   max_length=40_96	,   padding='''max_length'''	,   truncation=__SCREAMING_SNAKE_CASE	,   )

										__SCREAMING_SNAKE_CASE					:							Dict			   = forward(**__SCREAMING_SNAKE_CASE       )

										__SCREAMING_SNAKE_CASE					:							str			   = CATEGORY_MAPPING[category.item()]

										__SCREAMING_SNAKE_CASE					:							Tuple			   = example['''long'''] + example['''short''']
										if example["category"] in ["yes", "no", "null"]:
															__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = [example['''category''']]
										__SCREAMING_SNAKE_CASE					:							Dict			   = example['''category'''] != '''null'''
										# Now target can be: "yes", "no", "null", "list of long & short answers"

										if predicted_category in ["yes", "no", "null"]:
															__SCREAMING_SNAKE_CASE					:							Optional[int]			   = [predicted_category]
															__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = example['''output'''] == example['''targets''']
															__SCREAMING_SNAKE_CASE					:							int			   = predicted_category != '''null'''
															return example

										__SCREAMING_SNAKE_CASE					:							Dict			   = 38 if predicted_category == '''short''' else 10_24
										__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = get_best_valid_start_end_idx(
										    start_scores[0]	,   end_scores[0]	,   top_k=8	,   max_size=__SCREAMING_SNAKE_CASE       )

										__SCREAMING_SNAKE_CASE					:							Any			   = inputs['''input_ids'''][0].tolist()
										__SCREAMING_SNAKE_CASE					:							List[Any]			   = [tokenizer.decode(input_ids[start_score : end_score + 1]       )]

										__SCREAMING_SNAKE_CASE					:							Optional[int]			   = expand_to_aliases(example['''targets''']	,   make_sub_answers=__SCREAMING_SNAKE_CASE       )
										__SCREAMING_SNAKE_CASE					:							int			   = expand_to_aliases(example['''output''']       )

										# some preprocessing to both prediction and answer
										__SCREAMING_SNAKE_CASE					:							Optional[int]			   = {''''''.join(a.split()       ) for a in answers}
										__SCREAMING_SNAKE_CASE					:							Dict			   = {''''''.join(p.split()       ) for p in predictions}
										__SCREAMING_SNAKE_CASE					:							List[Any]			   = {s for s in predictions if s not in ['''``''', '''\'\'''', '''`''', '''\'''']}

										# if there is a common element, it's a exact match
										__SCREAMING_SNAKE_CASE					:							List[str]			   = len(list(answers & predictions       )       ) > 0
										__SCREAMING_SNAKE_CASE					:							str			   = predicted_category != '''null''' and len(__SCREAMING_SNAKE_CASE       ) > 0

										return example

					__SCREAMING_SNAKE_CASE					:							Any			   = short_validation_dataset.map(__SCREAMING_SNAKE_CASE       )

					__SCREAMING_SNAKE_CASE					:							int			   = len(__SCREAMING_SNAKE_CASE       )
					__SCREAMING_SNAKE_CASE					:							List[Any]			   = len(short_validation_dataset.filter(lambda __SCREAMING_SNAKE_CASE       : x["match"] == 1       )       )
					print('''EM score:'''	,   (matched / total) * 1_00	,   '''%'''       )


if __name__ == "__main__":
						main()






'''simple docstring'''






from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin






@dataclass
class      SCREAMING_SNAKE_CASE_	( snake_case		):
			__a					:      torch.FloatTensor






class      SCREAMING_SNAKE_CASE_	( snake_case					,					snake_case		):




			@register_to_config
			def __init__(      self					,      lowercase = 1_6					,      lowercase = 8_8					,      lowercase = None					,      lowercase = None					,      lowercase = 1					,      lowercase = 0.0					,      lowercase = 3_2					,      lowercase = None					,      lowercase = False					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = "geglu"					,      lowercase = None					,      lowercase = False					,      lowercase = False					,      lowercase = False					,      lowercase = "layer_norm"					,      lowercase = True					,      )			->						Optional[int]:


								'''simple docstring'''






								super().__init__()
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = use_linear_projection
								__SCREAMING_SNAKE_CASE					:							Dict			   = num_attention_heads
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = attention_head_dim
								__SCREAMING_SNAKE_CASE					:							Any			   = num_attention_heads * attention_head_dim

								# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
								# Define whether input is continuous or discrete depending on configuration
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = (in_channels is not None) and (patch_size is None)
								__SCREAMING_SNAKE_CASE					:							List[str]			   = num_vector_embeds is not None
								__SCREAMING_SNAKE_CASE					:							Any			   = in_channels is not None and patch_size is not None

								if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
													__SCREAMING_SNAKE_CASE					:							Any			   = (
													    f"""The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"""
													    ''' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'''
													    ''' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'''
													    ''' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'''
													    ''' would be very nice if you could open a Pull request for the `transformer/config.json` file'''
													)
													deprecate('''norm_type!=num_embeds_ada_norm'''					,      '''1.0.0'''					,      lowercase					,      standard_warn=lowercase       )
													__SCREAMING_SNAKE_CASE					:							str			   = '''ada_norm'''

								if self.is_input_continuous and self.is_input_vectorized:
													raise ValueError(
													    f"""Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"""
													    ''' sure that either `in_channels` or `num_vector_embeds` is None.'''       )
								elif self.is_input_vectorized and self.is_input_patches:
													raise ValueError(
													    f"""Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"""
													    ''' sure that either `num_vector_embeds` or `num_patches` is None.'''       )
								elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
													raise ValueError(
													    f"""Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"""
													    f""" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."""       )

								# 2. Define input layers
								if self.is_input_continuous:
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = in_channels

													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = torch.nn.GroupNorm(num_groups=lowercase					,      num_channels=lowercase					,      eps=1e-6					,      affine=lowercase       )
													if use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							List[Any]			   = nn.Linear(lowercase					,      lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							Any			   = nn.Convad(lowercase					,      lowercase					,      kernel_size=1					,      stride=1					,      padding=0       )
								elif self.is_input_vectorized:
													assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
													assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = sample_size
													__SCREAMING_SNAKE_CASE					:							str			   = sample_size
													__SCREAMING_SNAKE_CASE					:							int			   = num_vector_embeds
													__SCREAMING_SNAKE_CASE					:							str			   = self.height * self.width

													__SCREAMING_SNAKE_CASE					:							int			   = ImagePositionalEmbeddings(
													    num_embed=lowercase					,      embed_dim=lowercase					,      height=self.height					,      width=self.width       )
								elif self.is_input_patches:
													assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

													__SCREAMING_SNAKE_CASE					:							Any			   = sample_size
													__SCREAMING_SNAKE_CASE					:							Any			   = sample_size

													__SCREAMING_SNAKE_CASE					:							int			   = patch_size
													__SCREAMING_SNAKE_CASE					:							Any			   = PatchEmbed(
													    height=lowercase					,      width=lowercase					,      patch_size=lowercase					,      in_channels=lowercase					,      embed_dim=lowercase					,      )

								# 3. Define transformers blocks
								__SCREAMING_SNAKE_CASE					:							Tuple			   = nn.ModuleList(
								    [
								        BasicTransformerBlock(
								            lowercase					,      lowercase					,      lowercase					,      dropout=lowercase					,      cross_attention_dim=lowercase					,      activation_fn=lowercase					,      num_embeds_ada_norm=lowercase					,      attention_bias=lowercase					,      only_cross_attention=lowercase					,      upcast_attention=lowercase					,      norm_type=lowercase					,      norm_elementwise_affine=lowercase					,      )
								        for d in range(lowercase       )
								    ]       )

								# 4. Define output layers
								__SCREAMING_SNAKE_CASE					:							List[str]			   = in_channels if out_channels is None else out_channels
								if self.is_input_continuous:
													# TODO: should use out_channels for continuous projections
													if use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							str			   = nn.Linear(lowercase					,      lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = nn.Convad(lowercase					,      lowercase					,      kernel_size=1					,      stride=1					,      padding=0       )
								elif self.is_input_vectorized:
													__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = nn.LayerNorm(lowercase       )
													__SCREAMING_SNAKE_CASE					:							int			   = nn.Linear(lowercase					,      self.num_vector_embeds - 1       )
								elif self.is_input_patches:
													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = nn.LayerNorm(lowercase					,      elementwise_affine=lowercase					,      eps=1e-6       )
													__SCREAMING_SNAKE_CASE					:							Any			   = nn.Linear(lowercase					,      2 * inner_dim       )
													__SCREAMING_SNAKE_CASE					:							Any			   = nn.Linear(lowercase					,      patch_size * patch_size * self.out_channels       )




			def 			_snake_case   (      self					,      lowercase					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = True					,      )			->						Optional[Any]:


								'''simple docstring'''






								if attention_mask is not None and attention_mask.ndim == 2:
													# assume that mask is expressed as:
													#   (1 = keep,      0 = discard)
													# convert mask into a bias that can be added to attention scores:
													#       (keep = +0,     discard = -10000.0)
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = (1 - attention_mask.to(hidden_states.dtype       )) * -1_0_0_0_0.0
													__SCREAMING_SNAKE_CASE					:							Dict			   = attention_mask.unsqueeze(1       )

								# convert encoder_attention_mask to a bias the same way we do for attention_mask
								if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
													__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = (1 - encoder_attention_mask.to(hidden_states.dtype       )) * -1_0_0_0_0.0
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = encoder_attention_mask.unsqueeze(1       )

								# 1. Input
								if self.is_input_continuous:
													__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Optional[Any]			   = hidden_states.shape
													__SCREAMING_SNAKE_CASE					:							Dict			   = hidden_states

													__SCREAMING_SNAKE_CASE					:							str			   = self.norm(lowercase       )
													if not self.use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							int			   = self.proj_in(lowercase       )
																		__SCREAMING_SNAKE_CASE					:							List[str]			   = hidden_states.shape[1]
																		__SCREAMING_SNAKE_CASE					:							List[str]			   = hidden_states.permute(0					,      2					,      3					,      1       ).reshape(lowercase					,      height * width					,      lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							Tuple			   = hidden_states.shape[1]
																		__SCREAMING_SNAKE_CASE					:							Optional[int]			   = hidden_states.permute(0					,      2					,      3					,      1       ).reshape(lowercase					,      height * width					,      lowercase       )
																		__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.proj_in(lowercase       )
								elif self.is_input_vectorized:
													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = self.latent_image_embedding(lowercase       )
								elif self.is_input_patches:
													__SCREAMING_SNAKE_CASE					:							int			   = self.pos_embed(lowercase       )

								# 2. Blocks
								for block in self.transformer_blocks:
													__SCREAMING_SNAKE_CASE					:							Tuple			   = block(
													    lowercase					,      attention_mask=lowercase					,      encoder_hidden_states=lowercase					,      encoder_attention_mask=lowercase					,      timestep=lowercase					,      cross_attention_kwargs=lowercase					,      class_labels=lowercase					,      )

								# 3. Output
								if self.is_input_continuous:
													if not self.use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							Tuple			   = hidden_states.reshape(lowercase					,      lowercase					,      lowercase					,      lowercase       ).permute(0					,      3					,      1					,      2       ).contiguous()
																		__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = self.proj_out(lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							str			   = self.proj_out(lowercase       )
																		__SCREAMING_SNAKE_CASE					:							str			   = hidden_states.reshape(lowercase					,      lowercase					,      lowercase					,      lowercase       ).permute(0					,      3					,      1					,      2       ).contiguous()

													__SCREAMING_SNAKE_CASE					:							str			   = hidden_states + residual
								elif self.is_input_vectorized:
													__SCREAMING_SNAKE_CASE					:							Dict			   = self.norm_out(lowercase       )
													__SCREAMING_SNAKE_CASE					:							str			   = self.out(lowercase       )
													# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = logits.permute(0					,      2					,      1       )

													# log(p(x_0))
													__SCREAMING_SNAKE_CASE					:							int			   = F.log_softmax(logits.double()					,      dim=1       ).float()
								elif self.is_input_patches:
													# TODO: cleanup!
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.transformer_blocks[0].norma.emb(
													    lowercase					,      lowercase					,      hidden_dtype=hidden_states.dtype       )
													__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.proj_out_a(F.silu(lowercase       )       ).chunk(2					,      dim=1       )
													__SCREAMING_SNAKE_CASE					:							int			   = self.norm_out(lowercase       ) * (1 + scale[:, None]) + shift[:, None]
													__SCREAMING_SNAKE_CASE					:							int			   = self.proj_out_a(lowercase       )

													# unpatchify
													__SCREAMING_SNAKE_CASE					:							List[str]			   = int(hidden_states.shape[1] ** 0.5       )
													__SCREAMING_SNAKE_CASE					:							int			   = hidden_states.reshape(
													    shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)       )
													__SCREAMING_SNAKE_CASE					:							int			   = torch.einsum('''nhwpqc->nchpwq'''					,      lowercase       )
													__SCREAMING_SNAKE_CASE					:							Any			   = hidden_states.reshape(
													    shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)       )

								if not return_dict:
													return (output,)

								return TransformeraDModelOutput(sample=lowercase       )


