

"""simple docstring"""






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    :
  _UpperCAmelCase          =		LEDConfig
  _UpperCAmelCase          =		{}
  _UpperCAmelCase          =		"gelu"

  def __init__(		self    ,_A    ,_A=13    ,_A=7    ,_A=True    ,_A=False    ,_A=99    ,_A=32    ,_A=2    ,_A=4    ,_A=37    ,_A=0.1    ,_A=0.1    ,_A=20    ,_A=2    ,_A=1    ,_A=0    ,_A=4    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   str  =    parent
      _lowerCAmelCase       :   Dict  =    batch_size
      _lowerCAmelCase       :   Any  =    seq_length
      _lowerCAmelCase       :   Dict  =    is_training
      _lowerCAmelCase       :   Optional[int]  =    use_labels
      _lowerCAmelCase       :   List[str]  =    vocab_size
      _lowerCAmelCase       :   str  =    hidden_size
      _lowerCAmelCase       :   Tuple  =    num_hidden_layers
      _lowerCAmelCase       :   Union[str, Any]  =    num_attention_heads
      _lowerCAmelCase       :   int  =    intermediate_size

      _lowerCAmelCase       :   List[Any]  =    hidden_dropout_prob
      _lowerCAmelCase       :   Union[str, Any]  =    attention_probs_dropout_prob
      _lowerCAmelCase       :   Optional[Any]  =    max_position_embeddings
      _lowerCAmelCase       :   List[str]  =    eos_token_id
      _lowerCAmelCase       :   Union[str, Any]  =    pad_token_id
      _lowerCAmelCase       :   List[Any]  =    bos_token_id
      _lowerCAmelCase       :   Dict  =    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
      _lowerCAmelCase       :   Tuple  =    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
      _lowerCAmelCase       :   int  =    (
          self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    ids_tensor([self.batch_size, self.seq_length - 1]    ,self.vocab_size      )
      _lowerCAmelCase       :   int  =    tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size      )    ,1      )
      _lowerCAmelCase       :   Any  =    tf.concat([input_ids, eos_tensor]    ,axis=1      )

      _lowerCAmelCase       :   Tuple  =    ids_tensor([self.batch_size, self.seq_length]    ,self.vocab_size      )

      _lowerCAmelCase       :   Any  =    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    ,)
      _lowerCAmelCase       :   Tuple  =    prepare_led_inputs_dict(_A    ,_A    ,_A      )
      _lowerCAmelCase       :   Dict  =    tf.concat(
          [tf.zeros_like(_A      )[:, :-1], tf.ones_like(_A      )[:, -1:]]    ,axis=-1    ,)
      _lowerCAmelCase       :   Any  =    global_attention_mask
      return config, inputs_dict







  def __lowerCamelCase		(		self    ,_A    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[int]  =    TFLEDModel(config=_A      ).get_decoder()
      _lowerCAmelCase       :   Tuple  =    inputs_dict['input_ids']

      _lowerCAmelCase       :   Union[str, Any]  =    input_ids[:1, :]
      _lowerCAmelCase       :   int  =    inputs_dict['attention_mask'][:1, :]
      _lowerCAmelCase       :   Dict  =    1

      # first forward pass
      _lowerCAmelCase       :   str  =    model(_A    ,attention_mask=_A    ,use_cache=_A      )

      _lowerCAmelCase,					_lowerCAmelCase       :   Dict  =    outputs.to_tuple()

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

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

      _lowerCAmelCase       :   Optional[int]  =    model(_A    ,attention_mask=_A      )[0]
      _lowerCAmelCase       :   int  =    model(_A    ,attention_mask=_A    ,past_key_values=_A      )[0]

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

      # select random slice
      _lowerCAmelCase       :   Tuple  =    int(ids_tensor((1,)    ,output_from_past.shape[-1]      )      )
      _lowerCAmelCase       :   Tuple  =    output_from_no_past[:, -3:, random_slice_idx]
      _lowerCAmelCase       :   Tuple  =    output_from_past[:, :, random_slice_idx]

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





def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase=None ,   _lowerCamelCase=None ,   _lowerCamelCase=None ,   _lowerCamelCase=None ,   ):

    '''simple docstring'''



    if attention_mask is None:
        _lowerCAmelCase       :   Optional[Any]  =    tf.cast(tf.math.not_equal(_lowerCamelCase ,   config.pad_token_id   ) ,   tf.inta   )
    if decoder_attention_mask is None:
        _lowerCAmelCase       :   List[str]  =    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:
        _lowerCAmelCase       :   Union[str, Any]  =    tf.ones((config.encoder_layers, config.encoder_attention_heads)   )
    if decoder_head_mask is None:
        _lowerCAmelCase       :   Dict  =    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    (    a__							,							a__							,							unittest.TestCase						):
  _UpperCAmelCase          =		(TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
  _UpperCAmelCase          =		(TFLEDForConditionalGeneration,) if is_tf_available() else ()
  _UpperCAmelCase          =		(
      {
          "conversational": TFLEDForConditionalGeneration,
          "feature-extraction": TFLEDModel,
          "summarization": TFLEDForConditionalGeneration,
          "text2text-generation": TFLEDForConditionalGeneration,
          "translation": TFLEDForConditionalGeneration,
      }
      if is_tf_available()
      else {}
  )
  _UpperCAmelCase          =		True
  _UpperCAmelCase          =		False
  _UpperCAmelCase          =		False
  _UpperCAmelCase          =		False

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    TFLEDModelTester(self      )
      _lowerCAmelCase       :   Optional[Any]  =    ConfigTester(self    ,config_class=_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      self.config_tester.run_common_tests()

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.check_decoder_model_past_large_inputs(*_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase,					_lowerCAmelCase       :   Dict  =    self.model_tester.prepare_config_and_inputs_for_common()
      _lowerCAmelCase       :   Tuple  =    tf.zeros_like(inputs_dict['attention_mask']      )
      _lowerCAmelCase       :   Tuple  =    2
      _lowerCAmelCase       :   Tuple  =    tf.where(
          tf.range(self.model_tester.seq_length      )[None, :] < num_global_attn_indices    ,1    ,inputs_dict['global_attention_mask']    ,)

      _lowerCAmelCase       :   int  =    True
      _lowerCAmelCase       :   int  =    self.model_tester.seq_length
      _lowerCAmelCase       :   Union[str, Any]  =    self.model_tester.encoder_seq_length

      def check_decoder_attentions_output(_A      ):
          _lowerCAmelCase       :   Optional[Any]  =    outputs.decoder_attentions
          self.assertEqual(len(_A      )    ,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(_A      ):
          _lowerCAmelCase       :   List[Any]  =    [t.numpy() for t in outputs.encoder_attentions]
          _lowerCAmelCase       :   Dict  =    [t.numpy() for t in outputs.encoder_global_attentions]
          self.assertEqual(len(_A      )    ,self.model_tester.num_hidden_layers      )
          self.assertEqual(len(_A      )    ,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:
          _lowerCAmelCase       :   Dict  =    True
          _lowerCAmelCase       :   Any  =    False
          _lowerCAmelCase       :   List[Any]  =    False
          _lowerCAmelCase       :   Any  =    model_class(_A      )
          _lowerCAmelCase       :   Union[str, Any]  =    model(self._prepare_for_class(_A    ,_A      )      )
          _lowerCAmelCase       :   List[str]  =    len(_A      )
          self.assertEqual(config.output_hidden_states    ,_A      )
          check_encoder_attentions_output(_A      )

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

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

          # Check attention is always last and order is fine
          _lowerCAmelCase       :   Any  =    True
          _lowerCAmelCase       :   Union[str, Any]  =    True
          _lowerCAmelCase       :   Union[str, Any]  =    model_class(_A      )
          _lowerCAmelCase       :   Optional[Any]  =    model(self._prepare_for_class(_A    ,_A      )      )

          self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)    ,len(_A      )      )
          self.assertEqual(model.config.output_hidden_states    ,_A      )
          check_encoder_attentions_output(_A      )

  @unittest.skip('LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'      )
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      pass





def 		lowerCamelCase__		(					_lowerCamelCase   ):

    '''simple docstring'''



    return tf.constant(_lowerCamelCase ,   dtype=tf.intaa   )


_lowerCAmelCase						=		1E-4







@slow
@require_tf
class 	__UpperCamelCase    (    unittest.TestCase						):

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Tuple  =    TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'      ).led

      # change to intended input here
      _lowerCAmelCase       :   Dict  =    _long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
      _lowerCAmelCase       :   Dict  =    _long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
      _lowerCAmelCase       :   Union[str, Any]  =    prepare_led_inputs_dict(model.config    ,_A    ,_A      )
      _lowerCAmelCase       :   Union[str, Any]  =    model(**_A      )[0]
      _lowerCAmelCase       :   Optional[Any]  =    (1, 1024, 768)
      self.assertEqual(output.shape    ,_A      )
      # change to expected output here
      _lowerCAmelCase       :   Tuple  =    tf.convert_to_tensor(
          [[2.3_0_5_0, 2.8_2_7_9, 0.6_5_3_1], [-1.8_4_5_7, -0.1_4_5_5, -3.5_6_6_1], [-1.0_1_8_6, 0.4_5_8_6, -2.2_0_4_3]]    ,)
      tf.debugging.assert_near(output[:, :3, :3]    ,_A    ,atol=1E-3      )







  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Any  =    TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'      )

      # change to intended input here
      _lowerCAmelCase       :   List[Any]  =    _long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
      _lowerCAmelCase       :   Any  =    _long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
      _lowerCAmelCase       :   List[str]  =    prepare_led_inputs_dict(model.config    ,_A    ,_A      )
      _lowerCAmelCase       :   List[str]  =    model(**_A      )[0]
      _lowerCAmelCase       :   List[Any]  =    (1, 1024, model.config.vocab_size)
      self.assertEqual(output.shape    ,_A      )
      # change to expected output here
      _lowerCAmelCase       :   Dict  =    tf.convert_to_tensor(
          [[3_3.6_5_0_7, 6.4_5_7_2, 1_6.8_0_8_9], [5.8_7_3_9, -2.4_2_3_8, 1_1.2_9_0_2], [-3.2_1_3_9, -4.3_1_4_9, 4.2_7_8_3]]    ,)
      tf.debugging.assert_near(output[:, :3, :3]    ,_A    ,atol=1E-3    ,rtol=1E-3      )








"""simple docstring"""






from dataclasses import dataclass
from typing import Dict, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin







@dataclass
class 	__UpperCamelCase    (    a__						):
  _UpperCAmelCase          =		42







class 	__UpperCamelCase    (    a__							,							a__						):

  @register_to_config
  def __init__(		self    ,_A = 32    ,_A = 64    ,_A = 20    ,_A = 768    ,_A=77    ,_A=4    ,_A = 0.0    ,_A = "silu"    ,_A = None    ,_A = None    ,_A = "linear"    ,_A = "prd"    ,_A = None    ,_A = None    ,_A = None    ,):
      '''simple docstring'''

      super().__init__()
      _lowerCAmelCase       :   Any  =    num_attention_heads
      _lowerCAmelCase       :   Optional[int]  =    attention_head_dim
      _lowerCAmelCase       :   Tuple  =    num_attention_heads * attention_head_dim
      _lowerCAmelCase       :   Optional[Any]  =    additional_embeddings

      _lowerCAmelCase       :   Union[str, Any]  =    time_embed_dim or inner_dim
      _lowerCAmelCase       :   Union[str, Any]  =    embedding_proj_dim or embedding_dim
      _lowerCAmelCase       :   Optional[int]  =    clip_embed_dim or embedding_dim

      _lowerCAmelCase       :   int  =    Timesteps(_A    ,_A    ,0      )
      _lowerCAmelCase       :   int  =    TimestepEmbedding(_A    ,_A    ,out_dim=_A    ,act_fn=_A      )

      _lowerCAmelCase       :   List[Any]  =    nn.Linear(_A    ,_A      )

      if embedding_proj_norm_type is None:
          _lowerCAmelCase       :   Optional[Any]  =    None
      elif embedding_proj_norm_type == "layer":
          _lowerCAmelCase       :   List[Any]  =    nn.LayerNorm(_A      )
      else:
          raise ValueError(F"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"""      )

      _lowerCAmelCase       :   Tuple  =    nn.Linear(_A    ,_A      )

      if encoder_hid_proj_type is None:
          _lowerCAmelCase       :   int  =    None
      elif encoder_hid_proj_type == "linear":
          _lowerCAmelCase       :   List[Any]  =    nn.Linear(_A    ,_A      )
      else:
          raise ValueError(F"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"""      )

      _lowerCAmelCase       :   Dict  =    nn.Parameter(torch.zeros(1    ,num_embeddings + additional_embeddings    ,_A      )      )

      if added_emb_type == "prd":
          _lowerCAmelCase       :   Dict  =    nn.Parameter(torch.zeros(1    ,1    ,_A      )      )
      elif added_emb_type is None:
          _lowerCAmelCase       :   List[Any]  =    None
      else:
          raise ValueError(
              F"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."""      )

      _lowerCAmelCase       :   List[Any]  =    nn.ModuleList(
          [
              BasicTransformerBlock(
                  _A    ,_A    ,_A    ,dropout=_A    ,activation_fn='gelu'    ,attention_bias=_A    ,)
              for d in range(_A      )
          ]      )

      if norm_in_type == "layer":
          _lowerCAmelCase       :   Any  =    nn.LayerNorm(_A      )
      elif norm_in_type is None:
          _lowerCAmelCase       :   Any  =    None
      else:
          raise ValueError(F"""Unsupported norm_in_type: {norm_in_type}."""      )

      _lowerCAmelCase       :   Union[str, Any]  =    nn.LayerNorm(_A      )

      _lowerCAmelCase       :   int  =    nn.Linear(_A    ,_A      )

      _lowerCAmelCase       :   Any  =    torch.full(
          [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]    ,-1_0_0_0_0.0      )
      causal_attention_mask.triu_(1      )
      _lowerCAmelCase       :   Tuple  =    causal_attention_mask[None, ...]
      self.register_buffer('causal_attention_mask'    ,_A    ,persistent=_A      )

      _lowerCAmelCase       :   Tuple  =    nn.Parameter(torch.zeros(1    ,_A      )      )
      _lowerCAmelCase       :   Dict  =    nn.Parameter(torch.zeros(1    ,_A      )      )

  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    {}

      def fn_recursive_add_processors(_A    ,_A    ,_A      ):
          if hasattr(_A    ,'set_processor'      ):
              _lowerCAmelCase       :   str  =    module.processor

          for sub_name, child in module.named_children():
              fn_recursive_add_processors(F"""{name}.{sub_name}"""    ,_A    ,_A      )

          return processors

      for name, module in self.named_children():
          fn_recursive_add_processors(_A    ,_A    ,_A      )

      return processors

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[int]  =    len(self.attn_processors.keys()      )

      if isinstance(_A    ,_A      ) and len(_A      ) != count:
          raise ValueError(
              F"""A dict of processors was passed, but the number of processors {len(_A      )} does not match the"""
              F""" number of attention layers: {count}. Please make sure to pass {count} processor classes."""      )

      def fn_recursive_attn_processor(_A    ,_A    ,_A      ):
          if hasattr(_A    ,'set_processor'      ):
              if not isinstance(_A    ,_A      ):
                  module.set_processor(_A      )
              else:
                  module.set_processor(processor.pop(F"""{name}.processor"""      )      )

          for sub_name, child in module.named_children():
              fn_recursive_attn_processor(F"""{name}.{sub_name}"""    ,_A    ,_A      )

      for name, module in self.named_children():
          fn_recursive_attn_processor(_A    ,_A    ,_A      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      self.set_attn_processor(AttnProcessor()      )

  def __lowerCamelCase		(		self    ,_A    ,_A    ,_A    ,_A = None    ,_A = None    ,_A = True    ,):
      '''simple docstring'''

      _lowerCAmelCase       :   str  =    hidden_states.shape[0]

      _lowerCAmelCase       :   int  =    timestep
      if not torch.is_tensor(_A      ):
          _lowerCAmelCase       :   str  =    torch.tensor([timesteps]    ,dtype=torch.long    ,device=hidden_states.device      )
      elif torch.is_tensor(_A      ) and len(timesteps.shape      ) == 0:
          _lowerCAmelCase       :   Dict  =    timesteps[None].to(hidden_states.device      )

      # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
      _lowerCAmelCase       :   Optional[int]  =    timesteps * torch.ones(_A    ,dtype=timesteps.dtype    ,device=timesteps.device      )

      _lowerCAmelCase       :   Dict  =    self.time_proj(_A      )

      # timesteps does not contain any weights and will always return f32 tensors
      # but time_embedding might be fp16, so we need to cast here.
      _lowerCAmelCase       :   Any  =    timesteps_projected.to(dtype=self.dtype      )
      _lowerCAmelCase       :   Optional[Any]  =    self.time_embedding(_A      )

      if self.embedding_proj_norm is not None:
          _lowerCAmelCase       :   int  =    self.embedding_proj_norm(_A      )

      _lowerCAmelCase       :   str  =    self.embedding_proj(_A      )
      if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
          _lowerCAmelCase       :   str  =    self.encoder_hidden_states_proj(_A      )
      elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
          raise ValueError('`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'      )

      _lowerCAmelCase       :   Any  =    self.proj_in(_A      )

      _lowerCAmelCase       :   Dict  =    self.positional_embedding.to(hidden_states.dtype      )

      _lowerCAmelCase       :   List[Any]  =    []
      _lowerCAmelCase       :   Optional[Any]  =    0

      if encoder_hidden_states is not None:
          additional_embeds.append(_A      )
          additional_embeddings_len += encoder_hidden_states.shape[1]

      if len(proj_embeddings.shape      ) == 2:
          _lowerCAmelCase       :   int  =    proj_embeddings[:, None, :]

      if len(hidden_states.shape      ) == 2:
          _lowerCAmelCase       :   Any  =    hidden_states[:, None, :]

      _lowerCAmelCase       :   int  =    additional_embeds + [
          proj_embeddings,
          time_embeddings[:, None, :],
          hidden_states,
      ]

      if self.prd_embedding is not None:
          _lowerCAmelCase       :   List[str]  =    self.prd_embedding.to(hidden_states.dtype      ).expand(_A    ,-1    ,-1      )
          additional_embeds.append(_A      )

      _lowerCAmelCase       :   List[str]  =    torch.cat(
          _A    ,dim=1    ,)

      # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
      _lowerCAmelCase       :   Tuple  =    additional_embeddings_len + proj_embeddings.shape[1] + 1
      if positional_embeddings.shape[1] < hidden_states.shape[1]:
          _lowerCAmelCase       :   Any  =    F.pad(
              _A    ,(
                  0,
                  0,
                  additional_embeddings_len,
                  self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
              )    ,value=0.0    ,)

      _lowerCAmelCase       :   int  =    hidden_states + positional_embeddings

      if attention_mask is not None:
          _lowerCAmelCase       :   Optional[Any]  =    (1 - attention_mask.to(hidden_states.dtype      )) * -1_0_0_0_0.0
          _lowerCAmelCase       :   Union[str, Any]  =    F.pad(_A    ,(0, self.additional_embeddings)    ,value=0.0      )
          _lowerCAmelCase       :   Tuple  =    (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype      )
          _lowerCAmelCase       :   Optional[Any]  =    attention_mask.repeat_interleave(self.config.num_attention_heads    ,dim=0      )

      if self.norm_in is not None:
          _lowerCAmelCase       :   Any  =    self.norm_in(_A      )

      for block in self.transformer_blocks:
          _lowerCAmelCase       :   int  =    block(_A    ,attention_mask=_A      )

      _lowerCAmelCase       :   Union[str, Any]  =    self.norm_out(_A      )

      if self.prd_embedding is not None:
          _lowerCAmelCase       :   Optional[int]  =    hidden_states[:, -1]
      else:
          _lowerCAmelCase       :   Any  =    hidden_states[:, additional_embeddings_len:]

      _lowerCAmelCase       :   Optional[int]  =    self.proj_to_clip_embeddings(_A      )

      if not return_dict:
          return (predicted_image_embedding,)

      return PriorTransformerOutput(predicted_image_embedding=_A      )







  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    (prior_latents * self.clip_std) + self.clip_mean
      return prior_latents





