



"""simple docstring"""






import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()






class 			__snake_case			(      _lowercase    ,	_lowercase    ,	unittest.TestCase):
  snake_case__							:  Tuple         =		AutoencoderKL
  snake_case__							:  Optional[int]         =		"sample"
  snake_case__							:  Optional[Any]         =		1e-2




  @property
  def 		SCREAMING_SNAKE_CASE							(					self	:		Union[str, Any]						):






         """simple docstring"""






         _lowerCamelCase		:							Optional[Any]            =					4
         _lowerCamelCase		:							Optional[Any]            =					3
         _lowerCamelCase		:							List[str]            =					(3_2, 3_2)

         _lowerCamelCase		:							Any            =					floats_tensor((batch_size, num_channels) + sizes						).to(__lowerCAmelCase						)

         return {"sample": image}




  @property
  def 		SCREAMING_SNAKE_CASE							(					self	:		List[Any]						):






         """simple docstring"""






         return (3, 3_2, 3_2)




  @property
  def 		SCREAMING_SNAKE_CASE							(					self	:		Dict						):






         """simple docstring"""






         return (3, 3_2, 3_2)




  def 		SCREAMING_SNAKE_CASE							(					self	:		Optional[int]						):






         """simple docstring"""






         _lowerCamelCase		:							Any            =					{
             '''block_out_channels''': [3_2, 6_4],
             '''in_channels''': 3,
             '''out_channels''': 3,
             '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
             '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
             '''latent_channels''': 4,
         }
         _lowerCamelCase		:							Tuple            =					self.dummy_input
         return init_dict, inputs_dict




  def 		SCREAMING_SNAKE_CASE							(					self	:		int						):






         """simple docstring"""






         pass




  def 		SCREAMING_SNAKE_CASE							(					self	:		Optional[int]						):






         """simple docstring"""






         pass




  @unittest.skipIf(torch_device == '''mps'''		,   '''Gradient checkpointing skipped on MPS'''						)
  def 		SCREAMING_SNAKE_CASE							(					self	:		List[Any]						):






         """simple docstring"""






         _lowerCamelCase     ,       _lowerCamelCase		:							Union[str, Any]            =					self.prepare_init_args_and_inputs_for_common()
         _lowerCamelCase		:							List[str]            =					self.model_class(**__lowerCAmelCase						)
         model.to(__lowerCAmelCase						)

         assert not model.is_gradient_checkpointing and model.training

         _lowerCamelCase		:							Union[str, Any]            =					model(**__lowerCAmelCase						).sample
         # run the backwards pass on the model. For backwards pass, for simplicity purpose,
         # we won't calculate the loss and rather backprop on out.sum()
         model.zero_grad()

         _lowerCamelCase		:							List[Any]            =					torch.randn_like(__lowerCAmelCase						)
         _lowerCamelCase		:							Tuple            =					(out - labels).mean()
         loss.backward()

         # re-instantiate the model now enabling gradient checkpointing
         _lowerCamelCase		:							List[Any]            =					self.model_class(**__lowerCAmelCase						)
         # clone model
         model_a.load_state_dict(model.state_dict()						)
         model_a.to(__lowerCAmelCase						)
         model_a.enable_gradient_checkpointing()

         assert model_a.is_gradient_checkpointing and model_a.training

         _lowerCamelCase		:							List[Any]            =					model_a(**__lowerCAmelCase						).sample
         # run the backwards pass on the model. For backwards pass, for simplicity purpose,
         # we won't calculate the loss and rather backprop on out.sum()
         model_a.zero_grad()
         _lowerCamelCase		:							Any            =					(out_a - labels).mean()
         loss_a.backward()

         # compare the output and parameters gradients
         self.assertTrue((loss - loss_a).abs() < 1E-5						)
         _lowerCamelCase		:							Optional[int]            =					dict(model.named_parameters()						)
         _lowerCamelCase		:							Any            =					dict(model_a.named_parameters()						)
         for name, param in named_params.items():
                self.assertTrue(torch_all_close(param.grad.data		,   named_params_a[name].grad.data		,   atol=5E-5						)						)




  def 		SCREAMING_SNAKE_CASE							(					self	:		str						):






         """simple docstring"""






         _lowerCamelCase     ,       _lowerCamelCase		:							List[Any]            =					AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''		,   output_loading_info=__lowerCAmelCase						)
         self.assertIsNotNone(__lowerCAmelCase						)
         self.assertEqual(len(loading_info['''missing_keys''']						)		,   0						)

         model.to(__lowerCAmelCase						)
         _lowerCamelCase		:							Union[str, Any]            =					model(**self.dummy_input						)

         assert image is not None, "Make sure output is not None"




  def 		SCREAMING_SNAKE_CASE							(					self	:		List[Any]						):






         """simple docstring"""






         _lowerCamelCase		:							List[Any]            =					AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''						)
         _lowerCamelCase		:							int            =					model.to(__lowerCAmelCase						)
         model.eval()

         if torch_device == "mps":
                _lowerCamelCase		:							int            =					torch.manual_seed(0						)
         else:
                _lowerCamelCase		:							str            =					torch.Generator(device=__lowerCAmelCase						).manual_seed(0						)

         _lowerCamelCase		:							int            =					torch.randn(
             1		,   model.config.in_channels		,   model.config.sample_size		,   model.config.sample_size		,   generator=torch.manual_seed(0						)		,   )
         _lowerCamelCase		:							Optional[int]            =					image.to(__lowerCAmelCase						)
         with torch.no_grad():
                _lowerCamelCase		:							str            =					model(__lowerCAmelCase		,   sample_posterior=__lowerCAmelCase		,   generator=__lowerCAmelCase						).sample

         _lowerCamelCase		:							List[str]            =					output[0, -1, -3:, -3:].flatten().cpu()

         # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
         # the expected output slices are not the same for CPU and GPU.
         if torch_device == "mps":
                _lowerCamelCase		:							str            =					torch.tensor(
                    [
                        -4.0078E-01,
                        -3.8323E-04,
                        -1.2681E-01,
                        -1.1462E-01,
                        2.0095E-01,
                        1.0893E-01,
                        -8.8247E-02,
                        -3.0361E-01,
                        -9.8644E-03,
                    ]						)
         elif torch_device == "cpu":
                _lowerCamelCase		:							Optional[int]            =					torch.tensor(
                    [-0.13_52, 0.08_78, 0.04_19, -0.08_18, -0.10_69, 0.06_88, -0.14_58, -0.44_46, -0.00_26]						)
         else:
                _lowerCamelCase		:							List[str]            =					torch.tensor(
                    [-0.24_21, 0.46_42, 0.25_07, -0.04_38, 0.06_82, 0.31_60, -0.20_18, -0.07_27, 0.24_85]						)

         self.assertTrue(torch_all_close(__lowerCAmelCase		,   __lowerCAmelCase		,   rtol=1E-2						)						)






@slow
class 			__snake_case			(      unittest.TestCase):




  def 		SCREAMING_SNAKE_CASE							(					self	:		Any		,   __lowerCAmelCase	:		Tuple		,   __lowerCAmelCase	:		Any						):






         """simple docstring"""






         return f'''gaussian_noise_s={seed}_shape={"_".join([str(__lowerCAmelCase						) for s in shape]						)}.npy'''




  def 		SCREAMING_SNAKE_CASE							(					self	:		str						):






         """simple docstring"""






         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()




  def 		SCREAMING_SNAKE_CASE							(					self	:		int		,   __lowerCAmelCase	:		Dict=0		,   __lowerCAmelCase	:		Tuple=(4, 3, 5_1_2, 5_1_2)		,   __lowerCAmelCase	:		Tuple=False						):






         """simple docstring"""






         _lowerCamelCase		:							str            =					torch.floataa if fpaa else torch.floataa
         _lowerCamelCase		:							List[Any]            =					torch.from_numpy(load_hf_numpy(self.get_file_format(__lowerCAmelCase		,   __lowerCAmelCase						)						)						).to(__lowerCAmelCase						).to(__lowerCAmelCase						)
         return image




  def 		SCREAMING_SNAKE_CASE							(					self	:		Dict		,   __lowerCAmelCase	:		List[str]="CompVis/stable-diffusion-v1-4"		,   __lowerCAmelCase	:		List[str]=False						):






         """simple docstring"""






         _lowerCamelCase		:							Any            =					'''fp16''' if fpaa else None
         _lowerCamelCase		:							str            =					torch.floataa if fpaa else torch.floataa

         _lowerCamelCase		:							Any            =					AutoencoderKL.from_pretrained(
             __lowerCAmelCase		,   subfolder='''vae'''		,   torch_dtype=__lowerCAmelCase		,   revision=__lowerCAmelCase		,   )
         model.to(__lowerCAmelCase						).eval()

         return model




  def 		SCREAMING_SNAKE_CASE							(					self	:		int		,   __lowerCAmelCase	:		int=0						):






         """simple docstring"""






         if torch_device == "mps":
                return torch.manual_seed(__lowerCAmelCase						)
         return torch.Generator(device=__lowerCAmelCase						).manual_seed(__lowerCAmelCase						)




  @parameterized.expand(
      [
          # fmt: off
          [3_3, [-0.16_03, 0.98_78, -0.04_95, -0.07_90, -0.27_09, 0.83_75, -0.20_60, -0.08_24], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
          [4_7, [-0.23_76, 0.11_68, 0.13_32, -0.48_40, -0.25_08, -0.07_91, -0.04_93, -0.40_89], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
          # fmt: on
      ]						)
  def 		SCREAMING_SNAKE_CASE							(					self	:		str		,   __lowerCAmelCase	:		Tuple		,   __lowerCAmelCase	:		List[Any]		,   __lowerCAmelCase	:		Union[str, Any]						):






         """simple docstring"""






         _lowerCamelCase		:							Dict            =					self.get_sd_vae_model()
         _lowerCamelCase		:							List[str]            =					self.get_sd_image(__lowerCAmelCase						)
         _lowerCamelCase		:							Optional[int]            =					self.get_generator(__lowerCAmelCase						)

         with torch.no_grad():
                _lowerCamelCase		:							Any            =					model(__lowerCAmelCase		,   generator=__lowerCAmelCase		,   sample_posterior=__lowerCAmelCase						).sample

         assert sample.shape == image.shape

         _lowerCamelCase		:							Optional[int]            =					sample[-1, -2:, -2:, :2].flatten().float().cpu()
         _lowerCamelCase		:							Tuple            =					torch.tensor(expected_slice_mps if torch_device == '''mps''' else expected_slice						)

         assert torch_all_close(__lowerCAmelCase		,   __lowerCAmelCase		,   atol=3E-3						)




  @parameterized.expand(
      [
          # fmt: off
          [3_3, [-0.05_13, 0.02_89, 1.37_99, 0.21_66, -0.25_73, -0.08_71, 0.51_03, -0.09_99]],
          [4_7, [-0.41_28, -0.13_20, -0.37_04, 0.19_65, -0.41_16, -0.23_32, -0.33_40, 0.22_47]],
          # fmt: on
      ]						)
  @require_torch_gpu
  def 		SCREAMING_SNAKE_CASE							(					self	:		Union[str, Any]		,   __lowerCAmelCase	:		Dict		,   __lowerCAmelCase	:		Tuple						):






         """simple docstring"""






         _lowerCamelCase		:							List[str]            =					self.get_sd_vae_model(fpaa=__lowerCAmelCase						)
         _lowerCamelCase		:							int            =					self.get_sd_image(__lowerCAmelCase		,   fpaa=__lowerCAmelCase						)
         _lowerCamelCase		:							Optional[int]            =					self.get_generator(__lowerCAmelCase						)

         with torch.no_grad():
                _lowerCamelCase		:							Optional[Any]            =					model(__lowerCAmelCase		,   generator=__lowerCAmelCase		,   sample_posterior=__lowerCAmelCase						).sample

         assert sample.shape == image.shape

         _lowerCamelCase		:							Union[str, Any]            =					sample[-1, -2:, :2, -2:].flatten().float().cpu()
         _lowerCamelCase		:							Optional[Any]            =					torch.tensor(__lowerCAmelCase						)

         assert torch_all_close(__lowerCAmelCase		,   __lowerCAmelCase		,   atol=1E-2						)




  @parameterized.expand(
      [
          # fmt: off
          [3_3, [-0.16_09, 0.98_66, -0.04_87, -0.07_77, -0.27_16, 0.83_68, -0.20_55, -0.08_14], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
          [4_7, [-0.23_77, 0.11_47, 0.13_33, -0.48_41, -0.25_06, -0.08_05, -0.04_91, -0.40_85], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
          # fmt: on
      ]						)
  def 		SCREAMING_SNAKE_CASE							(					self	:		Optional[Any]		,   __lowerCAmelCase	:		Any		,   __lowerCAmelCase	:		Tuple		,   __lowerCAmelCase	:		Union[str, Any]						):






         """simple docstring"""






         _lowerCamelCase		:							Union[str, Any]            =					self.get_sd_vae_model()
         _lowerCamelCase		:							int            =					self.get_sd_image(__lowerCAmelCase						)

         with torch.no_grad():
                _lowerCamelCase		:							List[str]            =					model(__lowerCAmelCase						).sample

         assert sample.shape == image.shape

         _lowerCamelCase		:							Optional[int]            =					sample[-1, -2:, -2:, :2].flatten().float().cpu()
         _lowerCamelCase		:							Tuple            =					torch.tensor(expected_slice_mps if torch_device == '''mps''' else expected_slice						)

         assert torch_all_close(__lowerCAmelCase		,   __lowerCAmelCase		,   atol=3E-3						)




  @parameterized.expand(
      [
          # fmt: off
          [1_3, [-0.20_51, -0.18_03, -0.23_11, -0.21_14, -0.32_92, -0.35_74, -0.29_53, -0.33_23]],
          [3_7, [-0.26_32, -0.26_25, -0.21_99, -0.27_41, -0.45_39, -0.49_90, -0.37_20, -0.49_25]],
          # fmt: on
      ]						)
  @require_torch_gpu
  def 		SCREAMING_SNAKE_CASE							(					self	:		Optional[Any]		,   __lowerCAmelCase	:		Union[str, Any]		,   __lowerCAmelCase	:		Optional[Any]						):






         """simple docstring"""






         _lowerCamelCase		:							Dict            =					self.get_sd_vae_model()
         _lowerCamelCase		:							Any            =					self.get_sd_image(__lowerCAmelCase		,   shape=(3, 4, 6_4, 6_4)						)

         with torch.no_grad():
                _lowerCamelCase		:							List[Any]            =					model.decode(__lowerCAmelCase						).sample

         assert list(sample.shape						) == [3, 3, 5_1_2, 5_1_2]

         _lowerCamelCase		:							Tuple            =					sample[-1, -2:, :2, -2:].flatten().cpu()
         _lowerCamelCase		:							List[Any]            =					torch.tensor(__lowerCAmelCase						)

         assert torch_all_close(__lowerCAmelCase		,   __lowerCAmelCase		,   atol=1E-3						)




  @parameterized.expand(
      [
          # fmt: off
          [2_7, [-0.03_69, 0.02_07, -0.07_76, -0.06_82, -0.17_47, -0.19_30, -0.14_65, -0.20_39]],
          [1_6, [-0.16_28, -0.21_34, -0.27_47, -0.26_42, -0.37_74, -0.44_04, -0.36_87, -0.42_77]],
          # fmt: on
      ]						)
  @require_torch_gpu
  def 		SCREAMING_SNAKE_CASE							(					self	:		Optional[Any]		,   __lowerCAmelCase	:		Optional[int]		,   __lowerCAmelCase	:		Any						):






         """simple docstring"""






         _lowerCamelCase		:							str            =					self.get_sd_vae_model(fpaa=__lowerCAmelCase						)
         _lowerCamelCase		:							int            =					self.get_sd_image(__lowerCAmelCase		,   shape=(3, 4, 6_4, 6_4)		,   fpaa=__lowerCAmelCase						)

         with torch.no_grad():
                _lowerCamelCase		:							Dict            =					model.decode(__lowerCAmelCase						).sample

         assert list(sample.shape						) == [3, 3, 5_1_2, 5_1_2]

         _lowerCamelCase		:							Dict            =					sample[-1, -2:, :2, -2:].flatten().float().cpu()
         _lowerCamelCase		:							Union[str, Any]            =					torch.tensor(__lowerCAmelCase						)

         assert torch_all_close(__lowerCAmelCase		,   __lowerCAmelCase		,   atol=5E-3						)




  @parameterized.expand([(1_3,), (1_6,), (2_7,)]						)
  @require_torch_gpu
  @unittest.skipIf(not is_xformers_available()		,   reason='''xformers is not required when using PyTorch 2.0.'''						)
  def 		SCREAMING_SNAKE_CASE							(					self	:		Any		,   __lowerCAmelCase	:		Optional[Any]						):






         """simple docstring"""






         _lowerCamelCase		:							str            =					self.get_sd_vae_model(fpaa=__lowerCAmelCase						)
         _lowerCamelCase		:							Union[str, Any]            =					self.get_sd_image(__lowerCAmelCase		,   shape=(3, 4, 6_4, 6_4)		,   fpaa=__lowerCAmelCase						)

         with torch.no_grad():
                _lowerCamelCase		:							List[str]            =					model.decode(__lowerCAmelCase						).sample

         model.enable_xformers_memory_efficient_attention()
         with torch.no_grad():
                _lowerCamelCase		:							Optional[int]            =					model.decode(__lowerCAmelCase						).sample

         assert list(sample.shape						) == [3, 3, 5_1_2, 5_1_2]

         assert torch_all_close(__lowerCAmelCase		,   __lowerCAmelCase		,   atol=1E-1						)




  @parameterized.expand([(1_3,), (1_6,), (3_7,)]						)
  @require_torch_gpu
  @unittest.skipIf(not is_xformers_available()		,   reason='''xformers is not required when using PyTorch 2.0.'''						)
  def 		SCREAMING_SNAKE_CASE							(					self	:		Optional[int]		,   __lowerCAmelCase	:		List[Any]						):






         """simple docstring"""






         _lowerCamelCase		:							Dict            =					self.get_sd_vae_model()
         _lowerCamelCase		:							List[str]            =					self.get_sd_image(__lowerCAmelCase		,   shape=(3, 4, 6_4, 6_4)						)

         with torch.no_grad():
                _lowerCamelCase		:							Optional[int]            =					model.decode(__lowerCAmelCase						).sample

         model.enable_xformers_memory_efficient_attention()
         with torch.no_grad():
                _lowerCamelCase		:							Any            =					model.decode(__lowerCAmelCase						).sample

         assert list(sample.shape						) == [3, 3, 5_1_2, 5_1_2]

         assert torch_all_close(__lowerCAmelCase		,   __lowerCAmelCase		,   atol=1E-2						)




  @parameterized.expand(
      [
          # fmt: off
          [3_3, [-0.30_01, 0.09_18, -2.69_84, -3.97_20, -3.20_99, -5.03_53, 1.73_38, -0.20_65, 3.42_67]],
          [4_7, [-1.50_30, -4.38_71, -6.03_55, -9.11_57, -1.66_61, -2.78_53, 2.16_07, -5.08_23, 2.56_33]],
          # fmt: on
      ]						)
  def 		SCREAMING_SNAKE_CASE							(					self	:		Any		,   __lowerCAmelCase	:		int		,   __lowerCAmelCase	:		Optional[Any]						):






         """simple docstring"""






         _lowerCamelCase		:							Optional[int]            =					self.get_sd_vae_model()
         _lowerCamelCase		:							Optional[Any]            =					self.get_sd_image(__lowerCAmelCase						)
         _lowerCamelCase		:							List[str]            =					self.get_generator(__lowerCAmelCase						)

         with torch.no_grad():
                _lowerCamelCase		:							Tuple            =					model.encode(__lowerCAmelCase						).latent_dist
                _lowerCamelCase		:							Tuple            =					dist.sample(generator=__lowerCAmelCase						)

         assert list(sample.shape						) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

         _lowerCamelCase		:							Optional[Any]            =					sample[0, -1, -3:, -3:].flatten().cpu()
         _lowerCamelCase		:							str            =					torch.tensor(__lowerCAmelCase						)

         _lowerCamelCase		:							Optional[int]            =					3E-3 if torch_device != '''mps''' else 1E-2
         assert torch_all_close(__lowerCAmelCase		,   __lowerCAmelCase		,   atol=__lowerCAmelCase						)






'''simple docstring'''
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel
from transformers.utils import ModelOutput
@dataclass
class   lowercase      (				A__   ):

		"""simple docstring"""

		_a       					=							None
		_a       					=							None
		_a       					=							None
		_a       					=							None



class   lowercase      (				A__   ):

		"""simple docstring"""

		def __init__(						self       ,  UpperCamelCase_=1       ,  UpperCamelCase_=0       ,  UpperCamelCase_=2       ,  UpperCamelCase_=512       ,  UpperCamelCase_="cls"       ,  UpperCamelCase_=False       ,  UpperCamelCase_=True       ,  **UpperCamelCase_       ,  ):

							'''simple docstring'''


							super().__init__(pad_token_id=UpperCamelCase_       ,  bos_token_id=UpperCamelCase_       ,  eos_token_id=UpperCamelCase_       ,  **UpperCamelCase_				)
							UpperCamelCase__  :Optional[int]				   =					project_dim
							UpperCamelCase__  :List[str]				   =					pooler_fn
							UpperCamelCase__  :Dict				   =					learn_encoder
							UpperCamelCase__  :Any				   =					use_attention_mask



class   lowercase      (				A__   ):

		"""simple docstring"""

		_a       					=							[R'pooler', R'logit_scale']
		_a       					=							[R'position_ids', R'predictions.decoder.bias']
		_a       					=							'roberta'
		_a       					=							RobertaSeriesConfig




		def __init__(						self       ,  UpperCamelCase_				):

							'''simple docstring'''


							super().__init__(UpperCamelCase_				)
							UpperCamelCase__  :Optional[Any]				   =					XLMRobertaModel(UpperCamelCase_				)
							UpperCamelCase__  :Tuple				   =					nn.Linear(config.hidden_size       ,  config.project_dim				)
							UpperCamelCase__  :Union[str, Any]				   =					getattr(UpperCamelCase_       ,  '''has_pre_transformation'''       ,  UpperCamelCase_				)
							if self.has_pre_transformation:
												UpperCamelCase__  :Any				   =					nn.Linear(config.hidden_size       ,  config.project_dim				)
												UpperCamelCase__  :str				   =					nn.LayerNorm(config.hidden_size       ,  eps=config.layer_norm_eps				)
							self.post_init()




		def        lowerCAmelCase__							(						self       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  UpperCamelCase_ = None       ,  ):

							'''simple docstring'''


							UpperCamelCase__  :Optional[int]				   =					return_dict if return_dict is not None else self.config.use_return_dict

							UpperCamelCase__  :str				   =					self.base_model(
							    input_ids=UpperCamelCase_       ,  attention_mask=UpperCamelCase_       ,  token_type_ids=UpperCamelCase_       ,  position_ids=UpperCamelCase_       ,  head_mask=UpperCamelCase_       ,  inputs_embeds=UpperCamelCase_       ,  encoder_hidden_states=UpperCamelCase_       ,  encoder_attention_mask=UpperCamelCase_       ,  output_attentions=UpperCamelCase_       ,  output_hidden_states=True if self.has_pre_transformation else output_hidden_states       ,  return_dict=UpperCamelCase_       ,  )

							if self.has_pre_transformation:
												UpperCamelCase__  :Dict				   =					outputs['''hidden_states'''][-2]
												UpperCamelCase__  :Optional[int]				   =					self.pre_LN(UpperCamelCase_				)
												UpperCamelCase__  :Optional[Any]				   =					self.transformation_pre(UpperCamelCase_				)

												return TransformationModelOutput(
												    projection_state=UpperCamelCase_       ,  last_hidden_state=outputs.last_hidden_state       ,  hidden_states=outputs.hidden_states       ,  attentions=outputs.attentions       ,  )
							else:
												UpperCamelCase__  :List[Any]				   =					self.transformation(outputs.last_hidden_state				)
												return TransformationModelOutput(
												    projection_state=UpperCamelCase_       ,  last_hidden_state=outputs.last_hidden_state       ,  hidden_states=outputs.hidden_states       ,  attentions=outputs.attentions       ,  )