



"""simple docstring"""
import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)



def        lowerCAmelCase_(   lowercase_     :							Tuple						,   lowercase_     :							List[str]						,   lowercase_     :							Dict							)	->       Union[str, Any]:
				if isinstance(lowercase_						,   torch.Tensor							):
								return image
				elif isinstance(lowercase_						,   PIL.Image.Image							):
								_lowerCamelCase        =				[image]

				if isinstance(image[0]						,   PIL.Image.Image							):
								_lowerCamelCase        =				[np.array(i.resize((w, h)						,   resample=PIL_INTERPOLATION['''lanczos''']							)							)[None, :] for i in image]
								_lowerCamelCase        =				np.concatenate(lowercase_						,   axis=0							)
								_lowerCamelCase        =				np.array(lowercase_							).astype(np.floataa							) / 2_5_5.0
								_lowerCamelCase        =				image.transpose(0						,   3						,   1						,   2							)
								_lowerCamelCase        =				2.0 * image - 1.0
								_lowerCamelCase        =				torch.from_numpy(lowercase_							)
				elif isinstance(image[0]						,   torch.Tensor							):
								_lowerCamelCase        =				torch.cat(lowercase_						,   dim=0							)
				return image



def        lowerCAmelCase_(   lowercase_     :							Dict						,   lowercase_     :							List[Any]						,   lowercase_     :							Union[str, Any]						,   lowercase_     :							Any=0.9_9_9_5							)	->       Optional[Any]:
				if not isinstance(lowercase_						,   np.ndarray							):
								_lowerCamelCase        =				True
								_lowerCamelCase        =				va.device
								_lowerCamelCase        =				va.cpu().numpy()
								_lowerCamelCase        =				va.cpu().numpy()

				_lowerCamelCase        =				np.sum(va * va / (np.linalg.norm(lowercase_							) * np.linalg.norm(lowercase_							))							)
				if np.abs(lowercase_							) > DOT_THRESHOLD:
								_lowerCamelCase        =				(1 - t) * va + t * va
				else:
								_lowerCamelCase        =				np.arccos(lowercase_							)
								_lowerCamelCase        =				np.sin(lowercase_							)
								_lowerCamelCase        =				theta_a * t
								_lowerCamelCase        =				np.sin(lowercase_							)
								_lowerCamelCase        =				np.sin(theta_a - theta_t							) / sin_theta_a
								_lowerCamelCase        =				sin_theta_t / sin_theta_a
								_lowerCamelCase        =				sa * va + sa * va

				if inputs_are_torch:
								_lowerCamelCase        =				torch.from_numpy(lowercase_							).to(lowercase_							)

				return va



def        lowerCAmelCase_(   lowercase_     :							Optional[int]						,   lowercase_     :							Dict							)	->       str:
				_lowerCamelCase        =				F.normalize(lowercase_						,   dim=-1							)
				_lowerCamelCase        =				F.normalize(lowercase_						,   dim=-1							)
				return (x - y).norm(dim=-1							).div(2							).arcsin().pow(2							).mul(2							)



def        lowerCAmelCase_(   lowercase_     :							List[str]						,   lowercase_     :							Union[str, Any]							)	->       Optional[int]:
				for param in model.parameters():
								_lowerCamelCase        =				value




class 							lowerCamelCase_(   A__    ):





					'''simple docstring'''



					def __init__(    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__=None				,       lowerCamelCase__=None				,       lowerCamelCase__=None				,       ):
									super().__init__()
									self.register_modules(
									    vae=lowerCamelCase__				,       text_encoder=lowerCamelCase__				,       clip_model=lowerCamelCase__				,       tokenizer=lowerCamelCase__				,       unet=lowerCamelCase__				,       scheduler=lowerCamelCase__				,       feature_extractor=lowerCamelCase__				,       coca_model=lowerCamelCase__				,       coca_tokenizer=lowerCamelCase__				,       coca_transform=lowerCamelCase__				,       )
									_lowerCamelCase        =				(
									    feature_extractor.size
									    if isinstance(feature_extractor.size				,       lowerCamelCase__   )
									    else feature_extractor.size['''shortest_edge''']
									)
									_lowerCamelCase        =				transforms.Normalize(mean=feature_extractor.image_mean				,       std=feature_extractor.image_std   )
									set_requires_grad(self.text_encoder				,       lowerCamelCase__   )
									set_requires_grad(self.clip_model				,       lowerCamelCase__   )

					def 			snake_case__      (    self				,       lowerCamelCase__ = "auto"   ):
									if slice_size == "auto":
													# half the attention head size is usually a good trade-off between
													# speed and memory
													_lowerCamelCase        =				self.unet.config.attention_head_dim // 2
									self.unet.set_attention_slice(lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									self.enable_attention_slicing(lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									set_requires_grad(self.vae				,       lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									set_requires_grad(self.vae				,       lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									set_requires_grad(self.unet				,       lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									set_requires_grad(self.unet				,       lowerCamelCase__   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									# get the original timestep using init_timestep
									_lowerCamelCase        =				min(int(num_inference_steps * strength   )				,       lowerCamelCase__   )

									_lowerCamelCase        =				max(num_inference_steps - init_timestep				,       0   )
									_lowerCamelCase        =				self.scheduler.timesteps[t_start:]

									return timesteps, num_inference_steps - t_start

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__=None   ):
									if not isinstance(lowerCamelCase__				,       torch.Tensor   ):
													raise ValueError(F"""`image` has to be of type `torch.Tensor` but is {type(lowerCamelCase__   )}"""   )

									_lowerCamelCase        =				image.to(device=lowerCamelCase__				,       dtype=lowerCamelCase__   )

									if isinstance(lowerCamelCase__				,       lowerCamelCase__   ):
													_lowerCamelCase        =				[
													    self.vae.encode(image[i : i + 1]   ).latent_dist.sample(generator[i]   ) for i in range(lowerCamelCase__   )
													]
													_lowerCamelCase        =				torch.cat(lowerCamelCase__				,       dim=0   )
									else:
													_lowerCamelCase        =				self.vae.encode(lowerCamelCase__   ).latent_dist.sample(lowerCamelCase__   )

									# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
									_lowerCamelCase        =				0.1_8_2_1_5 * init_latents
									_lowerCamelCase        =				init_latents.repeat_interleave(lowerCamelCase__				,       dim=0   )

									_lowerCamelCase        =				randn_tensor(init_latents.shape				,       generator=lowerCamelCase__				,       device=lowerCamelCase__				,       dtype=lowerCamelCase__   )

									# get latents
									_lowerCamelCase        =				self.scheduler.add_noise(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )
									_lowerCamelCase        =				init_latents

									return latents

					def 			snake_case__      (    self				,       lowerCamelCase__   ):
									_lowerCamelCase        =				self.coca_transform(lowerCamelCase__   ).unsqueeze(0   )
									with torch.no_grad(), torch.cuda.amp.autocast():
													_lowerCamelCase        =				self.coca_model.generate(transformed_image.to(device=self.device				,       dtype=self.coca_model.dtype   )   )
									_lowerCamelCase        =				self.coca_tokenizer.decode(generated[0].cpu().numpy()   )
									return generated.split('''<end_of_text>'''   )[0].replace('''<start_of_text>'''				,       ''''''   ).rstrip(''' .,'''   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__   ):
									_lowerCamelCase        =				self.feature_extractor.preprocess(lowerCamelCase__   )
									_lowerCamelCase        =				torch.from_numpy(clip_image_input['''pixel_values'''][0]   ).unsqueeze(0   ).to(self.device   ).half()
									_lowerCamelCase        =				self.clip_model.get_image_features(lowerCamelCase__   )
									_lowerCamelCase        =				image_embeddings_clip / image_embeddings_clip.norm(p=2				,       dim=-1				,       keepdim=lowerCamelCase__   )
									_lowerCamelCase        =				image_embeddings_clip.repeat_interleave(lowerCamelCase__				,       dim=0   )
									return image_embeddings_clip

					@torch.enable_grad()
					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       ):
									_lowerCamelCase        =				latents.detach().requires_grad_()

									_lowerCamelCase        =				self.scheduler.scale_model_input(lowerCamelCase__				,       lowerCamelCase__   )

									# predict the noise residual
									_lowerCamelCase        =				self.unet(lowerCamelCase__				,       lowerCamelCase__				,       encoder_hidden_states=lowerCamelCase__   ).sample

									if isinstance(self.scheduler				,       (PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)   ):
													_lowerCamelCase        =				self.scheduler.alphas_cumprod[timestep]
													_lowerCamelCase        =				1 - alpha_prod_t
													# compute predicted original sample from predicted noise also called
													# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
													_lowerCamelCase        =				(latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

													_lowerCamelCase        =				torch.sqrt(lowerCamelCase__   )
													_lowerCamelCase        =				pred_original_sample * (fac) + latents * (1 - fac)
									elif isinstance(self.scheduler				,       lowerCamelCase__   ):
													_lowerCamelCase        =				self.scheduler.sigmas[index]
													_lowerCamelCase        =				latents - sigma * noise_pred
									else:
													raise ValueError(F"""scheduler type {type(self.scheduler   )} not supported"""   )

									# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
									_lowerCamelCase        =				1 / 0.1_8_2_1_5 * sample
									_lowerCamelCase        =				self.vae.decode(lowerCamelCase__   ).sample
									_lowerCamelCase        =				(image / 2 + 0.5).clamp(0				,       1   )

									_lowerCamelCase        =				transforms.Resize(self.feature_extractor_size   )(lowerCamelCase__   )
									_lowerCamelCase        =				self.normalize(lowerCamelCase__   ).to(latents.dtype   )

									_lowerCamelCase        =				self.clip_model.get_image_features(lowerCamelCase__   )
									_lowerCamelCase        =				image_embeddings_clip / image_embeddings_clip.norm(p=2				,       dim=-1				,       keepdim=lowerCamelCase__   )

									_lowerCamelCase        =				spherical_dist_loss(lowerCamelCase__				,       lowerCamelCase__   ).mean() * clip_guidance_scale

									_lowerCamelCase        =				-torch.autograd.grad(lowerCamelCase__				,       lowerCamelCase__   )[0]

									if isinstance(self.scheduler				,       lowerCamelCase__   ):
													_lowerCamelCase        =				latents.detach() + grads * (sigma**2)
													_lowerCamelCase        =				noise_pred_original
									else:
													_lowerCamelCase        =				noise_pred_original - torch.sqrt(lowerCamelCase__   ) * grads
									return noise_pred, latents





					@torch.no_grad()
					def __call__(    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__ = None				,       lowerCamelCase__ = None				,       lowerCamelCase__ = 5_1_2				,       lowerCamelCase__ = 5_1_2				,       lowerCamelCase__ = 0.6				,       lowerCamelCase__ = 5_0				,       lowerCamelCase__ = 7.5				,       lowerCamelCase__ = 1				,       lowerCamelCase__ = 0.0				,       lowerCamelCase__ = 1_0_0				,       lowerCamelCase__ = None				,       lowerCamelCase__ = "pil"				,       lowerCamelCase__ = True				,       lowerCamelCase__ = 0.8				,       lowerCamelCase__ = 0.1				,       lowerCamelCase__ = 0.1				,       ):
									if isinstance(lowerCamelCase__				,       lowerCamelCase__   ) and len(lowerCamelCase__   ) != batch_size:
													raise ValueError(F"""You have passed {batch_size} batch_size, but only {len(lowerCamelCase__   )} generators."""   )

									if height % 8 != 0 or width % 8 != 0:
													raise ValueError(F"""`height` and `width` have to be divisible by 8 but are {height} and {width}."""   )

									if isinstance(lowerCamelCase__				,       torch.Generator   ) and batch_size > 1:
													_lowerCamelCase        =				[generator] + [None] * (batch_size - 1)

									_lowerCamelCase        =				[
									    ('''model''', self.coca_model is None),
									    ('''tokenizer''', self.coca_tokenizer is None),
									    ('''transform''', self.coca_transform is None),
									]
									_lowerCamelCase        =				[x[0] for x in coca_is_none if x[1]]
									_lowerCamelCase        =				''', '''.join(lowerCamelCase__   )
									# generate prompts with coca model if prompt is None
									if content_prompt is None:
													if len(lowerCamelCase__   ):
																	raise ValueError(
																	    F"""Content prompt is None and CoCa [{coca_is_none_str}] is None."""
																	    F"""Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline."""   )
													_lowerCamelCase        =				self.get_image_description(lowerCamelCase__   )
									if style_prompt is None:
													if len(lowerCamelCase__   ):
																	raise ValueError(
																	    F"""Style prompt is None and CoCa [{coca_is_none_str}] is None."""
																	    F""" Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline."""   )
													_lowerCamelCase        =				self.get_image_description(lowerCamelCase__   )

									# get prompt text embeddings for content and style
									_lowerCamelCase        =				self.tokenizer(
									    lowerCamelCase__				,       padding='''max_length'''				,       max_length=self.tokenizer.model_max_length				,       truncation=lowerCamelCase__				,       return_tensors='''pt'''				,       )
									_lowerCamelCase        =				self.text_encoder(content_text_input.input_ids.to(self.device   )   )[0]

									_lowerCamelCase        =				self.tokenizer(
									    lowerCamelCase__				,       padding='''max_length'''				,       max_length=self.tokenizer.model_max_length				,       truncation=lowerCamelCase__				,       return_tensors='''pt'''				,       )
									_lowerCamelCase        =				self.text_encoder(style_text_input.input_ids.to(self.device   )   )[0]

									_lowerCamelCase        =				slerp(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )

									# duplicate text embeddings for each generation per prompt
									_lowerCamelCase        =				text_embeddings.repeat_interleave(lowerCamelCase__				,       dim=0   )

									# set timesteps
									_lowerCamelCase        =				'''offset''' in set(inspect.signature(self.scheduler.set_timesteps   ).parameters.keys()   )
									_lowerCamelCase        =				{}
									if accepts_offset:
													_lowerCamelCase        =				1

									self.scheduler.set_timesteps(lowerCamelCase__				,       **lowerCamelCase__   )
									# Some schedulers like PNDM have timesteps as arrays
									# It's more optimized to move all timesteps to correct device beforehand
									self.scheduler.timesteps.to(self.device   )

									_lowerCamelCase  ,     _lowerCamelCase        =				self.get_timesteps(lowerCamelCase__				,       lowerCamelCase__				,       self.device   )
									_lowerCamelCase        =				timesteps[:1].repeat(lowerCamelCase__   )

									# Preprocess image
									_lowerCamelCase        =				preprocess(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )
									_lowerCamelCase        =				self.prepare_latents(
									    lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       text_embeddings.dtype				,       self.device				,       lowerCamelCase__   )

									_lowerCamelCase        =				preprocess(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )
									_lowerCamelCase        =				self.prepare_latents(
									    lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       text_embeddings.dtype				,       self.device				,       lowerCamelCase__   )

									_lowerCamelCase        =				slerp(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )

									if clip_guidance_scale > 0:
													_lowerCamelCase        =				self.get_clip_image_embeddings(lowerCamelCase__				,       lowerCamelCase__   )
													_lowerCamelCase        =				self.get_clip_image_embeddings(lowerCamelCase__				,       lowerCamelCase__   )
													_lowerCamelCase        =				slerp(
													    lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )

									# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
									# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
									# corresponds to doing no classifier free guidance.
									_lowerCamelCase        =				guidance_scale > 1.0
									# get unconditional embeddings for classifier free guidance
									if do_classifier_free_guidance:
													_lowerCamelCase        =				content_text_input.input_ids.shape[-1]
													_lowerCamelCase        =				self.tokenizer(['''''']				,       padding='''max_length'''				,       max_length=lowerCamelCase__				,       return_tensors='''pt'''   )
													_lowerCamelCase        =				self.text_encoder(uncond_input.input_ids.to(self.device   )   )[0]
													# duplicate unconditional embeddings for each generation per prompt
													_lowerCamelCase        =				uncond_embeddings.repeat_interleave(lowerCamelCase__				,       dim=0   )

													# For classifier free guidance, we need to do two forward passes.
													# Here we concatenate the unconditional and text embeddings into a single batch
													# to avoid doing two forward passes
													_lowerCamelCase        =				torch.cat([uncond_embeddings, text_embeddings]   )

									# get the initial random noise unless the user supplied it

									# Unlike in other pipelines, latents need to be generated in the target device
									# for 1-to-1 results reproducibility with the CompVis implementation.
									# However this currently doesn't work in `mps`.
									_lowerCamelCase        =				(batch_size, self.unet.config.in_channels, height // 8, width // 8)
									_lowerCamelCase        =				text_embeddings.dtype
									if latents is None:
													if self.device.type == "mps":
																	# randn does not work reproducibly on mps
																	_lowerCamelCase        =				torch.randn(lowerCamelCase__				,       generator=lowerCamelCase__				,       device='''cpu'''				,       dtype=lowerCamelCase__   ).to(
																	    self.device   )
													else:
																	_lowerCamelCase        =				torch.randn(lowerCamelCase__				,       generator=lowerCamelCase__				,       device=self.device				,       dtype=lowerCamelCase__   )
									else:
													if latents.shape != latents_shape:
																	raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""   )
													_lowerCamelCase        =				latents.to(self.device   )

									# scale the initial noise by the standard deviation required by the scheduler
									_lowerCamelCase        =				latents * self.scheduler.init_noise_sigma

									# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
									# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
									# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
									# and should be between [0, 1]
									_lowerCamelCase        =				'''eta''' in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
									_lowerCamelCase        =				{}
									if accepts_eta:
													_lowerCamelCase        =				eta

									# check if the scheduler accepts generator
									_lowerCamelCase        =				'''generator''' in set(inspect.signature(self.scheduler.step   ).parameters.keys()   )
									if accepts_generator:
													_lowerCamelCase        =				generator

									with self.progress_bar(total=lowerCamelCase__   ):
													for i, t in enumerate(lowerCamelCase__   ):
																	# expand the latents if we are doing classifier free guidance
																	_lowerCamelCase        =				torch.cat([latents] * 2   ) if do_classifier_free_guidance else latents
																	_lowerCamelCase        =				self.scheduler.scale_model_input(lowerCamelCase__				,       lowerCamelCase__   )

																	# predict the noise residual
																	_lowerCamelCase        =				self.unet(lowerCamelCase__				,       lowerCamelCase__				,       encoder_hidden_states=lowerCamelCase__   ).sample

																	# perform classifier free guidance
																	if do_classifier_free_guidance:
																					_lowerCamelCase  ,     _lowerCamelCase        =				noise_pred.chunk(2   )
																					_lowerCamelCase        =				noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

																	# perform clip guidance
																	if clip_guidance_scale > 0:
																					_lowerCamelCase        =				(
																					    text_embeddings.chunk(2   )[1] if do_classifier_free_guidance else text_embeddings
																					)
																					_lowerCamelCase  ,     _lowerCamelCase        =				self.cond_fn(
																					    lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       )

																	# compute the previous noisy sample x_t -> x_t-1
																	_lowerCamelCase        =				self.scheduler.step(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       **lowerCamelCase__   ).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
									_lowerCamelCase        =				1 / 0.1_8_2_1_5 * latents
									_lowerCamelCase        =				self.vae.decode(lowerCamelCase__   ).sample

									_lowerCamelCase        =				(image / 2 + 0.5).clamp(0				,       1   )
									_lowerCamelCase        =				image.cpu().permute(0				,       2				,       3				,       1   ).numpy()

									if output_type == "pil":
													_lowerCamelCase        =				self.numpy_to_pil(lowerCamelCase__   )

									if not return_dict:
													return (image, None)

									return StableDiffusionPipelineOutput(images=lowerCamelCase__				,       nsfw_content_detected=lowerCamelCase__   )





"""simple docstring"""
import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
			import jax

			from transformers.models.big_bird.modeling_flax_big_bird import (
			    FlaxBigBirdForCausalLM,
			    FlaxBigBirdForMaskedLM,
			    FlaxBigBirdForMultipleChoice,
			    FlaxBigBirdForPreTraining,
			    FlaxBigBirdForQuestionAnswering,
			    FlaxBigBirdForSequenceClassification,
			    FlaxBigBirdForTokenClassification,
			    FlaxBigBirdModel,
			)


class 							lowerCamelCase_(   unittest.TestCase    ):





					'''simple docstring'''



					def __init__(    self				,       lowerCamelCase__				,       lowerCamelCase__=2				,       lowerCamelCase__=5_6				,       lowerCamelCase__=True				,       lowerCamelCase__=True				,       lowerCamelCase__=True				,       lowerCamelCase__=True				,       lowerCamelCase__=9_9				,       lowerCamelCase__=3_2				,       lowerCamelCase__=2				,       lowerCamelCase__=2				,       lowerCamelCase__=7				,       lowerCamelCase__="gelu_new"				,       lowerCamelCase__=0.1				,       lowerCamelCase__=0.1				,       lowerCamelCase__=5_1_2				,       lowerCamelCase__=1_6				,       lowerCamelCase__=2				,       lowerCamelCase__=0.0_2				,       lowerCamelCase__=4				,       lowerCamelCase__="block_sparse"				,       lowerCamelCase__=True				,       lowerCamelCase__=False				,       lowerCamelCase__=2				,       lowerCamelCase__=3				,       ):
									_lowerCamelCase        =				parent
									_lowerCamelCase        =				batch_size
									_lowerCamelCase        =				seq_length
									_lowerCamelCase        =				is_training
									_lowerCamelCase        =				use_attention_mask
									_lowerCamelCase        =				use_token_type_ids
									_lowerCamelCase        =				use_labels
									_lowerCamelCase        =				vocab_size
									_lowerCamelCase        =				hidden_size
									_lowerCamelCase        =				num_hidden_layers
									_lowerCamelCase        =				num_attention_heads
									_lowerCamelCase        =				intermediate_size
									_lowerCamelCase        =				hidden_act
									_lowerCamelCase        =				hidden_dropout_prob
									_lowerCamelCase        =				attention_probs_dropout_prob
									_lowerCamelCase        =				max_position_embeddings
									_lowerCamelCase        =				type_vocab_size
									_lowerCamelCase        =				type_sequence_label_size
									_lowerCamelCase        =				initializer_range
									_lowerCamelCase        =				num_choices

									_lowerCamelCase        =				rescale_embeddings
									_lowerCamelCase        =				attention_type
									_lowerCamelCase        =				use_bias
									_lowerCamelCase        =				block_size
									_lowerCamelCase        =				num_random_blocks

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size   )

									_lowerCamelCase        =				None
									if self.use_attention_mask:
													_lowerCamelCase        =				random_attention_mask([self.batch_size, self.seq_length]   )

									_lowerCamelCase        =				None
									if self.use_token_type_ids:
													_lowerCamelCase        =				ids_tensor([self.batch_size, self.seq_length]				,       self.type_vocab_size   )

									_lowerCamelCase        =				BigBirdConfig(
									    vocab_size=self.vocab_size				,       hidden_size=self.hidden_size				,       num_hidden_layers=self.num_hidden_layers				,       num_attention_heads=self.num_attention_heads				,       intermediate_size=self.intermediate_size				,       hidden_act=self.hidden_act				,       hidden_dropout_prob=self.hidden_dropout_prob				,       attention_probs_dropout_prob=self.attention_probs_dropout_prob				,       max_position_embeddings=self.max_position_embeddings				,       type_vocab_size=self.type_vocab_size				,       is_decoder=lowerCamelCase__				,       initializer_range=self.initializer_range				,       attention_type=self.attention_type				,       block_size=self.block_size				,       num_random_blocks=self.num_random_blocks				,       use_bias=self.use_bias				,       rescale_embeddings=self.rescale_embeddings				,       )

									return config, input_ids, token_type_ids, attention_mask





					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.prepare_config_and_inputs()
									_lowerCamelCase  ,     _lowerCamelCase  ,     _lowerCamelCase  ,     _lowerCamelCase        =				config_and_inputs
									_lowerCamelCase        =				{
									    '''input_ids''': input_ids,
									    '''token_type_ids''': token_type_ids,
									    '''attention_mask''': attention_mask,
									}
									return config, inputs_dict


@require_flax
class 							lowerCamelCase_(   A__,				unittest.TestCase    ):





					'''simple docstring'''



					lowercase__   :       List[str]										= (
					    (
					        FlaxBigBirdForCausalLM,
					        FlaxBigBirdModel,
					        FlaxBigBirdForPreTraining,
					        FlaxBigBirdForMaskedLM,
					        FlaxBigBirdForMultipleChoice,
					        FlaxBigBirdForQuestionAnswering,
					        FlaxBigBirdForSequenceClassification,
					        FlaxBigBirdForTokenClassification,
					    )
					    if is_flax_available()
					    else ()
					)

					lowercase__   :       Any										= False
					lowercase__   :       Optional[int]										= False

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				FlaxBigBirdModelTester(self   )

					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def 			snake_case__      (    self   ):
									super().test_from_pretrained_save_pretrained()

					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def 			snake_case__      (    self   ):
									super().test_from_pretrained_with_no_automatic_init()

					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def 			snake_case__      (    self   ):
									super().test_no_automatic_init()

					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def 			snake_case__      (    self   ):
									super().test_hidden_states_output()

					@slow
					def 			snake_case__      (    self   ):
									for model_class_name in self.all_model_classes:
													_lowerCamelCase        =				model_class_name.from_pretrained('''google/bigbird-roberta-base'''   )
													self.assertIsNotNone(lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									if self.test_attn_probs:
													super().test_attention_outputs()

					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def 			snake_case__      (    self   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
													with self.subTest(model_class.__name__   ):
																	_lowerCamelCase        =				self._prepare_for_class(lowerCamelCase__				,       lowerCamelCase__   )
																	_lowerCamelCase        =				model_class(lowerCamelCase__   )

																	@jax.jit
																	def model_jitted(lowerCamelCase__				,       lowerCamelCase__=None				,       **lowerCamelCase__   ):
																					return model(input_ids=lowerCamelCase__				,       attention_mask=lowerCamelCase__				,       **lowerCamelCase__   )

																	with self.subTest('''JIT Enabled'''   ):
																					_lowerCamelCase        =				model_jitted(**lowerCamelCase__   ).to_tuple()

																	with self.subTest('''JIT Disabled'''   ):
																					with jax.disable_jit():
																									_lowerCamelCase        =				model_jitted(**lowerCamelCase__   ).to_tuple()

																	self.assertEqual(len(lowerCamelCase__   )				,       len(lowerCamelCase__   )   )
																	for jitted_output, output in zip(lowerCamelCase__				,       lowerCamelCase__   ):
																					self.assertEqual(jitted_output.shape				,       output.shape   )





					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__=1e-5				,       lowerCamelCase__="outputs"				,       lowerCamelCase__=None   ):
									# `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
									# an effort was done to return `attention_probs` (yet to be verified).
									if name.startswith('''outputs.attentions'''   ):
													return
									else:
													super().check_pt_flax_outputs(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )
