



'''simple docstring'''





import inspect
import unittest

from transformers import ViTMSNConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import ViTMSNForImageClassification, ViTMSNModel
    from transformers.models.vit_msn.modeling_vit_msn import VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import ViTImageProcessor



class  snake_case__						:
 def __init__(   self						:							Optional[Any]      ,     __a						:							Dict      ,     __a						:							Any=13      ,     __a						:							int=30      ,     __a						:							List[Any]=2      ,     __a						:							str=3      ,     __a						:							Dict=True      ,     __a						:							Optional[int]=True      ,     __a						:							Optional[int]=32      ,     __a						:							Union[str, Any]=5      ,     __a						:							int=4      ,     __a						:							Union[str, Any]=37      ,     __a						:							str="gelu"      ,     __a						:							str=0.1      ,     __a						:							Tuple=0.1      ,     __a						:							Optional[Any]=10      ,     __a						:							Dict=0.0_2      ,     __a						:							Union[str, Any]=None      ,     )      ->   str:







       '''simple docstring'''
       __snake_case :   Dict													=    parent
       __snake_case :   List[Any]													=    batch_size
       __snake_case :   Optional[int]													=    image_size
       __snake_case :   Tuple													=    patch_size
       __snake_case :   str													=    num_channels
       __snake_case :   int													=    is_training
       __snake_case :   List[Any]													=    use_labels
       __snake_case :   Optional[Any]													=    hidden_size
       __snake_case :   str													=    num_hidden_layers
       __snake_case :   int													=    num_attention_heads
       __snake_case :   List[str]													=    intermediate_size
       __snake_case :   str													=    hidden_act
       __snake_case :   str													=    hidden_dropout_prob
       __snake_case :   List[str]													=    attention_probs_dropout_prob
       __snake_case :   Optional[Any]													=    type_sequence_label_size
       __snake_case :   List[str]													=    initializer_range
       __snake_case :   List[str]													=    scope

       # in ViT MSN, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
       __snake_case :   Optional[int]													=    (image_size // patch_size) ** 2
       __snake_case :   Any													=    num_patches + 1
 def A_   (   self						:							str  )      ->   Any:







       '''simple docstring'''
       __snake_case :   List[str]													=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

       __snake_case :   Any													=    None
       if self.use_labels:
             __snake_case :   Optional[int]													=    ids_tensor([self.batch_size]      ,     self.type_sequence_label_size  )

       __snake_case :   int													=    self.get_config()

       return config, pixel_values, labels
 def A_   (   self						:							Any  )      ->   int:







       '''simple docstring'''
       return ViTMSNConfig(
           image_size=self.image_size      ,     patch_size=self.patch_size      ,     num_channels=self.num_channels      ,     hidden_size=self.hidden_size      ,     num_hidden_layers=self.num_hidden_layers      ,     num_attention_heads=self.num_attention_heads      ,     intermediate_size=self.intermediate_size      ,     hidden_act=self.hidden_act      ,     hidden_dropout_prob=self.hidden_dropout_prob      ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,     initializer_range=self.initializer_range      ,     )
 def A_   (   self						:							List[Any]      ,     __a						:							Optional[int]      ,     __a						:							Any      ,     __a						:							str  )      ->   Any:







       '''simple docstring'''
       __snake_case :   Tuple													=    ViTMSNModel(config=_lowerCAmelCase  )
       model.to(_lowerCAmelCase  )
       model.eval()
       __snake_case :   Any													=    model(_lowerCAmelCase  )
       self.parent.assertEqual(result.last_hidden_state.shape      ,     (self.batch_size, self.seq_length, self.hidden_size)  )
 def A_   (   self						:							Tuple      ,     __a						:							str      ,     __a						:							Tuple      ,     __a						:							Union[str, Any]  )      ->   Optional[int]:







       '''simple docstring'''
       __snake_case :   int													=    self.type_sequence_label_size
       __snake_case :   List[Any]													=    ViTMSNForImageClassification(_lowerCAmelCase  )
       model.to(_lowerCAmelCase  )
       model.eval()
       __snake_case :   Dict													=    model(_lowerCAmelCase      ,     labels=_lowerCAmelCase  )
       print('Pixel and labels shape: {pixel_values.shape}, {labels.shape}'  )
       print('Labels: {labels}'  )
       self.parent.assertEqual(result.logits.shape      ,     (self.batch_size, self.type_sequence_label_size)  )

       # test greyscale images
       __snake_case :   Optional[Any]													=    1
       __snake_case :   Union[str, Any]													=    ViTMSNForImageClassification(_lowerCAmelCase  )
       model.to(_lowerCAmelCase  )
       model.eval()

       __snake_case :   Union[str, Any]													=    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
       __snake_case :   int													=    model(_lowerCAmelCase  )
       self.parent.assertEqual(result.logits.shape      ,     (self.batch_size, self.type_sequence_label_size)  )





 def A_   (   self						:							Optional[Any]  )      ->   int:







       '''simple docstring'''
       __snake_case :   str													=    self.prepare_config_and_inputs()
       __snake_case :   List[Any]													=    config_and_inputs
       __snake_case :   int													=    {'''pixel_values''': pixel_values}
       return config, inputs_dict



@require_torch
class  snake_case__						(    __lowercase  ,     __lowercase  ,     unittest.TestCase				):
 A__				= (ViTMSNModel, ViTMSNForImageClassification) if is_torch_available() else ()
 A__				= (
     {'''feature-extraction''': ViTMSNModel, '''image-classification''': ViTMSNForImageClassification}
     if is_torch_available()
     else {}
 )

 A__				= False
 A__				= False
 A__				= False
 A__				= False
 def A_   (   self						:							List[Any]  )      ->   str:







       '''simple docstring'''
       __snake_case :   Optional[Any]													=    ViTMSNModelTester(self  )
       __snake_case :   str													=    ConfigTester(self      ,     config_class=_lowerCAmelCase      ,     has_text_modality=_lowerCAmelCase      ,     hidden_size=37  )
 def A_   (   self						:							Tuple  )      ->   Union[str, Any]:







       '''simple docstring'''
       self.config_tester.run_common_tests()
 @unittest.skip(reason='ViTMSN does not use inputs_embeds'  )
 def A_   (   self						:							str  )      ->   Dict:







       '''simple docstring'''
       pass
 def A_   (   self						:							str  )      ->   Optional[Any]:







       '''simple docstring'''
       __snake_case :   Union[str, Any]													=    self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             __snake_case :   List[Any]													=    model_class(_lowerCAmelCase  )
             self.assertIsInstance(model.get_input_embeddings()      ,     (nn.Module)  )
             __snake_case :   Union[str, Any]													=    model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(_lowerCAmelCase      ,     nn.Linear  )  )
 def A_   (   self						:							List[Any]  )      ->   List[Any]:







       '''simple docstring'''
       __snake_case :   List[str]													=    self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             __snake_case :   int													=    model_class(_lowerCAmelCase  )
             __snake_case :   Dict													=    inspect.signature(model.forward  )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             __snake_case :   Union[str, Any]													=    [*signature.parameters.keys()]

             __snake_case :   Optional[int]													=    ['''pixel_values''']
             self.assertListEqual(arg_names[:1]      ,     _lowerCAmelCase  )
 def A_   (   self						:							List[str]  )      ->   Dict:







       '''simple docstring'''
       __snake_case :   Tuple													=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_lowerCAmelCase  )
 def A_   (   self						:							Tuple  )      ->   Optional[Any]:







       '''simple docstring'''
       __snake_case :   str													=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_lowerCAmelCase  )





 @slow
 def A_   (   self						:							List[str]  )      ->   Optional[Any]:







       '''simple docstring'''




       for model_name in VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __snake_case :   Tuple													=    ViTMSNModel.from_pretrained(_lowerCAmelCase  )
             self.assertIsNotNone(_lowerCAmelCase  )





def   a_   (		)    ->  Any:
      __snake_case :   Any													=    Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
      return image





@require_torch
@require_vision
class  snake_case__						(    unittest.TestCase				):
 @cached_property
 def A_   (   self						:							List[str]  )      ->   Any:







       '''simple docstring'''
       return ViTImageProcessor.from_pretrained('facebook/vit-msn-small'  ) if is_vision_available() else None





 @slow
 def A_   (   self						:							List[str]  )      ->   str:







       '''simple docstring'''
       torch.manual_seed(2  )
       __snake_case :   Union[str, Any]													=    ViTMSNForImageClassification.from_pretrained('facebook/vit-msn-small'  ).to(_lowerCAmelCase  )

       __snake_case :   str													=    self.default_image_processor
       __snake_case :   Optional[Any]													=    prepare_img()
       __snake_case :   int													=    image_processor(images=_lowerCAmelCase      ,     return_tensors='pt'  ).to(_lowerCAmelCase  )

       # forward pass
       with torch.no_grad():
             __snake_case :   Any													=    model(**_lowerCAmelCase  )

       # verify the logits
       __snake_case :   Optional[Any]													=    torch.Size((1, 1000)  )
       self.assertEqual(outputs.logits.shape      ,     _lowerCAmelCase  )

       __snake_case :   List[Any]													=    torch.tensor([-0.0_8_0_3, -0.4_4_5_4, -0.2_3_7_5]  ).to(_lowerCAmelCase  )

       self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,     _lowerCAmelCase      ,     atol=1e-4  )  )










from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)





@flax.struct.dataclass
class 				a			(      __lowercase			):
				SCREAMING_SNAKE_CASE__     :   jnp.ndarray





@flax_register_to_config
class 				a			(      nn.Module   ,__lowercase   ,__lowercase			):
				SCREAMING_SNAKE_CASE__     :   int		      =     32
				SCREAMING_SNAKE_CASE__     :   int		      =     4
				SCREAMING_SNAKE_CASE__     :   int		      =     4
				SCREAMING_SNAKE_CASE__     :   Tuple[str]		      =     (
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "DownBlock2D",
				)
				SCREAMING_SNAKE_CASE__     :   Tuple[str]		      =     ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
				SCREAMING_SNAKE_CASE__     :   Union[bool, Tuple[bool]]		      =     False
				SCREAMING_SNAKE_CASE__     :   Tuple[int]		      =     (320, 640, 1280, 1280)
				SCREAMING_SNAKE_CASE__     :   int		      =     2
				SCREAMING_SNAKE_CASE__     :   Union[int, Tuple[int]]		      =     8
				SCREAMING_SNAKE_CASE__     :   Optional[Union[int, Tuple[int]]]		      =     None
				SCREAMING_SNAKE_CASE__     :   int		      =     1280
				SCREAMING_SNAKE_CASE__     :   float		      =     0.0
				SCREAMING_SNAKE_CASE__     :   bool		      =     False
				SCREAMING_SNAKE_CASE__     :   jnp.dtype		      =     jnp.floataa
				SCREAMING_SNAKE_CASE__     :   bool		      =     True
				SCREAMING_SNAKE_CASE__     :   int		      =     0
				SCREAMING_SNAKE_CASE__     :   bool		      =     False






				def 			snake_case_    (  self			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				int           =   (1, self.in_channels, self.sample_size, self.sample_size)
										__SCREAMING_SNAKE_CASE:				Tuple           =   jnp.zeros(_lowerCAmelCase			, dtype=jnp.floataa					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   jnp.ones((1,)			, dtype=jnp.intaa					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   jnp.zeros((1, 1, self.cross_attention_dim)			, dtype=jnp.floataa					)

										__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				Optional[int]           =   jax.random.split(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   {'''params''': params_rng, '''dropout''': dropout_rng}

										return self.init(_lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					)["params"]






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				int           =   self.block_out_channels
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   block_out_channels[0] * 4

										if self.num_attention_heads is not None:
																raise ValueError(
																    '''At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.'''					)

										# If `num_attention_heads` is not defined (which is the case for most models)
										# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
										# The reason for this behavior is to correct for incorrectly named variables that were introduced
										# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
										# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
										# which is why we correct for the naming here.
										__SCREAMING_SNAKE_CASE:				Any           =   self.num_attention_heads or self.attention_head_dim

										# input
										__SCREAMING_SNAKE_CASE:				str           =   nn.Conv(
										    block_out_channels[0]			, kernel_size=(3, 3)			, strides=(1, 1)			, padding=((1, 1), (1, 1))			, dtype=self.dtype			, )

										# time
										__SCREAMING_SNAKE_CASE:				int           =   FlaxTimesteps(
										    block_out_channels[0]			, flip_sin_to_cos=self.flip_sin_to_cos			, freq_shift=self.config.freq_shift					)
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   FlaxTimestepEmbedding(_lowerCAmelCase			, dtype=self.dtype					)

										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.only_cross_attention
										if isinstance(_lowerCAmelCase			, _lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   (only_cross_attention,) * len(self.down_block_types					)

										if isinstance(_lowerCAmelCase			, _lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Any           =   (num_attention_heads,) * len(self.down_block_types					)

										# down
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   []
										__SCREAMING_SNAKE_CASE:				List[str]           =   block_out_channels[0]
										for i, down_block_type in enumerate(self.down_block_types					):
																__SCREAMING_SNAKE_CASE:				List[str]           =   output_channel
																__SCREAMING_SNAKE_CASE:				str           =   block_out_channels[i]
																__SCREAMING_SNAKE_CASE:				Any           =   i == len(_lowerCAmelCase					) - 1

																if down_block_type == "CrossAttnDownBlock2D":
																						__SCREAMING_SNAKE_CASE:				str           =   FlaxCrossAttnDownBlockaD(
																						    in_channels=_lowerCAmelCase			, out_channels=_lowerCAmelCase			, dropout=self.dropout			, num_layers=self.layers_per_block			, num_attention_heads=num_attention_heads[i]			, add_downsample=not is_final_block			, use_linear_projection=self.use_linear_projection			, only_cross_attention=only_cross_attention[i]			, use_memory_efficient_attention=self.use_memory_efficient_attention			, dtype=self.dtype			, )
																else:
																						__SCREAMING_SNAKE_CASE:				Tuple           =   FlaxDownBlockaD(
																						    in_channels=_lowerCAmelCase			, out_channels=_lowerCAmelCase			, dropout=self.dropout			, num_layers=self.layers_per_block			, add_downsample=not is_final_block			, dtype=self.dtype			, )

																down_blocks.append(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				str           =   down_blocks

										# mid
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   FlaxUNetMidBlockaDCrossAttn(
										    in_channels=block_out_channels[-1]			, dropout=self.dropout			, num_attention_heads=num_attention_heads[-1]			, use_linear_projection=self.use_linear_projection			, use_memory_efficient_attention=self.use_memory_efficient_attention			, dtype=self.dtype			, )

										# up
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   []
										__SCREAMING_SNAKE_CASE:				Tuple           =   list(reversed(_lowerCAmelCase					)					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   list(reversed(_lowerCAmelCase					)					)
										__SCREAMING_SNAKE_CASE:				str           =   list(reversed(_lowerCAmelCase					)					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   reversed_block_out_channels[0]
										for i, up_block_type in enumerate(self.up_block_types					):
																__SCREAMING_SNAKE_CASE:				int           =   output_channel
																__SCREAMING_SNAKE_CASE:				List[str]           =   reversed_block_out_channels[i]
																__SCREAMING_SNAKE_CASE:				List[str]           =   reversed_block_out_channels[min(i + 1			, len(_lowerCAmelCase					) - 1					)]

																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   i == len(_lowerCAmelCase					) - 1

																if up_block_type == "CrossAttnUpBlock2D":
																						__SCREAMING_SNAKE_CASE:				Optional[int]           =   FlaxCrossAttnUpBlockaD(
																						    in_channels=_lowerCAmelCase			, out_channels=_lowerCAmelCase			, prev_output_channel=_lowerCAmelCase			, num_layers=self.layers_per_block + 1			, num_attention_heads=reversed_num_attention_heads[i]			, add_upsample=not is_final_block			, dropout=self.dropout			, use_linear_projection=self.use_linear_projection			, only_cross_attention=only_cross_attention[i]			, use_memory_efficient_attention=self.use_memory_efficient_attention			, dtype=self.dtype			, )
																else:
																						__SCREAMING_SNAKE_CASE:				int           =   FlaxUpBlockaD(
																						    in_channels=_lowerCAmelCase			, out_channels=_lowerCAmelCase			, prev_output_channel=_lowerCAmelCase			, num_layers=self.layers_per_block + 1			, add_upsample=not is_final_block			, dropout=self.dropout			, dtype=self.dtype			, )

																up_blocks.append(_lowerCAmelCase					)
																__SCREAMING_SNAKE_CASE:				Any           =   output_channel
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   up_blocks

										# out
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   nn.GroupNorm(num_groups=32			, epsilon=1e-5					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   nn.Conv(
										    self.out_channels			, kernel_size=(3, 3)			, strides=(1, 1)			, padding=((1, 1), (1, 1))			, dtype=self.dtype			, )






				def __call__(  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase=None			, _lowerCAmelCase=None			, _lowerCAmelCase = True			, _lowerCAmelCase = False			, ):

										"""simple docstring"""




										if not isinstance(_lowerCAmelCase			, jnp.ndarray					):
																__SCREAMING_SNAKE_CASE:				Dict           =   jnp.array([timesteps]			, dtype=jnp.intaa					)
										elif isinstance(_lowerCAmelCase			, jnp.ndarray					) and len(timesteps.shape					) == 0:
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   timesteps.astype(dtype=jnp.floataa					)
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   jnp.expand_dims(_lowerCAmelCase			, 0					)

										__SCREAMING_SNAKE_CASE:				Any           =   self.time_proj(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   self.time_embedding(_lowerCAmelCase					)

										# 2. pre-process
										__SCREAMING_SNAKE_CASE:				int           =   jnp.transpose(_lowerCAmelCase			, (0, 2, 3, 1)					)
										__SCREAMING_SNAKE_CASE:				List[str]           =   self.conv_in(_lowerCAmelCase					)

										# 3. down
										__SCREAMING_SNAKE_CASE:				Dict           =   (sample,)
										for down_block in self.down_blocks:
																if isinstance(_lowerCAmelCase			, _lowerCAmelCase					):
																						__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				Optional[Any]           =   down_block(_lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, deterministic=not train					)
																else:
																						__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				Tuple           =   down_block(_lowerCAmelCase			, _lowerCAmelCase			, deterministic=not train					)
																down_block_res_samples += res_samples

										if down_block_additional_residuals is not None:
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   ()

																for down_block_res_sample, down_block_additional_residual in zip(
																    _lowerCAmelCase			, _lowerCAmelCase					):
																						down_block_res_sample += down_block_additional_residual
																						new_down_block_res_samples += (down_block_res_sample,)

																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   new_down_block_res_samples

										# 4. mid
										__SCREAMING_SNAKE_CASE:				Dict           =   self.mid_block(_lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, deterministic=not train					)

										if mid_block_additional_residual is not None:
																sample += mid_block_additional_residual

										# 5. up
										for up_block in self.up_blocks:
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   down_block_res_samples[-(self.layers_per_block + 1) :]
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   down_block_res_samples[: -(self.layers_per_block + 1)]
																if isinstance(_lowerCAmelCase			, _lowerCAmelCase					):
																						__SCREAMING_SNAKE_CASE:				Any           =   up_block(
																						    _lowerCAmelCase			, temb=_lowerCAmelCase			, encoder_hidden_states=_lowerCAmelCase			, res_hidden_states_tuple=_lowerCAmelCase			, deterministic=not train			, )
																else:
																						__SCREAMING_SNAKE_CASE:				List[str]           =   up_block(_lowerCAmelCase			, temb=_lowerCAmelCase			, res_hidden_states_tuple=_lowerCAmelCase			, deterministic=not train					)

        # 6. post-process
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   self.conv_norm_out(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   nn.silu(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Tuple           =   self.conv_out(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				int           =   jnp.transpose(_lowerCAmelCase			, (0, 3, 1, 2)					)

										if not return_dict:
																return (sample,)

										return FlaxUNetaDConditionOutput(sample=_lowerCAmelCase					)



