





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     UpperCamelCase__					(     lowerCamelCase__				, lowerCamelCase__				, lowerCamelCase__				, unittest.TestCase				):






		'''simple docstring'''






		__a       :       Union[str, Any]       =	StableDiffusionInstructPixaPixPipeline
		__a       :       Tuple       =	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width""", """cross_attention_kwargs"""}
		__a       :       Optional[Any]       =	TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
		__a       :       Any       =	IMAGE_TO_IMAGE_IMAGE_PARAMS
		__a       :       int       =	IMAGE_TO_IMAGE_IMAGE_PARAMS






		def 						snake_case__				(							self  )       -> Any:






									"""simple docstring"""
									torch.manual_seed(0  )
									lowercase_	:   Tuple  				=     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,   )
									lowercase_	:   str  				=     PNDMScheduler(skip_prk_steps=snake_case__  )
									torch.manual_seed(0  )
									lowercase_	:   int  				=     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  )
									lowercase_	:   int  				=     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=10_00,   )
									lowercase_	:   List[str]  				=     CLIPTextModel(snake_case__  )
									lowercase_	:   int  				=     CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

									lowercase_	:   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,   snake_case__,   snake_case__=0  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   List[str]  				=     floats_tensor((1, 3, 32, 32),   rng=random.Random(snake_case__  )  ).to(snake_case__  )
									lowercase_	:   int  				=     image.cpu().permute(0,   2,   3,   1  )[0]
									lowercase_	:   Dict  				=     Image.fromarray(np.uinta(snake_case__  )  ).convert("""RGB"""  )
									if str(snake_case__  ).startswith("""mps"""  ):
																lowercase_	:   int  				=     torch.manual_seed(snake_case__  )
									else:
																lowercase_	:   Any  				=     torch.Generator(device=snake_case__  ).manual_seed(snake_case__  )
									lowercase_	:   Optional[int]  				=     {
									    """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  )       -> str:






									"""simple docstring"""
									lowercase_	:   Dict  				=     """cpu"""  # ensure determinism for the device-dependent torch.Generator
									lowercase_	:   Optional[Any]  				=     self.get_dummy_components()
									lowercase_	:   int  				=     StableDiffusionInstructPixaPixPipeline(**snake_case__  )
									lowercase_	:   Any  				=     sd_pipe.to(snake_case__  )
									sd_pipe.set_progress_bar_config(disable=snake_case__  )

									lowercase_	:   List[Any]  				=     self.get_dummy_inputs(snake_case__  )
									lowercase_	:   Tuple  				=     sd_pipe(**snake_case__  ).images
									lowercase_	:   Dict  				=     image[0, -3:, -3:, -1]
									assert image.shape == (1, 32, 32, 3)
									lowercase_	:   Union[str, Any]  				=     np.array([0.7526, 0.3750, 0.4547, 0.6117, 0.5866, 0.5016, 0.4327, 0.5642, 0.4815]  )

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






		def 						snake_case__				(							self  )       -> Union[str, Any]:






									"""simple docstring"""
									lowercase_	:   List[str]  				=     """cpu"""  # ensure determinism for the device-dependent torch.Generator
									lowercase_	:   List[Any]  				=     self.get_dummy_components()
									lowercase_	:   Optional[int]  				=     StableDiffusionInstructPixaPixPipeline(**snake_case__  )
									lowercase_	:   str  				=     sd_pipe.to(snake_case__  )
									sd_pipe.set_progress_bar_config(disable=snake_case__  )

									lowercase_	:   List[Any]  				=     self.get_dummy_inputs(snake_case__  )
									lowercase_	:   Tuple  				=     """french fries"""
									lowercase_	:   int  				=     sd_pipe(**snake_case__,   negative_prompt=snake_case__  )
									lowercase_	:   List[str]  				=     output.images
									lowercase_	:   List[Any]  				=     image[0, -3:, -3:, -1]

									assert image.shape == (1, 32, 32, 3)
									lowercase_	:   Optional[Any]  				=     np.array([0.7511, 0.3642, 0.4553, 0.6236, 0.5797, 0.5013, 0.4343, 0.5611, 0.4831]  )

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






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									lowercase_	:   List[str]  				=     """cpu"""  # ensure determinism for the device-dependent torch.Generator
									lowercase_	:   Union[str, Any]  				=     self.get_dummy_components()
									lowercase_	:   Optional[int]  				=     StableDiffusionInstructPixaPixPipeline(**snake_case__  )
									lowercase_	:   Optional[Any]  				=     sd_pipe.to(snake_case__  )
									sd_pipe.set_progress_bar_config(disable=snake_case__  )

									lowercase_	:   Dict  				=     self.get_dummy_inputs(snake_case__  )
									lowercase_	:   Any  				=     [inputs["""prompt"""]] * 2

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

									lowercase_	:   Optional[Any]  				=     sd_pipe(**snake_case__  ).images
									lowercase_	:   List[Any]  				=     image[-1, -3:, -3:, -1]

									assert image.shape == (2, 32, 32, 3)
									lowercase_	:   Tuple  				=     np.array([0.5812, 0.5748, 0.5222, 0.5908, 0.5695, 0.7174, 0.6804, 0.5523, 0.5579]  )

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






		def 						snake_case__				(							self  )       -> Any:






									"""simple docstring"""
									lowercase_	:   str  				=     """cpu"""  # ensure determinism for the device-dependent torch.Generator
									lowercase_	:   Tuple  				=     self.get_dummy_components()
									lowercase_	:   int  				=     EulerAncestralDiscreteScheduler(
									    beta_start=0.00085,   beta_end=0.012,   beta_schedule="""scaled_linear"""  )
									lowercase_	:   List[Any]  				=     StableDiffusionInstructPixaPixPipeline(**snake_case__  )
									lowercase_	:   Any  				=     sd_pipe.to(snake_case__  )
									sd_pipe.set_progress_bar_config(disable=snake_case__  )

									lowercase_	:   Optional[Any]  				=     self.get_dummy_inputs(snake_case__  )
									lowercase_	:   Dict  				=     sd_pipe(**snake_case__  ).images
									lowercase_	:   str  				=     image[0, -3:, -3:, -1]

									lowercase_	:   Optional[int]  				=     [round(snake_case__,   4  ) for x in image_slice.flatten().tolist()]
									print(""",""".join([str(snake_case__  ) for x in slice]  )  )

									assert image.shape == (1, 32, 32, 3)
									lowercase_	:   str  				=     np.array([0.7417, 0.3842, 0.4732, 0.5776, 0.5891, 0.5139, 0.4052, 0.5673, 0.4986]  )

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






		def 						snake_case__				(							self  )       -> List[str]:






									"""simple docstring"""
									super().test_inference_batch_single_identical(expected_max_diff=3E-3  )






		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     self.get_dummy_components()
									lowercase_	:   Dict  				=     StableDiffusionInstructPixaPixPipeline(**snake_case__  )
									lowercase_	:   List[str]  				=     VaeImageProcessor(do_resize=snake_case__,   do_normalize=snake_case__  )
									lowercase_	:   Tuple  				=     pipe.to(snake_case__  )
									pipe.set_progress_bar_config(disable=snake_case__  )

									lowercase_	:   str  				=     pipe(**self.get_dummy_inputs_by_type(snake_case__,   input_image_type="""pt"""  )  )[0]

									lowercase_	:   Tuple  				=     components["""vae"""]
									lowercase_	:   List[Any]  				=     self.get_dummy_inputs_by_type(snake_case__,   input_image_type="""pt"""  )

									for image_param in self.image_latents_params:
																if image_param in inputs.keys():
																							lowercase_	:   Union[str, Any]  				=     vae.encode(inputs[image_param]  ).latent_dist.mode()

									lowercase_	:   int  				=     pipe(**snake_case__  )[0]

									lowercase_	:   Optional[Any]  				=     np.abs(out - out_latents_inputs  ).max()
									self.assertLess(snake_case__,   1E-4,   """passing latents as image input generate different result from passing image"""  )




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






		'''simple docstring'''






		def 						snake_case__				(							self  )       -> Union[str, Any]:






									"""simple docstring"""
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()






		def 						snake_case__				(							self,   snake_case__=0  )       -> Union[str, Any]:






									"""simple docstring"""
									lowercase_	:   str  				=     torch.manual_seed(snake_case__  )
									lowercase_	:   Any  				=     load_image(
									    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg"""  )
									lowercase_	:   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  )       -> Union[str, Any]:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     StableDiffusionInstructPixaPixPipeline.from_pretrained(
									    """timbrooks/instruct-pix2pix""",   safety_checker=snake_case__  )
									pipe.to(snake_case__  )
									pipe.set_progress_bar_config(disable=snake_case__  )
									pipe.enable_attention_slicing()

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

									assert image.shape == (1, 5_12, 5_12, 3)
									lowercase_	:   Any  				=     np.array([0.5902, 0.6015, 0.6027, 0.5983, 0.6092, 0.6061, 0.5765, 0.5785, 0.5555]  )

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






		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     StableDiffusionInstructPixaPixPipeline.from_pretrained(
									    """timbrooks/instruct-pix2pix""",   safety_checker=snake_case__  )
									lowercase_	:   Tuple  				=     LMSDiscreteScheduler.from_config(pipe.scheduler.config  )
									pipe.to(snake_case__  )
									pipe.set_progress_bar_config(disable=snake_case__  )
									pipe.enable_attention_slicing()

									lowercase_	:   str  				=     self.get_inputs()
									lowercase_	:   Optional[int]  				=     pipe(**snake_case__  ).images
									lowercase_	:   List[Any]  				=     image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 5_12, 5_12, 3)
									lowercase_	:   Optional[Any]  				=     np.array([0.6578, 0.6817, 0.6972, 0.6761, 0.6856, 0.6916, 0.6428, 0.6516, 0.6301]  )

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






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     StableDiffusionInstructPixaPixPipeline.from_pretrained(
									    """timbrooks/instruct-pix2pix""",   safety_checker=snake_case__  )
									lowercase_	:   Optional[int]  				=     DDIMScheduler.from_config(pipe.scheduler.config  )
									pipe.to(snake_case__  )
									pipe.set_progress_bar_config(disable=snake_case__  )
									pipe.enable_attention_slicing()

									lowercase_	:   List[Any]  				=     self.get_inputs()
									lowercase_	:   Optional[int]  				=     pipe(**snake_case__  ).images
									lowercase_	:   Union[str, Any]  				=     image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 5_12, 5_12, 3)
									lowercase_	:   List[str]  				=     np.array([0.3828, 0.3834, 0.3818, 0.3792, 0.3865, 0.3752, 0.3792, 0.3847, 0.3753]  )

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






		def 						snake_case__				(							self  )       -> Any:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     0

									def callback_fn(snake_case__,   snake_case__,   snake_case__  ) -> None:
																lowercase_	:   Any  				=     True
																nonlocal number_of_steps
																number_of_steps += 1
																if step == 1:
																							lowercase_	:   Tuple  				=     latents.detach().cpu().numpy()
																							assert latents.shape == (1, 4, 64, 64)
																							lowercase_	:   Tuple  				=     latents[0, -3:, -3:, -1]
																							lowercase_	:   Optional[Any]  				=     np.array([-0.2463, -0.4644, -0.9756, 1.5176, 1.4414, 0.7866, 0.9897, 0.8521, 0.7983]  )

																							assert np.abs(latents_slice.flatten() - expected_slice  ).max() < 5E-2
																elif step == 2:
																							lowercase_	:   Optional[int]  				=     latents.detach().cpu().numpy()
																							assert latents.shape == (1, 4, 64, 64)
																							lowercase_	:   Union[str, Any]  				=     latents[0, -3:, -3:, -1]
																							lowercase_	:   Any  				=     np.array([-0.2644, -0.4626, -0.9653, 1.5176, 1.4551, 0.7686, 0.9805, 0.8452, 0.8115]  )

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

									lowercase_	:   str  				=     False

									lowercase_	:   Tuple  				=     StableDiffusionInstructPixaPixPipeline.from_pretrained(
									    """timbrooks/instruct-pix2pix""",   safety_checker=snake_case__,   torch_dtype=torch.floataa  )
									lowercase_	:   Optional[int]  				=     pipe.to(snake_case__  )
									pipe.set_progress_bar_config(disable=snake_case__  )
									pipe.enable_attention_slicing()

									lowercase_	:   Optional[Any]  				=     self.get_inputs()
									pipe(**snake_case__,   callback=snake_case__,   callback_steps=1  )
									assert callback_fn.has_been_called
									assert number_of_steps == 3






		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

									lowercase_	:   int  				=     StableDiffusionInstructPixaPixPipeline.from_pretrained(
									    """timbrooks/instruct-pix2pix""",   safety_checker=snake_case__,   torch_dtype=torch.floataa  )
									lowercase_	:   Optional[Any]  				=     pipe.to(snake_case__  )
									pipe.set_progress_bar_config(disable=snake_case__  )
									pipe.enable_attention_slicing(1  )
									pipe.enable_sequential_cpu_offload()

									lowercase_	:   str  				=     self.get_inputs()
									lowercase_	:   Any  				=     pipe(**snake_case__  )

									lowercase_	:   Optional[Any]  				=     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  )       -> Any:






									"""simple docstring"""
									lowercase_	:   Optional[Any]  				=     self.get_inputs()
									# resize to resolution that is divisible by 8 but not 16 or 32
									lowercase_	:   Any  				=     inputs["""image"""].resize((5_04, 5_04)  )

									lowercase_	:   List[Any]  				=     """timbrooks/instruct-pix2pix"""
									lowercase_	:   Optional[int]  				=     StableDiffusionInstructPixaPixPipeline.from_pretrained(
									    snake_case__,   safety_checker=snake_case__,   )
									pipe.to(snake_case__  )
									pipe.set_progress_bar_config(disable=snake_case__  )
									pipe.enable_attention_slicing()

									lowercase_	:   List[Any]  				=     pipe(**snake_case__  )
									lowercase_	:   str  				=     output.images[0]

									lowercase_	:   Optional[int]  				=     image[2_55:2_58, 3_83:3_86, -1]

									assert image.shape == (5_04, 5_04, 3)
									lowercase_	:   str  				=     np.array([0.2726, 0.2529, 0.2664, 0.2655, 0.2641, 0.2642, 0.2591, 0.2649, 0.2590]  )

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






import tensorflow as tf

from ...tf_utils import shape_list
class     UpperCamelCase__					(     tf.keras.layers.Layer				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__=1,   snake_case__=False,   **snake_case__  )       -> Optional[int]:






									"""simple docstring"""
									super().__init__(**snake_case__  )

									lowercase_	:   Tuple  				=     vocab_size
									lowercase_	:   Union[str, Any]  				=     d_embed
									lowercase_	:   Optional[Any]  				=     d_proj

									lowercase_	:   Optional[Any]  				=     cutoffs + [vocab_size]
									lowercase_	:   Optional[int]  				=     [0] + self.cutoffs
									lowercase_	:   Union[str, Any]  				=     div_val

									lowercase_	:   Union[str, Any]  				=     self.cutoffs[0]
									lowercase_	:   List[str]  				=     len(self.cutoffs  ) - 1
									lowercase_	:   List[str]  				=     self.shortlist_size + self.n_clusters
									lowercase_	:   Tuple  				=     keep_order

									lowercase_	:   Dict  				=     []
									lowercase_	:   int  				=     []






		def 						snake_case__				(							self,   snake_case__  )       -> Optional[int]:






									"""simple docstring"""
									if self.n_clusters > 0:
																lowercase_	:   List[str]  				=     self.add_weight(
																    shape=(self.n_clusters, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name="""cluster_weight"""  )
																lowercase_	:   List[Any]  				=     self.add_weight(
																    shape=(self.n_clusters,),   initializer="""zeros""",   trainable=snake_case__,   name="""cluster_bias"""  )

									if self.div_val == 1:
																for i in range(len(self.cutoffs  )  ):
																							if self.d_proj != self.d_embed:
																														lowercase_	:   Any  				=     self.add_weight(
																														    shape=(self.d_embed, self.d_proj),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_projs_._{i}""",   )
																														self.out_projs.append(snake_case__  )
																							else:
																														self.out_projs.append(snake_case__  )
																							lowercase_	:   List[str]  				=     self.add_weight(
																							    shape=(self.vocab_size, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   List[str]  				=     self.add_weight(
																							    shape=(self.vocab_size,),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._bias""",   )
																							self.out_layers.append((weight, bias)  )
									else:
																for i in range(len(self.cutoffs  )  ):
																							lowercase_				,		lowercase_	:   Optional[Any]  				=     self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							lowercase_	:   Union[str, Any]  				=     self.d_embed // (self.div_val**i)

																							lowercase_	:   int  				=     self.add_weight(
																							    shape=(d_emb_i, self.d_proj),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_projs_._{i}"""  )
																							self.out_projs.append(snake_case__  )
																							lowercase_	:   Any  				=     self.add_weight(
																							    shape=(r_idx - l_idx, d_emb_i),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   Tuple  				=     self.add_weight(
																							    shape=(r_idx - l_idx,),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._bias""",   )
																							self.out_layers.append((weight, bias)  )
									super().build(snake_case__  )






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__,   snake_case__,   snake_case__=None  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   Dict  				=     x
									if proj is not None:
																lowercase_	:   List[Any]  				=     tf.einsum("""ibd,ed->ibe""",   snake_case__,   snake_case__  )
									return tf.einsum("""ibd,nd->ibn""",   snake_case__,   snake_case__  ) + b






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__  )       -> List[str]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     shape_list(snake_case__  )
									lowercase_	:   Optional[Any]  				=     tf.range(lp_size[0],   dtype=target.dtype  )
									lowercase_	:   str  				=     tf.stack([r, target],   1  )
									return tf.gather_nd(snake_case__,   snake_case__  )






		def 						snake_case__				(							self,   snake_case__,   snake_case__,   snake_case__=True,   snake_case__=False  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   Any  				=     0
									if self.n_clusters == 0:
																lowercase_	:   Union[str, Any]  				=     self._logit(snake_case__,   self.out_layers[0][0],   self.out_layers[0][1],   self.out_projs[0]  )
																if target is not None:
																							lowercase_	:   str  				=     tf.nn.sparse_softmax_cross_entropy_with_logits(labels=snake_case__,   logits=snake_case__  )
																lowercase_	:   Optional[int]  				=     tf.nn.log_softmax(snake_case__,   axis=-1  )
									else:
																lowercase_	:   int  				=     shape_list(snake_case__  )
																lowercase_	:   List[Any]  				=     []
																lowercase_	:   Tuple  				=     tf.zeros(hidden_sizes[:2]  )
																for i in range(len(self.cutoffs  )  ):
																							lowercase_				,		lowercase_	:   Optional[int]  				=     self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							if target is not None:
																														lowercase_	:   Tuple  				=     (target >= l_idx) & (target < r_idx)
																														lowercase_	:   Union[str, Any]  				=     tf.where(snake_case__  )
																														lowercase_	:   Optional[Any]  				=     tf.boolean_mask(snake_case__,   snake_case__  ) - l_idx

																							if self.div_val == 1:
																														lowercase_	:   Optional[Any]  				=     self.out_layers[0][0][l_idx:r_idx]
																														lowercase_	:   Union[str, Any]  				=     self.out_layers[0][1][l_idx:r_idx]
																							else:
																														lowercase_	:   Dict  				=     self.out_layers[i][0]
																														lowercase_	:   int  				=     self.out_layers[i][1]

																							if i == 0:
																														lowercase_	:   Optional[int]  				=     tf.concat([cur_W, self.cluster_weight],   0  )
																														lowercase_	:   Optional[Any]  				=     tf.concat([cur_b, self.cluster_bias],   0  )

																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[0]  )
																														lowercase_	:   List[str]  				=     tf.nn.log_softmax(snake_case__  )
																														out.append(head_logprob[..., : self.cutoffs[0]]  )
																														if target is not None:
																																					lowercase_	:   int  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     self._gather_logprob(snake_case__,   snake_case__  )
																							else:
																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[i]  )
																														lowercase_	:   Dict  				=     tf.nn.log_softmax(snake_case__  )
																														lowercase_	:   Optional[int]  				=     self.cutoffs[0] + i - 1  # No probability for the head cluster
																														lowercase_	:   int  				=     head_logprob[..., cluster_prob_idx, None] + tail_logprob
																														out.append(snake_case__  )
																														if target is not None:
																																					lowercase_	:   Optional[Any]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Union[str, Any]  				=     self._gather_logprob(snake_case__,   snake_case__  )
																																					cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
																							if target is not None:
																														loss += tf.scatter_nd(snake_case__,   -cur_logprob,   shape_list(snake_case__  )  )
																lowercase_	:   List[str]  				=     tf.concat(snake_case__,   axis=-1  )

									if target is not None:
																if return_mean:
																							lowercase_	:   Tuple  				=     tf.reduce_mean(snake_case__  )
																# Add the training-time loss value to the layer using `self.add_loss()`.
																self.add_loss(snake_case__  )

																# Log the loss as a metric (we could log arbitrary metrics,
																# including different metrics for training and inference.
																self.add_metric(snake_case__,   name=self.name,   aggregation="""mean""" if return_mean else """"""  )

									return out