
from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


lowercase__       =logging.get_logger(__name__)  # pylint: disable=invalid-name

lowercase__       ='\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline\n        >>> from diffusers.utils import load_image\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior.to("cuda")\n\n        >>> prompt = "A red cartoon frog, 4k"\n        >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)\n\n        >>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16\n        ... )\n        >>> pipe.to("cuda")\n\n        >>> init_image = load_image(\n        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"\n        ...     "/kandinsky/frog.png"\n        ... )\n\n        >>> image = pipe(\n        ...     image=init_image,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ...     strength=0.2,\n        ... ).images\n\n        >>> image[0].save("red_frog.png")\n        ```\n'




def    __UpperCamelCase     (				lowerCAmelCase__							:				Optional[int]  ,   lowerCAmelCase__							:				Optional[int]  ,   lowerCAmelCase__							:				Any=8  ):
 __a :  List[Any]						=  height // scale_factor**2
 if height % scale_factor**2 != 0:
  new_height += 1
 __a :  int						=  width // scale_factor**2
 if width % scale_factor**2 != 0:
  new_width += 1
 return new_height * scale_factor, new_width * scale_factor




def    __UpperCamelCase     (				lowerCAmelCase__							:				Tuple  ,   lowerCAmelCase__							:				List[Any]=5_1_2  ,   lowerCAmelCase__							:				Any=5_1_2  ):
 __a :  List[str]						=  pil_image.resize((w, h)  ,   resample=Image.BICUBIC  ,   reducing_gap=1  )
 __a :  str						=  np.array(pil_image.convert('''RGB'''  )  )
 __a :  List[Any]						=  arr.astype(np.floataa  ) / 1_27.5 - 1
 __a :  Union[str, Any]						=  np.transpose(lowerCAmelCase__  ,   [2, 0, 1]  )
 __a :  List[str]						=  torch.from_numpy(lowerCAmelCase__  ).unsqueeze(0  )
 return image



class 		UpperCamelCase__  (       __lowercase    ):
   def __init__(self :					Dict					,				snake_case_ :					UNetaDConditionModel					,				snake_case_ :					DDPMScheduler					,				snake_case_ :					VQModel					,				):
    super().__init__()

    self.register_modules(
        unet=snake_case_					,				scheduler=snake_case_					,				movq=snake_case_					,				)
    __a :  Any						=  2 ** (len(self.movq.config.block_out_channels   ) - 1)
   def       lowerCAmelCase			(self :					Any					,				snake_case_ :					Tuple					,				snake_case_ :					Optional[Any]					,				snake_case_ :					Tuple   ):
    # get the original timestep using init_timestep
    __a :  List[Any]						=  min(int(num_inference_steps * strength   )					,				snake_case_   )

    __a :  List[str]						=  max(num_inference_steps - init_timestep					,				0   )
    __a :  List[Any]						=  self.scheduler.timesteps[t_start:]

    return timesteps, num_inference_steps - t_start
   def       lowerCAmelCase			(self :					List[str]					,				snake_case_ :					Dict					,				snake_case_ :					Optional[int]					,				snake_case_ :					Dict					,				snake_case_ :					Optional[Any]					,				snake_case_ :					Tuple					,				snake_case_ :					List[Any]					,				snake_case_ :					Any=None   ):
    if not isinstance(snake_case_					,				(torch.Tensor, PIL.Image.Image, list)   ):
     raise ValueError(
         f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(snake_case_   )}"   )

    __a :  List[str]						=  image.to(device=snake_case_					,				dtype=snake_case_   )

    __a :  Optional[Any]						=  batch_size * num_images_per_prompt

    if image.shape[1] == 4:
     __a :  List[Any]						=  image

    else:
     if isinstance(snake_case_					,				snake_case_   ) and len(snake_case_   ) != batch_size:
      raise ValueError(
          f"You have passed a list of generators of length {len(snake_case_   )}, but requested an effective batch"
          f" size of {batch_size}. Make sure the batch size matches the length of the generators."   )

     elif isinstance(snake_case_					,				snake_case_   ):
      __a :  Tuple						=  [
          self.movq.encode(image[i : i + 1]   ).latent_dist.sample(generator[i]   ) for i in range(snake_case_   )
      ]
      __a :  Union[str, Any]						=  torch.cat(snake_case_					,				dim=0   )
     else:
      __a :  Tuple						=  self.movq.encode(snake_case_   ).latent_dist.sample(snake_case_   )

     __a :  Dict						=  self.movq.config.scaling_factor * init_latents

    __a :  List[str]						=  torch.cat([init_latents]					,				dim=0   )

    __a :  Dict						=  init_latents.shape
    __a :  Union[str, Any]						=  randn_tensor(snake_case_					,				generator=snake_case_					,				device=snake_case_					,				dtype=snake_case_   )

    # get latents
    __a :  Optional[Any]						=  self.scheduler.add_noise(snake_case_					,				snake_case_					,				snake_case_   )

    __a :  Dict						=  init_latents

    return latents
   def       lowerCAmelCase			(self :					Any					,				snake_case_ :					List[str]=0   ):
    if is_accelerate_available():
     from accelerate import cpu_offload
    else:
     raise ImportError('''Please install accelerate via `pip install accelerate`'''   )

    __a :  Union[str, Any]						=  torch.device(f"cuda:{gpu_id}"   )

    __a :  Tuple						=  [
        self.unet,
        self.movq,
    ]
    for cpu_offloaded_model in models:
     if cpu_offloaded_model is not None:
      cpu_offload(snake_case_					,				snake_case_   )
   def       lowerCAmelCase			(self :					List[Any]					,				snake_case_ :					Tuple=0   ):
    if is_accelerate_available() and is_accelerate_version('''>='''					,				'''0.17.0.dev0'''   ):
     from accelerate import cpu_offload_with_hook
    else:
     raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''   )

    __a :  Union[str, Any]						=  torch.device(f"cuda:{gpu_id}"   )

    if self.device.type != "cpu":
     self.to('''cpu'''					,				silence_dtype_warnings=snake_case_   )
     torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

    __a :  List[Any]						=  None
    for cpu_offloaded_model in [self.unet, self.movq]:
     __a				,       __a :  Dict						=  cpu_offload_with_hook(snake_case_					,				snake_case_					,				prev_module_hook=snake_case_   )

    # We'll offload the last model manually.
    __a :  Optional[Any]						=  hook
   @property
   # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
   def       lowerCAmelCase			(self :					Dict   ):
    if not hasattr(self.unet					,				'''_hf_hook'''   ):
     return self.device
    for module in self.unet.modules():
     if (
         hasattr(snake_case_					,				'''_hf_hook'''   )
         and hasattr(module._hf_hook					,				'''execution_device'''   )
         and module._hf_hook.execution_device is not None
     ):
      return torch.device(module._hf_hook.execution_device   )
    return self.device







   @torch.no_grad()
   @replace_example_docstring(snake_case_   )
   def __call__(self :					Optional[Any]					,				snake_case_ :					Union[torch.FloatTensor, List[torch.FloatTensor]]					,				snake_case_ :					Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]]					,				snake_case_ :					Union[torch.FloatTensor, List[torch.FloatTensor]]					,				snake_case_ :					int = 5_1_2					,				snake_case_ :					int = 5_1_2					,				snake_case_ :					int = 1_0_0					,				snake_case_ :					float = 4.0					,				snake_case_ :					float = 0.3					,				snake_case_ :					int = 1					,				snake_case_ :					Optional[Union[torch.Generator, List[torch.Generator]]] = None					,				snake_case_ :					Optional[str] = "pil"					,				snake_case_ :					bool = True					,				):
    __a :  Optional[Any]						=  self._execution_device

    __a :  Any						=  guidance_scale > 1.0

    if isinstance(snake_case_					,				snake_case_   ):
     __a :  str						=  torch.cat(snake_case_					,				dim=0   )
    __a :  int						=  image_embeds.shape[0]
    if isinstance(snake_case_					,				snake_case_   ):
     __a :  Dict						=  torch.cat(snake_case_					,				dim=0   )

    if do_classifier_free_guidance:
     __a :  Optional[int]						=  image_embeds.repeat_interleave(snake_case_					,				dim=0   )
     __a :  Optional[Any]						=  negative_image_embeds.repeat_interleave(snake_case_					,				dim=0   )

    __a :  List[str]						=  torch.cat([negative_image_embeds, image_embeds]					,				dim=0   ).to(dtype=self.unet.dtype					,				device=snake_case_   )

    if not isinstance(snake_case_					,				snake_case_   ):
     __a :  str						=  [image]
    if not all(isinstance(snake_case_					,				(PIL.Image.Image, torch.Tensor)   ) for i in image   ):
     raise ValueError(
         f"Input is in incorrect format: {[type(snake_case_   ) for i in image]}. Currently, we only support  PIL image and pytorch tensor"   )

    __a :  Tuple						=  torch.cat([prepare_image(snake_case_					,				snake_case_					,				snake_case_   ) for i in image]					,				dim=0   )
    __a :  Optional[int]						=  image.to(dtype=image_embeds.dtype					,				device=snake_case_   )

    __a :  List[Any]						=  self.movq.encode(snake_case_   )['''latents''']
    __a :  List[str]						=  latents.repeat_interleave(snake_case_					,				dim=0   )
    self.scheduler.set_timesteps(snake_case_					,				device=snake_case_   )
    __a				,       __a :  Dict						=  self.get_timesteps(snake_case_					,				snake_case_					,				snake_case_   )
    __a :  Dict						=  timesteps[:1].repeat(batch_size * num_images_per_prompt   )
    __a				,       __a :  int						=  downscale_height_and_width(snake_case_					,				snake_case_					,				self.movq_scale_factor   )
    __a :  int						=  self.prepare_latents(
        snake_case_					,				snake_case_					,				snake_case_					,				snake_case_					,				image_embeds.dtype					,				snake_case_					,				snake_case_   )
    for i, t in enumerate(self.progress_bar(snake_case_   )   ):
     # expand the latents if we are doing classifier free guidance
     __a :  Union[str, Any]						=  torch.cat([latents] * 2   ) if do_classifier_free_guidance else latents

     __a :  int						=  {'''image_embeds''': image_embeds}
     __a :  List[str]						=  self.unet(
         sample=snake_case_					,				timestep=snake_case_					,				encoder_hidden_states=snake_case_					,				added_cond_kwargs=snake_case_					,				return_dict=snake_case_					,				)[0]

     if do_classifier_free_guidance:
      __a				,       __a :  Tuple						=  noise_pred.split(latents.shape[1]					,				dim=1   )
      __a				,       __a :  List[Any]						=  noise_pred.chunk(2   )
      __a				,       __a :  Dict						=  variance_pred.chunk(2   )
      __a :  Tuple						=  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
      __a :  List[Any]						=  torch.cat([noise_pred, variance_pred_text]					,				dim=1   )

     if not (
         hasattr(self.scheduler.config					,				'''variance_type'''   )
         and self.scheduler.config.variance_type in ["learned", "learned_range"]
     ):
      __a				,       __a :  List[str]						=  noise_pred.split(latents.shape[1]					,				dim=1   )

     # compute the previous noisy sample x_t -> x_t-1
     __a :  Optional[Any]						=  self.scheduler.step(
         snake_case_					,				snake_case_					,				snake_case_					,				generator=snake_case_					,				)[0]

    # post-processing
    __a :  str						=  self.movq.decode(snake_case_					,				force_not_quantize=snake_case_   )['''sample''']

    if output_type not in ["pt", "np", "pil"]:
     raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"   )

    if output_type in ["np", "pil"]:
     __a :  List[Any]						=  image * 0.5 + 0.5
     __a :  Tuple						=  image.clamp(0					,				1   )
     __a :  Union[str, Any]						=  image.cpu().permute(0					,				2					,				3					,				1   ).float().numpy()

    if output_type == "pil":
     __a :  Tuple						=  self.numpy_to_pil(snake_case_   )

    if not return_dict:
     return (image,)

    return ImagePipelineOutput(images=snake_case_   )


from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import TFLEDForConditionalGeneration, TFLEDModel



@require_tf
class 		UpperCamelCase__  :
   _SCREAMING_SNAKE_CASE      :							Union[str, Any]	       =      LEDConfig
   _SCREAMING_SNAKE_CASE      :							Optional[int]	       =      {}
   _SCREAMING_SNAKE_CASE      :							int	       =      "gelu"
   def __init__(self :					List[str]					,				snake_case_ :					Union[str, Any]					,				snake_case_ :					Union[str, Any]=1_3					,				snake_case_ :					Optional[Any]=7					,				snake_case_ :					Any=True					,				snake_case_ :					List[Any]=False					,				snake_case_ :					str=9_9					,				snake_case_ :					Any=3_2					,				snake_case_ :					Dict=2					,				snake_case_ :					List[Any]=4					,				snake_case_ :					Optional[int]=3_7					,				snake_case_ :					Dict=0.1					,				snake_case_ :					int=0.1					,				snake_case_ :					Optional[Any]=2_0					,				snake_case_ :					Optional[Any]=2					,				snake_case_ :					Optional[int]=1					,				snake_case_ :					Optional[int]=0					,				snake_case_ :					str=4					,				):
    __a :  List[Any]						=  parent
    __a :  Union[str, Any]						=  batch_size
    __a :  List[str]						=  seq_length
    __a :  Any						=  is_training
    __a :  Tuple						=  use_labels
    __a :  List[Any]						=  vocab_size
    __a :  Optional[Any]						=  hidden_size
    __a :  int						=  num_hidden_layers
    __a :  Optional[int]						=  num_attention_heads
    __a :  int						=  intermediate_size

    __a :  Union[str, Any]						=  hidden_dropout_prob
    __a :  Dict						=  attention_probs_dropout_prob
    __a :  int						=  max_position_embeddings
    __a :  Tuple						=  eos_token_id
    __a :  Optional[Any]						=  pad_token_id
    __a :  List[str]						=  bos_token_id
    __a :  List[str]						=  attention_window

    # `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
    # [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
    # returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
    # because its local attention only attends to `self.attention_window` and one before and one after
    __a :  Union[str, Any]						=  self.attention_window + 2

    # because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
    # the `test_attention_outputs` and `test_hidden_states_output` tests
    __a :  List[str]						=  (
        self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
    )
   def       lowerCAmelCase			(self :					Optional[int]   ):
    __a :  Dict						=  ids_tensor([self.batch_size, self.seq_length - 1]					,				self.vocab_size   )
    __a :  Union[str, Any]						=  tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size   )					,				1   )
    __a :  int						=  tf.concat([input_ids, eos_tensor]					,				axis=1   )

    __a :  List[str]						=  ids_tensor([self.batch_size, self.seq_length]					,				self.vocab_size   )

    __a :  Optional[int]						=  self.config_cls(
        vocab_size=self.vocab_size					,				d_model=self.hidden_size					,				encoder_layers=self.num_hidden_layers					,				decoder_layers=self.num_hidden_layers					,				encoder_attention_heads=self.num_attention_heads					,				decoder_attention_heads=self.num_attention_heads					,				encoder_ffn_dim=self.intermediate_size					,				decoder_ffn_dim=self.intermediate_size					,				dropout=self.hidden_dropout_prob					,				attention_dropout=self.attention_probs_dropout_prob					,				max_position_embeddings=self.max_position_embeddings					,				eos_token_ids=[2]					,				bos_token_id=self.bos_token_id					,				pad_token_id=self.pad_token_id					,				decoder_start_token_id=self.pad_token_id					,				attention_window=self.attention_window					,				**self.config_updates					,				)
    __a :  Optional[int]						=  prepare_led_inputs_dict(snake_case_					,				snake_case_					,				snake_case_   )
    __a :  Dict						=  tf.concat(
        [tf.zeros_like(snake_case_   )[:, :-1], tf.ones_like(snake_case_   )[:, -1:]]					,				axis=-1					,				)
    __a :  Tuple						=  global_attention_mask
    return config, inputs_dict







   def       lowerCAmelCase			(self :					List[Any]					,				snake_case_ :					Dict					,				snake_case_ :					int   ):
    __a :  List[str]						=  TFLEDModel(config=snake_case_   ).get_decoder()
    __a :  Dict						=  inputs_dict['''input_ids''']

    __a :  Dict						=  input_ids[:1, :]
    __a :  Any						=  inputs_dict['''attention_mask'''][:1, :]
    __a :  List[str]						=  1

    # first forward pass
    __a :  Optional[Any]						=  model(snake_case_					,				attention_mask=snake_case_					,				use_cache=snake_case_   )

    __a				,       __a :  Tuple						=  outputs.to_tuple()

    # create hypothetical next token and extent to next_input_ids
    __a :  List[Any]						=  ids_tensor((self.batch_size, 3)					,				config.vocab_size   )
    __a :  str						=  tf.cast(ids_tensor((self.batch_size, 3)					,				2   )					,				tf.inta   )

    # append to next input_ids and
    __a :  Optional[Any]						=  tf.concat([input_ids, next_tokens]					,				axis=-1   )
    __a :  Optional[Any]						=  tf.concat([attention_mask, next_attn_mask]					,				axis=-1   )

    __a :  Optional[int]						=  model(snake_case_					,				attention_mask=snake_case_   )[0]
    __a :  int						=  model(snake_case_					,				attention_mask=snake_case_					,				past_key_values=snake_case_   )[0]

    self.parent.assertEqual(next_tokens.shape[1]					,				output_from_past.shape[1]   )

    # select random slice
    __a :  Any						=  int(ids_tensor((1,)					,				output_from_past.shape[-1]   )   )
    __a :  List[Any]						=  output_from_no_past[:, -3:, random_slice_idx]
    __a :  Dict						=  output_from_past[:, :, random_slice_idx]

    # test that outputs are equal for slice
    tf.debugging.assert_near(snake_case_					,				snake_case_					,				rtol=1E-3   )





def    __UpperCamelCase     (				lowerCAmelCase__							:				Union[str, Any]  ,   lowerCAmelCase__							:				List[Any]  ,   lowerCAmelCase__							:				int  ,   lowerCAmelCase__							:				List[Any]=None  ,   lowerCAmelCase__							:				int=None  ,   lowerCAmelCase__							:				Union[str, Any]=None  ,   lowerCAmelCase__							:				List[str]=None  ,   ):
 if attention_mask is None:
  __a :  Any						=  tf.cast(tf.math.not_equal(lowerCAmelCase__  ,   config.pad_token_id  )  ,   tf.inta  )
 if decoder_attention_mask is None:
  __a :  Dict						=  tf.concat(
      [
          tf.ones(decoder_input_ids[:, :1].shape  ,   dtype=tf.inta  ),
          tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]  ,   config.pad_token_id  )  ,   tf.inta  ),
      ]  ,   axis=-1  ,   )
 if head_mask is None:
  __a :  Dict						=  tf.ones((config.encoder_layers, config.encoder_attention_heads)  )
 if decoder_head_mask is None:
  __a :  str						=  tf.ones((config.decoder_layers, config.decoder_attention_heads)  )
 return {
     "input_ids": input_ids,
     "attention_mask": attention_mask,
     "decoder_input_ids": decoder_input_ids,
     "decoder_attention_mask": decoder_attention_mask,
     "head_mask": head_mask,
     "decoder_head_mask": decoder_head_mask,
 }



@require_tf
class 		UpperCamelCase__  (       __lowercase    ,__lowercase    ,unittest.TestCase    ):
   _SCREAMING_SNAKE_CASE      :							Dict	       =      (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
   _SCREAMING_SNAKE_CASE      :							Tuple	       =      (TFLEDForConditionalGeneration,) if is_tf_available() else ()
   _SCREAMING_SNAKE_CASE      :							Optional[Any]	       =      (
       {
           "conversational": TFLEDForConditionalGeneration,
           "feature-extraction": TFLEDModel,
           "summarization": TFLEDForConditionalGeneration,
           "text2text-generation": TFLEDForConditionalGeneration,
           "translation": TFLEDForConditionalGeneration,
       }
       if is_tf_available()
       else {}
   )
   _SCREAMING_SNAKE_CASE      :							Union[str, Any]	       =      True
   _SCREAMING_SNAKE_CASE      :							Optional[int]	       =      False
   _SCREAMING_SNAKE_CASE      :							Optional[Any]	       =      False
   _SCREAMING_SNAKE_CASE      :							Optional[Any]	       =      False
   def       lowerCAmelCase			(self :					Optional[int]   ):
    __a :  List[str]						=  TFLEDModelTester(self   )
    __a :  Optional[int]						=  ConfigTester(self					,				config_class=snake_case_   )
   def       lowerCAmelCase			(self :					Any   ):
    self.config_tester.run_common_tests()
   def       lowerCAmelCase			(self :					Optional[Any]   ):
    __a :  Union[str, Any]						=  self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.check_decoder_model_past_large_inputs(*snake_case_   )
   def       lowerCAmelCase			(self :					Any   ):
    __a				,       __a :  Tuple						=  self.model_tester.prepare_config_and_inputs_for_common()
    __a :  Any						=  tf.zeros_like(inputs_dict['''attention_mask''']   )
    __a :  Tuple						=  2
    __a :  Dict						=  tf.where(
        tf.range(self.model_tester.seq_length   )[None, :] < num_global_attn_indices					,				1					,				inputs_dict['''global_attention_mask''']					,				)

    __a :  List[str]						=  True
    __a :  Tuple						=  self.model_tester.seq_length
    __a :  Any						=  self.model_tester.encoder_seq_length

    def check_decoder_attentions_output(snake_case_ :					Any   ):
     __a :  str						=  outputs.decoder_attentions
     self.assertEqual(len(snake_case_   )					,				self.model_tester.num_hidden_layers   )
     self.assertListEqual(
         list(decoder_attentions[0].shape[-3:]   )					,				[self.model_tester.num_attention_heads, seq_length, seq_length]					,				)

    def check_encoder_attentions_output(snake_case_ :					Optional[int]   ):
     __a :  int						=  [t.numpy() for t in outputs.encoder_attentions]
     __a :  int						=  [t.numpy() for t in outputs.encoder_global_attentions]
     self.assertEqual(len(snake_case_   )					,				self.model_tester.num_hidden_layers   )
     self.assertEqual(len(snake_case_   )					,				self.model_tester.num_hidden_layers   )
     self.assertListEqual(
         list(attentions[0].shape[-3:]   )					,				[self.model_tester.num_attention_heads, seq_length, seq_length]					,				)
     self.assertListEqual(
         list(global_attentions[0].shape[-3:]   )					,				[self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices]					,				)

    for model_class in self.all_model_classes:
     __a :  Dict						=  True
     __a :  Optional[Any]						=  False
     __a :  List[str]						=  False
     __a :  List[Any]						=  model_class(snake_case_   )
     __a :  List[str]						=  model(self._prepare_for_class(snake_case_					,				snake_case_   )   )
     __a :  List[str]						=  len(snake_case_   )
     self.assertEqual(config.output_hidden_states					,				snake_case_   )
     check_encoder_attentions_output(snake_case_   )

     if self.is_encoder_decoder:
      __a :  List[str]						=  model_class(snake_case_   )
      __a :  int						=  model(self._prepare_for_class(snake_case_					,				snake_case_   )   )
      self.assertEqual(config.output_hidden_states					,				snake_case_   )
      check_decoder_attentions_output(snake_case_   )

     # Check that output attentions can also be changed via the config
     del inputs_dict["output_attentions"]
     __a :  List[Any]						=  True
     __a :  Dict						=  model_class(snake_case_   )
     __a :  Tuple						=  model(self._prepare_for_class(snake_case_					,				snake_case_   )   )
     self.assertEqual(config.output_hidden_states					,				snake_case_   )
     check_encoder_attentions_output(snake_case_   )

     # Check attention is always last and order is fine
     __a :  List[str]						=  True
     __a :  Any						=  True
     __a :  Tuple						=  model_class(snake_case_   )
     __a :  int						=  model(self._prepare_for_class(snake_case_					,				snake_case_   )   )

     self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)					,				len(snake_case_   )   )
     self.assertEqual(model.config.output_hidden_states					,				snake_case_   )
     check_encoder_attentions_output(snake_case_   )
   @unittest.skip('''LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'''   )
   def       lowerCAmelCase			(self :					List[str]   ):
    pass







   def       lowerCAmelCase			(self :					List[Any]   ):
    # TODO: Head-masking not yet implement
    pass





def    __UpperCamelCase     (				lowerCAmelCase__							:				Optional[Any]  ):
 return tf.constant(lowerCAmelCase__  ,   dtype=tf.intaa  )


lowercase__       =1e-4



@slow
@require_tf
class 		UpperCamelCase__  (       unittest.TestCase    ):
   def       lowerCAmelCase			(self :					Any   ):
    __a :  Dict						=  TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''   ).led

    # change to intended input here
    __a :  Union[str, Any]						=  _long_tensor([5_1_2 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
    __a :  Dict						=  _long_tensor([1_2_8 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
    __a :  List[str]						=  prepare_led_inputs_dict(model.config					,				snake_case_					,				snake_case_   )
    __a :  List[str]						=  model(**snake_case_   )[0]
    __a :  Any						=  (1, 1_0_2_4, 7_6_8)
    self.assertEqual(output.shape					,				snake_case_   )
    # change to expected output here
    __a :  Dict						=  tf.convert_to_tensor(
        [[2.3050, 2.8279, 0.6531], [-1.8457, -0.1455, -3.5661], [-1.0186, 0.4586, -2.2043]]					,				)
    tf.debugging.assert_near(output[:, :3, :3]					,				snake_case_					,				atol=1E-3   )







   def       lowerCAmelCase			(self :					int   ):
    __a :  Optional[Any]						=  TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''   )

    # change to intended input here
    __a :  int						=  _long_tensor([5_1_2 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
    __a :  Tuple						=  _long_tensor([1_2_8 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
    __a :  Dict						=  prepare_led_inputs_dict(model.config					,				snake_case_					,				snake_case_   )
    __a :  List[str]						=  model(**snake_case_   )[0]
    __a :  List[Any]						=  (1, 1_0_2_4, model.config.vocab_size)
    self.assertEqual(output.shape					,				snake_case_   )
    # change to expected output here
    __a :  str						=  tf.convert_to_tensor(
        [[33.6507, 6.4572, 16.8089], [5.8739, -2.4238, 11.2902], [-3.2139, -4.3149, 4.2783]]					,				)
    tf.debugging.assert_near(output[:, :3, :3]					,				snake_case_					,				atol=1E-3					,				rtol=1E-3   )
