



"""simple docstring"""






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class 	__SCREAMING_SNAKE_CASE     (				lowerCAmelCase_		,       lowerCAmelCase_		,       lowerCAmelCase_		,       unittest.TestCase       ):


					'''simple docstring'''





					_a             =							StableDiffusionInstructPixaPixPipeline
					_a             =							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width', 'cross_attention_kwargs'}
					_a             =							TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
					_a             =							IMAGE_TO_IMAGE_IMAGE_PARAMS
					_a             =							IMAGE_TO_IMAGE_IMAGE_PARAMS



					def 		snake_case	(	self							:  Optional[Any]					)-> Dict:
								torch.manual_seed(0					)
								lowerCamelCase__		:		str		     =UNetaDConditionModel(
								    block_out_channels=(32, 64),  layers_per_block=2,  sample_size=32,  in_channels=8,  out_channels=4,  down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),  up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D'''),  cross_attention_dim=32,  )
								lowerCamelCase__		:		Optional[int]		     =PNDMScheduler(skip_prk_steps=lowerCamelCase					)
								torch.manual_seed(0					)
								lowerCamelCase__		:		Optional[Any]		     =AutoencoderKL(
								    block_out_channels=[32, 64],  in_channels=3,  out_channels=3,  down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],  up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],  latent_channels=4,  )
								torch.manual_seed(0					)
								lowerCamelCase__		:		Any		     =CLIPTextConfig(
								    bos_token_id=0,  eos_token_id=2,  hidden_size=32,  intermediate_size=37,  layer_norm_eps=1E-05,  num_attention_heads=4,  num_hidden_layers=5,  pad_token_id=1,  vocab_size=1000,  )
								lowerCamelCase__		:		Dict		     =CLIPTextModel(lowerCamelCase					)
								lowerCamelCase__		:		List[Any]		     =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)

								lowerCamelCase__		:		Union[str, Any]		     ={
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								    '''text_encoder''': text_encoder,
								    '''tokenizer''': tokenizer,
								    '''safety_checker''': None,
								    '''feature_extractor''': None,
								}
								return components



					def 		snake_case	(	self							:  Any,  lowerCamelCase							:  Tuple,  lowerCamelCase							:  str=0					)-> Dict:
								lowerCamelCase__		:		Tuple		     =floats_tensor((1, 3, 32, 32),  rng=random.Random(lowerCamelCase					)					).to(lowerCamelCase					)
								lowerCamelCase__		:		List[str]		     =image.cpu().permute(0,  2,  3,  1					)[0]
								lowerCamelCase__		:		str		     =Image.fromarray(np.uinta(lowerCamelCase					)					).convert('''RGB'''					)
								if str(lowerCamelCase					).startswith('''mps'''					):
											lowerCamelCase__		:		str		     =torch.manual_seed(lowerCamelCase					)
								else:
											lowerCamelCase__		:		Optional[int]		     =torch.Generator(device=lowerCamelCase					).manual_seed(lowerCamelCase					)
								lowerCamelCase__		:		Any		     ={
								    '''prompt''': '''A painting of a squirrel eating a burger''',
								    '''image''': image,
								    '''generator''': generator,
								    '''num_inference_steps''': 2,
								    '''guidance_scale''': 6.0,
								    '''image_guidance_scale''': 1,
								    '''output_type''': '''numpy''',
								}
								return inputs



					def 		snake_case	(	self							:  Optional[Any]					)-> Optional[int]:
								lowerCamelCase__		:		int		     ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__		:		int		     =self.get_dummy_components()
								lowerCamelCase__		:		Tuple		     =StableDiffusionInstructPixaPixPipeline(**lowerCamelCase					)
								lowerCamelCase__		:		List[Any]		     =sd_pipe.to(lowerCamelCase					)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase					)

								lowerCamelCase__		:		Any		     =self.get_dummy_inputs(lowerCamelCase					)
								lowerCamelCase__		:		int		     =sd_pipe(**lowerCamelCase					).images
								lowerCamelCase__		:		Dict		     =image[0, -3:, -3:, -1]
								assert image.shape == (1, 32, 32, 3)
								lowerCamelCase__		:		Union[str, Any]		     =np.array([0.7_526, 0.3_750, 0.4_547, 0.6_117, 0.5_866, 0.5_016, 0.4_327, 0.5_642, 0.4_815]					)

								assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3



					def 		snake_case	(	self							:  Optional[int]					)-> str:
								lowerCamelCase__		:		int		     ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__		:		Optional[Any]		     =self.get_dummy_components()
								lowerCamelCase__		:		List[Any]		     =StableDiffusionInstructPixaPixPipeline(**lowerCamelCase					)
								lowerCamelCase__		:		Optional[Any]		     =sd_pipe.to(lowerCamelCase					)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase					)

								lowerCamelCase__		:		Optional[int]		     =self.get_dummy_inputs(lowerCamelCase					)
								lowerCamelCase__		:		str		     ='''french fries'''
								lowerCamelCase__		:		int		     =sd_pipe(**lowerCamelCase,  negative_prompt=lowerCamelCase					)
								lowerCamelCase__		:		List[Any]		     =output.images
								lowerCamelCase__		:		Union[str, Any]		     =image[0, -3:, -3:, -1]

								assert image.shape == (1, 32, 32, 3)
								lowerCamelCase__		:		List[Any]		     =np.array([0.7_511, 0.3_642, 0.4_553, 0.6_236, 0.5_797, 0.5_013, 0.4_343, 0.5_611, 0.4_831]					)

								assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3



					def 		snake_case	(	self							:  Union[str, Any]					)-> Union[str, Any]:
								lowerCamelCase__		:		List[Any]		     ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__		:		Union[str, Any]		     =self.get_dummy_components()
								lowerCamelCase__		:		Any		     =StableDiffusionInstructPixaPixPipeline(**lowerCamelCase					)
								lowerCamelCase__		:		Union[str, Any]		     =sd_pipe.to(lowerCamelCase					)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase					)

								lowerCamelCase__		:		Optional[Any]		     =self.get_dummy_inputs(lowerCamelCase					)
								lowerCamelCase__		:		Optional[Any]		     =[inputs['''prompt''']] * 2

								lowerCamelCase__		:		Tuple		     =np.array(inputs['''image''']					).astype(np.floataa					) / 255.0
								lowerCamelCase__		:		int		     =torch.from_numpy(lowerCamelCase					).unsqueeze(0					).to(lowerCamelCase					)
								lowerCamelCase__		:		Any		     =image / 2 + 0.5
								lowerCamelCase__		:		List[str]		     =image.permute(0,  3,  1,  2					)
								lowerCamelCase__		:		str		     =image.repeat(2,  1,  1,  1					)

								lowerCamelCase__		:		Optional[Any]		     =sd_pipe(**lowerCamelCase					).images
								lowerCamelCase__		:		Any		     =image[-1, -3:, -3:, -1]

								assert image.shape == (2, 32, 32, 3)
								lowerCamelCase__		:		Dict		     =np.array([0.5_812, 0.5_748, 0.5_222, 0.5_908, 0.5_695, 0.7_174, 0.6_804, 0.5_523, 0.5_579]					)

								assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3



					def 		snake_case	(	self							:  List[str]					)-> Dict:
								lowerCamelCase__		:		Optional[int]		     ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase__		:		Optional[Any]		     =self.get_dummy_components()
								lowerCamelCase__		:		List[Any]		     =EulerAncestralDiscreteScheduler(
								    beta_start=0.00_085,  beta_end=0.012,  beta_schedule='''scaled_linear'''					)
								lowerCamelCase__		:		Dict		     =StableDiffusionInstructPixaPixPipeline(**lowerCamelCase					)
								lowerCamelCase__		:		Optional[int]		     =sd_pipe.to(lowerCamelCase					)
								sd_pipe.set_progress_bar_config(disable=lowerCamelCase					)

								lowerCamelCase__		:		str		     =self.get_dummy_inputs(lowerCamelCase					)
								lowerCamelCase__		:		Optional[Any]		     =sd_pipe(**lowerCamelCase					).images
								lowerCamelCase__		:		Union[str, Any]		     =image[0, -3:, -3:, -1]

								lowerCamelCase__		:		Tuple		     =[round(lowerCamelCase,  4					) for x in image_slice.flatten().tolist()]
								print(''','''.join([str(lowerCamelCase					) for x in slice]					)					)

								assert image.shape == (1, 32, 32, 3)
								lowerCamelCase__		:		int		     =np.array([0.7_417, 0.3_842, 0.4_732, 0.5_776, 0.5_891, 0.5_139, 0.4_052, 0.5_673, 0.4_986]					)

								assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3



					def 		snake_case	(	self							:  int					)-> Optional[Any]:
								super().test_inference_batch_single_identical(expected_max_diff=3E-3					)




					def 		snake_case	(	self							:  Any					)-> int:
								lowerCamelCase__		:		List[str]		     =self.get_dummy_components()
								lowerCamelCase__		:		Any		     =StableDiffusionInstructPixaPixPipeline(**lowerCamelCase					)
								lowerCamelCase__		:		Tuple		     =VaeImageProcessor(do_resize=lowerCamelCase,  do_normalize=lowerCamelCase					)
								lowerCamelCase__		:		List[str]		     =pipe.to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)

								lowerCamelCase__		:		str		     =pipe(**self.get_dummy_inputs_by_type(lowerCamelCase,  input_image_type='''pt'''					)					)[0]

								lowerCamelCase__		:		Any		     =components['''vae''']
								lowerCamelCase__		:		int		     =self.get_dummy_inputs_by_type(lowerCamelCase,  input_image_type='''pt'''					)

								for image_param in self.image_latents_params:
											if image_param in inputs.keys():
														lowerCamelCase__		:		List[str]		     =vae.encode(inputs[image_param]					).latent_dist.mode()

								lowerCamelCase__		:		int		     =pipe(**lowerCamelCase					)[0]

								lowerCamelCase__		:		List[str]		     =np.abs(out - out_latents_inputs					).max()
								self.assertLess(lowerCamelCase,  1E-4,  '''passing latents as image input generate different result from passing image'''					)





@slow
@require_torch_gpu
class 	__SCREAMING_SNAKE_CASE     (				unittest.TestCase       ):


					'''simple docstring'''





					def 		snake_case	(	self							:  Tuple					)-> List[str]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



					def 		snake_case	(	self							:  Dict,  lowerCamelCase							:  Optional[int]=0					)-> str:
								lowerCamelCase__		:		Any		     =torch.manual_seed(lowerCamelCase					)
								lowerCamelCase__		:		Any		     =load_image(
								    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'''					)
								lowerCamelCase__		:		Union[str, Any]		     ={
								    '''prompt''': '''turn him into a cyborg''',
								    '''image''': image,
								    '''generator''': generator,
								    '''num_inference_steps''': 3,
								    '''guidance_scale''': 7.5,
								    '''image_guidance_scale''': 1.0,
								    '''output_type''': '''numpy''',
								}
								return inputs



					def 		snake_case	(	self							:  Tuple					)-> Optional[int]:
								lowerCamelCase__		:		Tuple		     =StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    '''timbrooks/instruct-pix2pix''',  safety_checker=lowerCamelCase					)
								pipe.to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)
								pipe.enable_attention_slicing()

								lowerCamelCase__		:		Any		     =self.get_inputs()
								lowerCamelCase__		:		Union[str, Any]		     =pipe(**lowerCamelCase					).images
								lowerCamelCase__		:		Any		     =image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								lowerCamelCase__		:		str		     =np.array([0.5_902, 0.6_015, 0.6_027, 0.5_983, 0.6_092, 0.6_061, 0.5_765, 0.5_785, 0.5_555]					)

								assert np.abs(expected_slice - image_slice					).max() < 1E-3



					def 		snake_case	(	self							:  Optional[int]					)-> Optional[int]:
								lowerCamelCase__		:		Dict		     =StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    '''timbrooks/instruct-pix2pix''',  safety_checker=lowerCamelCase					)
								lowerCamelCase__		:		Union[str, Any]		     =LMSDiscreteScheduler.from_config(pipe.scheduler.config					)
								pipe.to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)
								pipe.enable_attention_slicing()

								lowerCamelCase__		:		Tuple		     =self.get_inputs()
								lowerCamelCase__		:		Optional[Any]		     =pipe(**lowerCamelCase					).images
								lowerCamelCase__		:		Dict		     =image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								lowerCamelCase__		:		List[Any]		     =np.array([0.6_578, 0.6_817, 0.6_972, 0.6_761, 0.6_856, 0.6_916, 0.6_428, 0.6_516, 0.6_301]					)

								assert np.abs(expected_slice - image_slice					).max() < 1E-3



					def 		snake_case	(	self							:  int					)-> str:
								lowerCamelCase__		:		List[Any]		     =StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    '''timbrooks/instruct-pix2pix''',  safety_checker=lowerCamelCase					)
								lowerCamelCase__		:		Optional[int]		     =DDIMScheduler.from_config(pipe.scheduler.config					)
								pipe.to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)
								pipe.enable_attention_slicing()

								lowerCamelCase__		:		List[str]		     =self.get_inputs()
								lowerCamelCase__		:		Dict		     =pipe(**lowerCamelCase					).images
								lowerCamelCase__		:		Tuple		     =image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								lowerCamelCase__		:		List[Any]		     =np.array([0.3_828, 0.3_834, 0.3_818, 0.3_792, 0.3_865, 0.3_752, 0.3_792, 0.3_847, 0.3_753]					)

								assert np.abs(expected_slice - image_slice					).max() < 1E-3



					def 		snake_case	(	self							:  List[str]					)-> Any:
								lowerCamelCase__		:		Union[str, Any]		     =0

								def callback_fn(lowerCamelCase							:  int,  lowerCamelCase							:  int,  lowerCamelCase							:  torch.FloatTensor					) -> None:
											lowerCamelCase__		:		List[Any]		     =True
											nonlocal number_of_steps
											number_of_steps += 1
											if step == 1:
														lowerCamelCase__		:		Union[str, Any]		     =latents.detach().cpu().numpy()
														assert latents.shape == (1, 4, 64, 64)
														lowerCamelCase__		:		str		     =latents[0, -3:, -3:, -1]
														lowerCamelCase__		:		Optional[Any]		     =np.array([-0.2_463, -0.4_644, -0.9_756, 1.5_176, 1.4_414, 0.7_866, 0.9_897, 0.8_521, 0.7_983]					)

														assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5E-2
											elif step == 2:
														lowerCamelCase__		:		Dict		     =latents.detach().cpu().numpy()
														assert latents.shape == (1, 4, 64, 64)
														lowerCamelCase__		:		Optional[Any]		     =latents[0, -3:, -3:, -1]
														lowerCamelCase__		:		Any		     =np.array([-0.2_644, -0.4_626, -0.9_653, 1.5_176, 1.4_551, 0.7_686, 0.9_805, 0.8_452, 0.8_115]					)

														assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5E-2

								lowerCamelCase__		:		Union[str, Any]		     =False

								lowerCamelCase__		:		int		     =StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    '''timbrooks/instruct-pix2pix''',  safety_checker=lowerCamelCase,  torch_dtype=torch.floataa					)
								lowerCamelCase__		:		Union[str, Any]		     =pipe.to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)
								pipe.enable_attention_slicing()

								lowerCamelCase__		:		Optional[int]		     =self.get_inputs()
								pipe(**lowerCamelCase,  callback=lowerCamelCase,  callback_steps=1					)
								assert callback_fn.has_been_called
								assert number_of_steps == 3



					def 		snake_case	(	self							:  Dict					)-> Optional[int]:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								lowerCamelCase__		:		Any		     =StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    '''timbrooks/instruct-pix2pix''',  safety_checker=lowerCamelCase,  torch_dtype=torch.floataa					)
								lowerCamelCase__		:		int		     =pipe.to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)
								pipe.enable_attention_slicing(1					)
								pipe.enable_sequential_cpu_offload()

								lowerCamelCase__		:		str		     =self.get_inputs()
								lowerCamelCase__		:		Dict		     =pipe(**lowerCamelCase					)

								lowerCamelCase__		:		List[str]		     =torch.cuda.max_memory_allocated()
								# make sure that less than 2.2 GB is allocated
								assert mem_bytes < 2.2 * 10**9




					def 		snake_case	(	self							:  Optional[int]					)-> Tuple:
								lowerCamelCase__		:		Tuple		     =self.get_inputs()
								# resize to resolution that is divisible by 8 but not 16 or 32
								lowerCamelCase__		:		Optional[Any]		     =inputs['''image'''].resize((504, 504)					)

								lowerCamelCase__		:		Any		     ='''timbrooks/instruct-pix2pix'''
								lowerCamelCase__		:		Dict		     =StableDiffusionInstructPixaPixPipeline.from_pretrained(
								    lowerCamelCase,  safety_checker=lowerCamelCase,  )
								pipe.to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)
								pipe.enable_attention_slicing()

								lowerCamelCase__		:		List[Any]		     =pipe(**lowerCamelCase					)
								lowerCamelCase__		:		str		     =output.images[0]

								lowerCamelCase__		:		Any		     =image[255:258, 383:386, -1]

								assert image.shape == (504, 504, 3)
								lowerCamelCase__		:		Optional[Any]		     =np.array([0.2_726, 0.2_529, 0.2_664, 0.2_655, 0.2_641, 0.2_642, 0.2_591, 0.2_649, 0.2_590]					)

								assert np.abs(image_slice.flatten() - expected_slice					).max() < 5E-3









"""simple docstring"""






from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


_lowercase			:			List[str]										=						logging.get_logger(__name__)






def 					snake_case__       (		__lowerCamelCase    :   Union[tf.Tensor, np.ndarray]     ):
			"""simple docstring"""
			if isinstance(__lowerCamelCase							,		np.ndarray     ):
						return list(tensor.shape     )

			lowerCamelCase__		:		List[Any]		     =tf.shape(__lowerCamelCase     )

			if tensor.shape == tf.TensorShape(__lowerCamelCase     ):
						return dynamic

			lowerCamelCase__		:		List[str]		     =tensor.shape.as_list()

			return [dynamic[i] if s is None else s for i, s in enumerate(__lowerCamelCase     )]






def 					snake_case__       (		__lowerCamelCase    :   tf.Tensor							,		__lowerCamelCase    :   Optional[int] = None							,		__lowerCamelCase    :   Optional[str] = None     ):
			"""simple docstring"""
			return tf.nn.softmax(logits=logits + 1e-9							,		axis=__lowerCamelCase							,		name=__lowerCamelCase     )






def 					snake_case__       (		__lowerCamelCase    :   Optional[Any]							,		__lowerCamelCase    :   str							,		__lowerCamelCase    :   Optional[int]							,		__lowerCamelCase    :   Any=1e-5							,		__lowerCamelCase    :   str=-1     ):
			"""simple docstring"""
			# This is a very simplified functional layernorm, designed to duplicate
			# the functionality of PyTorch nn.functional.layer_norm when this is needed to port
			# models in Transformers.

			if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(__lowerCamelCase							,		__lowerCamelCase     ):
						raise NotImplementedError('''Only 1D weight and bias tensors are supported for now, with only a single axis.'''     )

			# Get mean and variance on the axis to be normalized
			lowerCamelCase__     ,       lowerCamelCase__		:		Optional[int]		     =tf.nn.moments(__lowerCamelCase							,		axes=[axis]							,		keepdims=__lowerCamelCase     )

			if axis != -1:
						# Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
						# on every dimension except axis
						lowerCamelCase__		:		Optional[Any]		     =[1] * inputs.shape.rank
						lowerCamelCase__		:		Union[str, Any]		     =shape_list(__lowerCamelCase     )[axis]
						lowerCamelCase__		:		Optional[Any]		     =tf.reshape(__lowerCamelCase							,		__lowerCamelCase     )
						lowerCamelCase__		:		Any		     =tf.reshape(__lowerCamelCase							,		__lowerCamelCase     )

			# Compute layer normalization using the batch_normalization
			# function.
			lowerCamelCase__		:		List[str]		     =tf.nn.batch_normalization(
			    __lowerCamelCase							,		__lowerCamelCase							,		__lowerCamelCase							,		offset=__lowerCamelCase							,		scale=__lowerCamelCase							,		variance_epsilon=__lowerCamelCase							,		)
			return outputs






def 					snake_case__       (		__lowerCamelCase    :   Any							,		__lowerCamelCase    :   List[Any]=0							,		__lowerCamelCase    :   int=-1     ):
			"""simple docstring"""
			# Replicates the behavior of torch.flatten in TF

			# If end_dim or start_dim is negative, count them from the end
			if end_dim < 0:
						end_dim += input.shape.rank
			if start_dim < 0:
						start_dim += input.shape.rank

			if start_dim == end_dim:
						return input

			lowerCamelCase__		:		Optional[int]		     =tf.shape(__lowerCamelCase     )
			lowerCamelCase__		:		Optional[int]		     =tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]     )
			lowerCamelCase__		:		List[Any]		     =tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]							,		axis=0     )
			return tf.reshape(__lowerCamelCase							,		__lowerCamelCase     )






def 					snake_case__       (		__lowerCamelCase    :   tf.Tensor     ):
			"""simple docstring"""
			if not isinstance(__lowerCamelCase							,		tf.Tensor     ):
						lowerCamelCase__		:		List[Any]		     =tf.convert_to_tensor(__lowerCamelCase     )  # Catches stray NumPy inputs
			if encoder_attention_mask.shape.rank == 3:
						lowerCamelCase__		:		int		     =encoder_attention_mask[:, None, :, :]
			if encoder_attention_mask.shape.rank == 2:
						lowerCamelCase__		:		int		     =encoder_attention_mask[:, None, None, :]
			# T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
			# Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
			# /transformer/transformer_layers.py#L270
			# encoder_extended_attention_mask = (encoder_extended_attention_mask ==
			# encoder_extended_attention_mask.transpose(-1, -2))
			lowerCamelCase__		:		str		     =(
			    tf.cast(1							,		encoder_attention_mask.dtype     ) - encoder_extended_attention_mask
			) * encoder_extended_attention_mask.dtype.min

			return encoder_extended_attention_mask






def 					snake_case__       (		__lowerCamelCase    :   tf.Tensor							,		__lowerCamelCase    :   int							,		__lowerCamelCase    :   str = "input_ids"     ):
			"""simple docstring"""
			tf.debugging.assert_less(
			    __lowerCamelCase							,		tf.cast(__lowerCamelCase							,		dtype=tensor.dtype     )							,		message=(
			        f'''The maximum value of {tensor_name} ({tf.math.reduce_max(__lowerCamelCase     )}) must be smaller than the embedding '''
			        f'''layer\'s input dimension ({embed_dim}). The likely cause is some problem at tokenization time.'''
			    )							,		)






def 					snake_case__       (		__lowerCamelCase    :   Tuple							,		__lowerCamelCase    :   List[str]							,		__lowerCamelCase    :   Dict     ):
			"""simple docstring"""
			lowerCamelCase__		:		Any		     =64512
			# Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
			# because in that case even chunking the array would not make the saving
			# possible.
			lowerCamelCase__		:		Tuple		     =[x for x in data if len(__lowerCamelCase     ) > HDF5_OBJECT_HEADER_LIMIT]

			# Expecting this to never be true.
			if bad_attributes:
						raise RuntimeError(
						    '''The following attributes cannot be saved to HDF5 file because '''
						    f'''they are larger than {HDF5_OBJECT_HEADER_LIMIT} '''
						    f'''bytes: {bad_attributes}'''     )

			lowerCamelCase__		:		Optional[Any]		     =np.asarray(__lowerCamelCase     )

			lowerCamelCase__		:		str		     =1
			lowerCamelCase__		:		List[Any]		     =np.array_split(__lowerCamelCase							,		__lowerCamelCase     )

			# This will never loop forever thanks to the test above.
			while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data     ):
						num_chunks += 1
						lowerCamelCase__		:		Union[str, Any]		     =np.array_split(__lowerCamelCase							,		__lowerCamelCase     )

			if num_chunks > 1:
						for chunk_id, chunk_data in enumerate(__lowerCamelCase     ):
									lowerCamelCase__		:		List[str]		     =chunk_data
			else:
						lowerCamelCase__		:		Dict		     =data






def 					snake_case__       (		__lowerCamelCase    :   Dict							,		__lowerCamelCase    :   Optional[Any]     ):
			"""simple docstring"""
			if name in group.attrs:
						lowerCamelCase__		:		Optional[int]		     =[n.decode('''utf8'''     ) if hasattr(__lowerCamelCase							,		'''decode'''     ) else n for n in group.attrs[name]]
			else:
						lowerCamelCase__		:		str		     =[]
						lowerCamelCase__		:		str		     =0
						while "%s%d" % (name, chunk_id) in group.attrs:
									data.extend(
									    [n.decode('''utf8'''     ) if hasattr(__lowerCamelCase							,		'''decode'''     ) else n for n in group.attrs['''%s%d''' % (name, chunk_id)]]     )
									chunk_id += 1
			return data






def 					snake_case__       (		__lowerCamelCase    :   Dict     ):
			"""simple docstring"""

			def _expand_single_ad_tensor(__lowerCamelCase    :   List[Any]     ):
						if isinstance(__lowerCamelCase							,		tf.Tensor     ) and t.shape.rank == 1:
									return tf.expand_dims(__lowerCamelCase							,		axis=-1     )
						return t

			return tf.nest.map_structure(_expand_single_ad_tensor							,		__lowerCamelCase     )




