



import colorsys
import io

import cva
import matplotlib as mpl
import matplotlib.colors as mplc
import matplotlib.figure as mplfigure
import numpy as np
import torch
from matplotlib.backends.backend_agg import FigureCanvasAgg

from utils import img_tensorize


UpperCAmelCase							=       1_000

class 			__snake_case   :







							'''simple docstring'''






							def __init__(				self	,	a_	,	a_=1.2	,	a_="g"	,	a_=0.5	,	a_="-"	,	a_="test_out.jpg"	,	a_=True	,	a_=False	,	a_=None	,	a_=None	,	a_=0.7	,	):
											if isinstance(a_	,	torch.Tensor		):
															a__					   =		img.numpy().astype("""np.uint8"""		)
											if isinstance(a_	,	a_		):
															a__					   =		img_tensorize(a_		)
											assert isinstance(a_	,	np.ndarray		)

											a__   ,      a__					   =		img.shape[1], img.shape[0]
											a__					   =		mplfigure.Figure(frameon=a_		)
											a__					   =		fig.get_dpi()
											a__					   =		(width * scale + 1E-2) / dpi
											a__					   =		(height * scale + 1E-2) / dpi
											fig.set_size_inches(a_	,	a_		)
											a__					   =		fig.add_axes([0.0, 0.0, 1.0, 1.0]		)
											ax.axis("""off"""		)
											ax.set_xlim(0.0	,	a_		)
											ax.set_ylim(a_		)

											a__					   =		saveas
											a__					   =		rgb
											a__					   =		pynb
											a__					   =		img
											a__					   =		edgecolor
											a__					   =		0.5
											a__					   =		linestyle
											a__					   =		int(np.sqrt(min(a_	,	a_		)		) * scale // 3		)
											a__					   =		width
											a__					   =		height
											a__					   =		scale
											a__					   =		fig
											a__					   =		ax
											a__					   =		pad
											a__					   =		idaobj
											a__					   =		idaattr
											a__					   =		FigureCanvasAgg(a_		)


							def _a   (				self	,	a_	,	a_=None		):
											if color is None:
															a__					   =		self.edgecolor
											((a__)   ,      (a__)   ,      (a__)   ,      (a__))					   =		box
											a__					   =		xa - xa
											a__					   =		ya - ya
											self.ax.add_patch(
											    mpl.patches.Rectangle(
											        (xa, ya)	,	a_	,	a_	,	fill=a_	,	edgecolor=a_	,	linewidth=self.font_size // 3	,	alpha=self.alpha	,	linestyle=self.linestyle	,	)		)


							def _a   (				self	,	a_	,	a_=None	,	a_=None	,	a_=None	,	a_=None		):
											if len(boxes.shape		) > 2:
															a__					   =		boxes[0]
											if len(obj_ids.shape		) > 1:
															a__					   =		obj_ids[0]
											if len(obj_scores.shape		) > 1:
															a__					   =		obj_scores[0]
											if len(attr_ids.shape		) > 1:
															a__					   =		attr_ids[0]
											if len(attr_scores.shape		) > 1:
															a__					   =		attr_scores[0]
											if isinstance(a_	,	torch.Tensor		):
															a__					   =		boxes.numpy()
											if isinstance(a_	,	a_		):
															a__					   =		np.array(a_		)
											assert isinstance(a_	,	np.ndarray		)
											a__					   =		np.prod(boxes[:, 2:] - boxes[:, :2]	,	axis=1		)
											a__					   =		np.argsort(-areas		).tolist()
											a__					   =		boxes[sorted_idxs] if boxes is not None else None
											a__					   =		obj_ids[sorted_idxs] if obj_ids is not None else None
											a__					   =		obj_scores[sorted_idxs] if obj_scores is not None else None
											a__					   =		attr_ids[sorted_idxs] if attr_ids is not None else None
											a__					   =		attr_scores[sorted_idxs] if attr_scores is not None else None

											a__					   =		[self._random_color(maximum=1		) for _ in range(len(a_		)		)]
											a__					   =		[assigned_colors[idx] for idx in sorted_idxs]
											if obj_ids is not None:
															a__					   =		self._create_text_labels_attr(a_	,	a_	,	a_	,	a_		)
															for i in range(len(a_		)		):
																			a__					   =		assigned_colors[i]
																			self.add_box(boxes[i]	,	a_		)
																			self.draw_labels(labels[i]	,	boxes[i]	,	a_		)


							def _a   (				self	,	a_	,	a_	,	a_		):
											a__   ,      a__   ,      a__   ,      a__					   =		box
											a__					   =		(xa, ya)
											a__					   =		(ya - ya) * (xa - xa)
											a__					   =		_SMALL_OBJ * self.scale
											if instance_area < small or ya - ya < 40 * self.scale:
															if ya >= self.height - 5:
																			a__					   =		(xa, ya)
															else:
																			a__					   =		(xa, ya)

											a__					   =		(ya - ya) / np.sqrt(self.height * self.width		)
											a__					   =		self._change_color_brightness(a_	,	brightness_factor=0.7		)
											a__					   =		np.clip((height_ratio - 0.02) / 0.08 + 1	,	1.2	,	2		)
											font_size *= 0.75 * self.font_size

											self.draw_text(
											    text=a_	,	position=a_	,	color=a_	,	)


							def _a   (				self	,	a_	,	a_	,	a_="g"	,	a_="left"	,	):
											a__					   =		0
											a__					   =		self.font_size
											a__					   =		np.maximum(list(mplc.to_rgb(a_		)		)	,	0.2		)
											a__					   =		max(0.8	,	np.max(a_		)		)
											a__					   =		{
											    """facecolor""": """black""",
											    """alpha""": self.alpha,
											    """pad""": self.pad,
											    """edgecolor""": """none""",
											}
											a__   ,      a__					   =		position
											self.ax.text(
											    a_	,	a_	,	a_	,	size=font_size * self.scale	,	family="""sans-serif"""	,	bbox=a_	,	verticalalignment="""top"""	,	horizontalalignment=a_	,	color=a_	,	zorder=10	,	rotation=a_	,	)


							def _a   (				self	,	a_=None		):
											if saveas is None:
															a__					   =		self.saveas
											if saveas.lower().endswith(""".jpg"""		) or saveas.lower().endswith(""".png"""		):
															cva.imwrite(
															    a_	,	self._get_buffer()[:, :, ::-1]	,	)
											else:
															self.fig.savefig(a_		)


							def _a   (				self	,	a_	,	a_	,	a_	,	a_		):
											a__					   =		[self.idaobj[i] for i in classes]
											a__					   =		[self.idaattr[i] for i in attr_classes]
											a__					   =		[
											    F'''{label} {score:.2f} {attr} {attr_score:.2f}'''
											    for label, score, attr, attr_score in zip(a_	,	a_	,	a_	,	a_		)
											]
											return labels


							def _a   (				self	,	a_	,	a_		):
											a__					   =		[self.idaobj[i] for i in classes]
											if scores is not None:
															if labels is None:
																			a__					   =		["""{:.0f}%""".format(s * 100		) for s in scores]
															else:
																			a__					   =		["""{} {:.0f}%""".format(a_	,	s * 100		) for li, s in zip(a_	,	a_		)]
											return labels


							def _a   (				self	,	a_=255		):
											a__					   =		np.random.randint(0	,	len(_COLORS		)		)
											a__					   =		_COLORS[idx] * maximum
											if not self.rgb:
															a__					   =		ret[::-1]
											return ret


							def _a   (				self		):
											if not self.pynb:
															a__   ,      a__					   =		self.canvas.print_to_buffer()
															if (width, height) != (self.width, self.height):
																			a__					   =		cva.resize(self.img	,	(width, height)		)
															else:
																			a__					   =		self.img
											else:
															a__					   =		io.BytesIO()  # works for cairo backend
															self.canvas.print_rgba(a_		)
															a__   ,      a__					   =		self.width, self.height
															a__					   =		buf.getvalue()
															a__					   =		self.img

											a__					   =		np.frombuffer(a_	,	dtype="""uint8"""		)
											a__					   =		buffer.reshape(a_	,	a_	,	4		)
											a__   ,      a__					   =		np.split(a_	,	[3]	,	axis=2		)

											try:
															import numexpr as ne  # fuse them with numexpr

															a__					   =		ne.evaluate("""img * (1 - alpha / 255.0) + rgb * (alpha / 255.0)"""		)
											except ImportError:
															a__					   =		alpha.astype("""float32"""		) / 255.0
															a__					   =		img * (1 - alpha) + rgb * alpha

											return visualized_image.astype("""uint8"""		)







							def _a   (				self	,	a_	,	a_		):
											assert brightness_factor >= -1.0 and brightness_factor <= 1.0
											a__					   =		mplc.to_rgb(a_		)
											a__					   =		colorsys.rgb_to_hls(*mplc.to_rgb(a_		)		)
											a__					   =		polygon_color[1] + (brightness_factor * polygon_color[1])
											a__					   =		0.0 if modified_lightness < 0.0 else modified_lightness
											a__					   =		1.0 if modified_lightness > 1.0 else modified_lightness
											a__					   =		colorsys.hls_to_rgb(polygon_color[0]	,	a_	,	polygon_color[2]		)
											return modified_color




UpperCAmelCase							=       (
    np.array(
        [
            0.0_00,
            0.4_47,
            0.7_41,
            0.8_50,
            0.3_25,
            0.0_98,
            0.9_29,
            0.6_94,
            0.1_25,
            0.4_94,
            0.1_84,
            0.5_56,
            0.4_66,
            0.6_74,
            0.1_88,
            0.3_01,
            0.7_45,
            0.9_33,
            0.6_35,
            0.0_78,
            0.1_84,
            0.3_00,
            0.3_00,
            0.3_00,
            0.6_00,
            0.6_00,
            0.6_00,
            1.0_00,
            0.0_00,
            0.0_00,
            1.0_00,
            0.5_00,
            0.0_00,
            0.7_49,
            0.7_49,
            0.0_00,
            0.0_00,
            1.0_00,
            0.0_00,
            0.0_00,
            0.0_00,
            1.0_00,
            0.6_67,
            0.0_00,
            1.0_00,
            0.3_33,
            0.3_33,
            0.0_00,
            0.3_33,
            0.6_67,
            0.0_00,
            0.3_33,
            1.0_00,
            0.0_00,
            0.6_67,
            0.3_33,
            0.0_00,
            0.6_67,
            0.6_67,
            0.0_00,
            0.6_67,
            1.0_00,
            0.0_00,
            1.0_00,
            0.3_33,
            0.0_00,
            1.0_00,
            0.6_67,
            0.0_00,
            1.0_00,
            1.0_00,
            0.0_00,
            0.0_00,
            0.3_33,
            0.5_00,
            0.0_00,
            0.6_67,
            0.5_00,
            0.0_00,
            1.0_00,
            0.5_00,
            0.3_33,
            0.0_00,
            0.5_00,
            0.3_33,
            0.3_33,
            0.5_00,
            0.3_33,
            0.6_67,
            0.5_00,
            0.3_33,
            1.0_00,
            0.5_00,
            0.6_67,
            0.0_00,
            0.5_00,
            0.6_67,
            0.3_33,
            0.5_00,
            0.6_67,
            0.6_67,
            0.5_00,
            0.6_67,
            1.0_00,
            0.5_00,
            1.0_00,
            0.0_00,
            0.5_00,
            1.0_00,
            0.3_33,
            0.5_00,
            1.0_00,
            0.6_67,
            0.5_00,
            1.0_00,
            1.0_00,
            0.5_00,
            0.0_00,
            0.3_33,
            1.0_00,
            0.0_00,
            0.6_67,
            1.0_00,
            0.0_00,
            1.0_00,
            1.0_00,
            0.3_33,
            0.0_00,
            1.0_00,
            0.3_33,
            0.3_33,
            1.0_00,
            0.3_33,
            0.6_67,
            1.0_00,
            0.3_33,
            1.0_00,
            1.0_00,
            0.6_67,
            0.0_00,
            1.0_00,
            0.6_67,
            0.3_33,
            1.0_00,
            0.6_67,
            0.6_67,
            1.0_00,
            0.6_67,
            1.0_00,
            1.0_00,
            1.0_00,
            0.0_00,
            1.0_00,
            1.0_00,
            0.3_33,
            1.0_00,
            1.0_00,
            0.6_67,
            1.0_00,
            0.3_33,
            0.0_00,
            0.0_00,
            0.5_00,
            0.0_00,
            0.0_00,
            0.6_67,
            0.0_00,
            0.0_00,
            0.8_33,
            0.0_00,
            0.0_00,
            1.0_00,
            0.0_00,
            0.0_00,
            0.0_00,
            0.1_67,
            0.0_00,
            0.0_00,
            0.3_33,
            0.0_00,
            0.0_00,
            0.5_00,
            0.0_00,
            0.0_00,
            0.6_67,
            0.0_00,
            0.0_00,
            0.8_33,
            0.0_00,
            0.0_00,
            1.0_00,
            0.0_00,
            0.0_00,
            0.0_00,
            0.1_67,
            0.0_00,
            0.0_00,
            0.3_33,
            0.0_00,
            0.0_00,
            0.5_00,
            0.0_00,
            0.0_00,
            0.6_67,
            0.0_00,
            0.0_00,
            0.8_33,
            0.0_00,
            0.0_00,
            1.0_00,
            0.0_00,
            0.0_00,
            0.0_00,
            0.1_43,
            0.1_43,
            0.1_43,
            0.8_57,
            0.8_57,
            0.8_57,
            1.0_00,
            1.0_00,
            1.0_00,
        ]
    )
    .astype(np.floataa)
    .reshape(-1, 3)
)






import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp

def       A_   (						__a  :					str			,		__a  :					int			,		__a  :					List[Any]			,		__a  :					Dict			,		__a  :					int = 4_096		):

				"""simple docstring"""
				a__   ,      a__   ,      a__					   =		key.shape[-3:]
				a__					   =		value.shape[-1]
				a__					   =		min(__a			,		__a		)
				a__					   =		query / jnp.sqrt(__a		)

				@functools.partial(jax.checkpoint			,		prevent_cse=__a		)
				def summarize_chunk(__a  :					Optional[Any]			,		__a  :					Tuple			,		__a  :					Dict		):
								a__					   =		jnp.einsum("""...qhd,...khd->...qhk"""			,		__a			,		__a			,		precision=__a		)

								a__					   =		jnp.max(__a			,		axis=-1			,		keepdims=__a		)
								a__					   =		jax.lax.stop_gradient(__a		)
								a__					   =		jnp.exp(attn_weights - max_score		)

								a__					   =		jnp.einsum("""...vhf,...qhv->...qhf"""			,		__a			,		__a			,		precision=__a		)
								a__					   =		jnp.einsum("""...qhk->...qh"""			,		__a		)

								return (exp_values, exp_weights.sum(axis=-1		), max_score)

				def chunk_scanner(__a  :					Tuple		):
								# julienne key array
								a__					   =		jax.lax.dynamic_slice(
								    operand=__a			,		start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]			,		slice_sizes=list(key.shape[:-3]		) + [key_chunk_size, num_heads, k_features]			,		)

								# julienne value array
								a__					   =		jax.lax.dynamic_slice(
								    operand=__a			,		start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]			,		slice_sizes=list(value.shape[:-3]		) + [key_chunk_size, num_heads, v_features]			,		)

								return summarize_chunk(__a			,		__a			,		__a		)

				a__   ,      a__   ,      a__					   =		jax.lax.map(f=__a			,		xs=jnp.arange(0			,		__a			,		__a		)		)

				a__					   =		jnp.max(__a			,		axis=0			,		keepdims=__a		)
				a__					   =		jnp.exp(chunk_max - global_max		)

				chunk_values *= jnp.expand_dims(__a			,		axis=-1		)
				chunk_weights *= max_diffs

				a__					   =		chunk_values.sum(axis=0		)
				a__					   =		jnp.expand_dims(__a			,		-1		).sum(axis=0		)

				return all_values / all_weights

def       A_   (						__a  :					Tuple			,		__a  :					List[Any]			,		__a  :					int			,		__a  :					List[str]=jax.lax.Precision.HIGHEST			,		__a  :					int = 1_024			,		__a  :					int = 4_096		):

				"""simple docstring"""
				a__   ,      a__   ,      a__					   =		query.shape[-3:]

				def chunk_scanner(__a  :					List[str]			,		__a  :					List[str]		):
								# julienne query array
								a__					   =		jax.lax.dynamic_slice(
								    operand=__a			,		start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]			,		slice_sizes=list(query.shape[:-3]		) + [min(__a			,		__a		), num_heads, q_features]			,		)

								return (
								    chunk_idx + query_chunk_size,  # unused ignore it
								    _query_chunk_attention(
								        query=__a			,		key=__a			,		value=__a			,		precision=__a			,		key_chunk_size=__a		),
								)

				a__   ,      a__					   =		jax.lax.scan(
				    f=__a			,		init=0			,		xs=__a			,		length=math.ceil(num_q / query_chunk_size		)		)

				return jnp.concatenate(__a			,		axis=-3		)  # fuse the chunked result back

class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   int       =       8
							UpperCamelCase__					:   int       =       6_4
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa


							def _a   (				self		):
											a__					   =		self.dim_head * self.heads
											a__					   =		self.dim_head**-0.5

											# Weights were exported with old names {to_q, to_k, to_v, to_out}
											a__					   =		nn.Dense(a_	,	use_bias=a_	,	dtype=self.dtype	,	name="""to_q"""		)
											a__					   =		nn.Dense(a_	,	use_bias=a_	,	dtype=self.dtype	,	name="""to_k"""		)
											a__					   =		nn.Dense(a_	,	use_bias=a_	,	dtype=self.dtype	,	name="""to_v"""		)

											a__					   =		nn.Dense(self.query_dim	,	dtype=self.dtype	,	name="""to_out_0"""		)
											a__					   =		nn.Dropout(rate=self.dropout		)


							def _a   (				self	,	a_		):
											a__   ,      a__   ,      a__					   =		tensor.shape
											a__					   =		self.heads
											a__					   =		tensor.reshape(a_	,	a_	,	a_	,	dim // head_size		)
											a__					   =		jnp.transpose(a_	,	(0, 2, 1, 3)		)
											a__					   =		tensor.reshape(batch_size * head_size	,	a_	,	dim // head_size		)
											return tensor


							def _a   (				self	,	a_		):
											a__   ,      a__   ,      a__					   =		tensor.shape
											a__					   =		self.heads
											a__					   =		tensor.reshape(batch_size // head_size	,	a_	,	a_	,	a_		)
											a__					   =		jnp.transpose(a_	,	(0, 2, 1, 3)		)
											a__					   =		tensor.reshape(batch_size // head_size	,	a_	,	dim * head_size		)
											return tensor







							def __call__(				self	,	a_	,	a_=None	,	a_=True		):
											a__					   =		hidden_states if context is None else context

											a__					   =		self.query(a_		)
											a__					   =		self.key(a_		)
											a__					   =		self.value(a_		)

											a__					   =		self.reshape_heads_to_batch_dim(a_		)
											a__					   =		self.reshape_heads_to_batch_dim(a_		)
											a__					   =		self.reshape_heads_to_batch_dim(a_		)

											if self.use_memory_efficient_attention:
															a__					   =		query_states.transpose(1	,	0	,	2		)
															a__					   =		key_states.transpose(1	,	0	,	2		)
															a__					   =		value_states.transpose(1	,	0	,	2		)

															# this if statement create a chunk size for each layer of the unet
															# the chunk size is equal to the query_length dimension of the deepest layer of the unet

															a__					   =		query_states.shape[-3]
															if flatten_latent_dim % 64 == 0:
																			a__					   =		int(flatten_latent_dim / 64		)
															elif flatten_latent_dim % 16 == 0:
																			a__					   =		int(flatten_latent_dim / 16		)
															elif flatten_latent_dim % 4 == 0:
																			a__					   =		int(flatten_latent_dim / 4		)
															else:
																			a__					   =		int(a_		)

															a__					   =		jax_memory_efficient_attention(
															    a_	,	a_	,	a_	,	query_chunk_size=a_	,	key_chunk_size=4_096 * 4		)

															a__					   =		hidden_states.transpose(1	,	0	,	2		)
											else:
															# compute attentions
															a__					   =		jnp.einsum("""b i d, b j d->b i j"""	,	a_	,	a_		)
															a__					   =		attention_scores * self.scale
															a__					   =		nn.softmax(a_	,	axis=2		)

															# attend to values
															a__					   =		jnp.einsum("""b i j, b j d -> b i d"""	,	a_	,	a_		)

											a__					   =		self.reshape_batch_dim_to_heads(a_		)
											a__					   =		self.proj_attn(a_		)
											return self.dropout_layer(a_	,	deterministic=a_		)




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa
							UpperCamelCase__					:   bool       =       False


							def _a   (				self		):
											# self attention (or cross_attention if only_cross_attention is True)
											a__					   =		FlaxAttention(
											    self.dim	,	self.n_heads	,	self.d_head	,	self.dropout	,	self.use_memory_efficient_attention	,	dtype=self.dtype		)
											# cross attention
											a__					   =		FlaxAttention(
											    self.dim	,	self.n_heads	,	self.d_head	,	self.dropout	,	self.use_memory_efficient_attention	,	dtype=self.dtype		)
											a__					   =		FlaxFeedForward(dim=self.dim	,	dropout=self.dropout	,	dtype=self.dtype		)
											a__					   =		nn.LayerNorm(epsilon=1E-5	,	dtype=self.dtype		)
											a__					   =		nn.LayerNorm(epsilon=1E-5	,	dtype=self.dtype		)
											a__					   =		nn.LayerNorm(epsilon=1E-5	,	dtype=self.dtype		)
											a__					   =		nn.Dropout(rate=self.dropout		)







							def __call__(				self	,	a_	,	a_	,	a_=True		):
											# self attention
											a__					   =		hidden_states
											if self.only_cross_attention:
															a__					   =		self.attna(self.norma(a_		)	,	a_	,	deterministic=a_		)
											else:
															a__					   =		self.attna(self.norma(a_		)	,	deterministic=a_		)
											a__					   =		hidden_states + residual

											# cross attention
											a__					   =		hidden_states
											a__					   =		self.attna(self.norma(a_		)	,	a_	,	deterministic=a_		)
											a__					   =		hidden_states + residual

											# feed forward
											a__					   =		hidden_states
											a__					   =		self.ff(self.norma(a_		)	,	deterministic=a_		)
											a__					   =		hidden_states + residual

											return self.dropout_layer(a_	,	deterministic=a_		)




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   int       =       1
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa
							UpperCamelCase__					:   bool       =       False


							def _a   (				self		):
											a__					   =		nn.GroupNorm(num_groups=32	,	epsilon=1E-5		)

											a__					   =		self.n_heads * self.d_head
											if self.use_linear_projection:
															a__					   =		nn.Dense(a_	,	dtype=self.dtype		)
											else:
															a__					   =		nn.Conv(
															    a_	,	kernel_size=(1, 1)	,	strides=(1, 1)	,	padding="""VALID"""	,	dtype=self.dtype	,	)

											a__					   =		[
											    FlaxBasicTransformerBlock(
											        a_	,	self.n_heads	,	self.d_head	,	dropout=self.dropout	,	only_cross_attention=self.only_cross_attention	,	dtype=self.dtype	,	use_memory_efficient_attention=self.use_memory_efficient_attention	,	)
											    for _ in range(self.depth		)
											]

											if self.use_linear_projection:
															a__					   =		nn.Dense(a_	,	dtype=self.dtype		)
											else:
															a__					   =		nn.Conv(
															    a_	,	kernel_size=(1, 1)	,	strides=(1, 1)	,	padding="""VALID"""	,	dtype=self.dtype	,	)

											a__					   =		nn.Dropout(rate=self.dropout		)







							def __call__(				self	,	a_	,	a_	,	a_=True		):
											a__   ,      a__   ,      a__   ,      a__					   =		hidden_states.shape
											a__					   =		hidden_states
											a__					   =		self.norm(a_		)
											if self.use_linear_projection:
															a__					   =		hidden_states.reshape(a_	,	height * width	,	a_		)
															a__					   =		self.proj_in(a_		)
											else:
															a__					   =		self.proj_in(a_		)
															a__					   =		hidden_states.reshape(a_	,	height * width	,	a_		)

											for transformer_block in self.transformer_blocks:
															a__					   =		transformer_block(a_	,	a_	,	deterministic=a_		)

											if self.use_linear_projection:
															a__					   =		self.proj_out(a_		)
															a__					   =		hidden_states.reshape(a_	,	a_	,	a_	,	a_		)
											else:
															a__					   =		hidden_states.reshape(a_	,	a_	,	a_	,	a_		)
															a__					   =		self.proj_out(a_		)

											a__					   =		hidden_states + residual
											return self.dropout_layer(a_	,	deterministic=a_		)




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa


							def _a   (				self		):
											# The second linear layer needs to be called
											# net_2 for now to match the index of the Sequential layer
											a__					   =		FlaxGEGLU(self.dim	,	self.dropout	,	self.dtype		)
											a__					   =		nn.Dense(self.dim	,	dtype=self.dtype		)







							def __call__(				self	,	a_	,	a_=True		):
											a__					   =		self.net_a(a_	,	deterministic=a_		)
											a__					   =		self.net_a(a_		)
											return hidden_states




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa


							def _a   (				self		):
											a__					   =		self.dim * 4
											a__					   =		nn.Dense(inner_dim * 2	,	dtype=self.dtype		)
											a__					   =		nn.Dropout(rate=self.dropout		)







							def __call__(				self	,	a_	,	a_=True		):
											a__					   =		self.proj(a_		)
											a__   ,      a__					   =		jnp.split(a_	,	2	,	axis=2		)
											return self.dropout_layer(hidden_linear * nn.gelu(a_		)	,	deterministic=a_		)

