





import inspect
import unittest

from transformers import MobileNetVaConfig
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 transformers import MobileNetVaForImageClassification, MobileNetVaModel
    from transformers.models.mobilenet_va.modeling_mobilenet_va import MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import MobileNetVaImageProcessor

class      lowerCAmelCase							(      __UpperCamelCase    ):
 def   A_			(   self						:   List[Any]			)    ->			Dict:
     lowerCamelCase__					:   int      				=	self.config_class(**self.inputs_dict			)
     self.parent.assertTrue(hasattr(UpperCAmelCase    ,  'tf_padding'			)			)
     self.parent.assertTrue(hasattr(UpperCAmelCase    ,  'depth_multiplier'			)			)


class      lowerCAmelCase							:
 def __init__(   self						:   Dict    ,  UpperCAmelCase						:   Tuple    ,  UpperCAmelCase						:   Any=13    ,  UpperCAmelCase						:   Union[str, Any]=3    ,  UpperCAmelCase						:   Optional[int]=32    ,  UpperCAmelCase						:   List[str]=0.2_5    ,  UpperCAmelCase						:   Any=8    ,  UpperCAmelCase						:   int=True    ,  UpperCAmelCase						:   Optional[Any]=1024    ,  UpperCAmelCase						:   Tuple=32    ,  UpperCAmelCase						:   str="relu6"    ,  UpperCAmelCase						:   List[Any]=0.1    ,  UpperCAmelCase						:   Any=0.0_2    ,  UpperCAmelCase						:   Tuple=True    ,  UpperCAmelCase						:   Optional[Any]=True    ,  UpperCAmelCase						:   Optional[int]=10    ,  UpperCAmelCase						:   List[Any]=None    ,  )    ->			Union[str, Any]:
     lowerCamelCase__					:   List[Any]      				=	parent
     lowerCamelCase__					:   List[Any]      				=	batch_size
     lowerCamelCase__					:   List[str]      				=	num_channels
     lowerCamelCase__					:   Optional[int]      				=	image_size
     lowerCamelCase__					:   str      				=	depth_multiplier
     lowerCamelCase__					:   Optional[int]      				=	min_depth
     lowerCamelCase__					:   List[str]      				=	tf_padding
     lowerCamelCase__					:   Optional[int]      				=	int(last_hidden_size * depth_multiplier			)
     lowerCamelCase__					:   Dict      				=	output_stride
     lowerCamelCase__					:   Dict      				=	hidden_act
     lowerCamelCase__					:   Optional[int]      				=	classifier_dropout_prob
     lowerCamelCase__					:   int      				=	use_labels
     lowerCamelCase__					:   Tuple      				=	is_training
     lowerCamelCase__					:   Union[str, Any]      				=	num_labels
     lowerCamelCase__					:   Optional[Any]      				=	initializer_range
     lowerCamelCase__					:   Tuple      				=	scope
 def   A_			(   self						:   Optional[int]			)    ->			Union[str, Any]:
     lowerCamelCase__					:   Any      				=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

     lowerCamelCase__					:   Any      				=	None
     lowerCamelCase__					:   Any      				=	None
     if self.use_labels:
         lowerCamelCase__					:   Optional[int]      				=	ids_tensor([self.batch_size]    ,  self.num_labels			)
         lowerCamelCase__					:   Any      				=	ids_tensor([self.batch_size, self.image_size, self.image_size]    ,  self.num_labels			)

     lowerCamelCase__					:   List[Any]      				=	self.get_config()

     return config, pixel_values, labels, pixel_labels
 def   A_			(   self						:   Union[str, Any]			)    ->			Tuple:
     return MobileNetVaConfig(
         num_channels=self.num_channels    ,  image_size=self.image_size    ,  depth_multiplier=self.depth_multiplier    ,  min_depth=self.min_depth    ,  tf_padding=self.tf_padding    ,  hidden_act=self.hidden_act    ,  classifier_dropout_prob=self.classifier_dropout_prob    ,  initializer_range=self.initializer_range    ,  )
 def   A_			(   self						:   Optional[Any]    ,  UpperCAmelCase						:   List[Any]    ,  UpperCAmelCase						:   Tuple    ,  UpperCAmelCase						:   Dict    ,  UpperCAmelCase						:   str			)    ->			Tuple:
     lowerCamelCase__					:   Union[str, Any]      				=	MobileNetVaModel(config=UpperCAmelCase			)
     model.to(UpperCAmelCase			)
     model.eval()
     lowerCamelCase__					:   Optional[int]      				=	model(UpperCAmelCase			)
     self.parent.assertEqual(
         result.last_hidden_state.shape    ,  (
             self.batch_size,
             self.last_hidden_size,
             self.image_size // self.output_stride,
             self.image_size // self.output_stride,
         )    ,  )
 def   A_			(   self						:   Any    ,  UpperCAmelCase						:   Dict    ,  UpperCAmelCase						:   Optional[int]    ,  UpperCAmelCase						:   Union[str, Any]    ,  UpperCAmelCase						:   Tuple			)    ->			Optional[Any]:
     lowerCamelCase__					:   Optional[Any]      				=	self.num_labels
     lowerCamelCase__					:   Optional[Any]      				=	MobileNetVaForImageClassification(UpperCAmelCase			)
     model.to(UpperCAmelCase			)
     model.eval()
     lowerCamelCase__					:   Dict      				=	model(UpperCAmelCase    ,  labels=UpperCAmelCase			)
     self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.num_labels)			)







 def   A_			(   self						:   str			)    ->			Any:
     lowerCamelCase__					:   str      				=	self.prepare_config_and_inputs()
     lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__					:   int      				=	config_and_inputs
     lowerCamelCase__					:   Tuple      				=	{'pixel_values': pixel_values}
     return config, inputs_dict


@require_torch
class      lowerCAmelCase							(      __UpperCamelCase,						__UpperCamelCase,						unittest.TestCase    ):
 UpperCAmelCase__									=	(MobileNetVaModel, MobileNetVaForImageClassification) if is_torch_available() else ()
 UpperCAmelCase__									=	(
     {"""feature-extraction""": MobileNetVaModel, """image-classification""": MobileNetVaForImageClassification}
     if is_torch_available()
     else {}
 )

 UpperCAmelCase__									=	False
 UpperCAmelCase__									=	False
 UpperCAmelCase__									=	False
 UpperCAmelCase__									=	False
 def   A_			(   self						:   Optional[int]			)    ->			Optional[Any]:
     lowerCamelCase__					:   str      				=	MobileNetVaModelTester(self			)
     lowerCamelCase__					:   Any      				=	MobileNetVaConfigTester(self    ,  config_class=UpperCAmelCase    ,  has_text_modality=UpperCAmelCase			)
 def   A_			(   self						:   Optional[Any]			)    ->			List[str]:
     self.config_tester.run_common_tests()
 @unittest.skip(reason='MobileNetV1 does not use inputs_embeds'			)
 def   A_			(   self						:   Union[str, Any]			)    ->			Union[str, Any]:
     pass
 @unittest.skip(reason='MobileNetV1 does not support input and output embeddings'			)
 def   A_			(   self						:   Dict			)    ->			str:
     pass
 @unittest.skip(reason='MobileNetV1 does not output attentions'			)
 def   A_			(   self						:   Tuple			)    ->			int:
     pass
 def   A_			(   self						:   Union[str, Any]			)    ->			Union[str, Any]:
     lowerCamelCase__				,							lowerCamelCase__					:   Optional[int]      				=	self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         lowerCamelCase__					:   str      				=	model_class(UpperCAmelCase			)
         lowerCamelCase__					:   Optional[int]      				=	inspect.signature(model.forward			)
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         lowerCamelCase__					:   Dict      				=	[*signature.parameters.keys()]

         lowerCamelCase__					:   int      				=	['pixel_values']
         self.assertListEqual(arg_names[:1]    ,  UpperCAmelCase			)
 def   A_			(   self						:   int			)    ->			List[str]:
     lowerCamelCase__					:   List[Any]      				=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*UpperCAmelCase			)
 def   A_			(   self						:   Any			)    ->			Optional[int]:
     def check_hidden_states_output(UpperCAmelCase						:   Optional[int]    ,  UpperCAmelCase						:   Dict    ,  UpperCAmelCase						:   List[Any]			):
         lowerCamelCase__					:   Union[str, Any]      				=	model_class(UpperCAmelCase			)
         model.to(UpperCAmelCase			)
         model.eval()

         with torch.no_grad():
             lowerCamelCase__					:   Optional[Any]      				=	model(**self._prepare_for_class(UpperCAmelCase    ,  UpperCAmelCase			)			)

         lowerCamelCase__					:   str      				=	outputs.hidden_states

         lowerCamelCase__					:   Union[str, Any]      				=	26
         self.assertEqual(len(UpperCAmelCase			)    ,  UpperCAmelCase			)

     lowerCamelCase__				,							lowerCamelCase__					:   int      				=	self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         lowerCamelCase__					:   List[str]      				=	True
         check_hidden_states_output(UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase			)

         # check that output_hidden_states also work using config
         del inputs_dict["output_hidden_states"]
         lowerCamelCase__					:   List[str]      				=	True

         check_hidden_states_output(UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase			)
 def   A_			(   self						:   Any			)    ->			Optional[Any]:
     lowerCamelCase__					:   List[str]      				=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase			)







 @slow
 def   A_			(   self						:   List[Any]			)    ->			int:
     for model_name in MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowerCamelCase__					:   Any      				=	MobileNetVaModel.from_pretrained(UpperCAmelCase			)
         self.assertIsNotNone(UpperCAmelCase			)


def  SCREAMING_SNAKE_CASE   ( )    ->  int:
    lowerCamelCase__					:   str      				=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
    return image


@require_torch
@require_vision
class      lowerCAmelCase							(      unittest.TestCase    ):
 @cached_property
 def   A_			(   self						:   List[str]			)    ->			Optional[Any]:
     return (
         MobileNetVaImageProcessor.from_pretrained('google/mobilenet_v1_1.0_224'			) if is_vision_available() else None
     )







 @slow
 def   A_			(   self						:   List[str]			)    ->			Dict:
     lowerCamelCase__					:   Union[str, Any]      				=	MobileNetVaForImageClassification.from_pretrained('google/mobilenet_v1_1.0_224'			).to(UpperCAmelCase			)

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

     # forward pass
     with torch.no_grad():
         lowerCamelCase__					:   str      				=	model(**UpperCAmelCase			)

     # verify the logits
     lowerCamelCase__					:   int      				=	torch.Size((1, 1001)			)
     self.assertEqual(outputs.logits.shape    ,  UpperCAmelCase			)

     lowerCamelCase__					:   Optional[int]      				=	torch.tensor([-4.1_7_3_9, -1.1_2_3_3, 3.1_2_0_5]			).to(UpperCAmelCase			)

     self.assertTrue(torch.allclose(outputs.logits[0, :3]    ,  UpperCAmelCase    ,  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      lowerCAmelCase							(      __UpperCamelCase    ):
 UpperCAmelCase__									=	42


@flax_register_to_config
class      lowerCAmelCase							(      nn.Module,						__UpperCamelCase,						__UpperCamelCase    ):
 UpperCAmelCase__									=	32
 UpperCAmelCase__									=	4
 UpperCAmelCase__									=	4
 UpperCAmelCase__									=	(
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "DownBlock2D",
 )
 UpperCAmelCase__									=	("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
 UpperCAmelCase__									=	False
 UpperCAmelCase__									=	(3_20, 6_40, 12_80, 12_80)
 UpperCAmelCase__									=	2
 UpperCAmelCase__									=	8
 UpperCAmelCase__									=	None
 UpperCAmelCase__									=	12_80
 UpperCAmelCase__									=	0.0
 UpperCAmelCase__									=	False
 UpperCAmelCase__									=	jnp.floataa
 UpperCAmelCase__									=	True
 UpperCAmelCase__									=	0
 UpperCAmelCase__									=	False
 def   A_			(   self						:   Tuple    ,  UpperCAmelCase						:   jax.random.KeyArray			)    ->			FrozenDict:
     # init input tensors
     lowerCamelCase__					:   int      				=	(1, self.in_channels, self.sample_size, self.sample_size)
     lowerCamelCase__					:   List[str]      				=	jnp.zeros(UpperCAmelCase    ,  dtype=jnp.floataa			)
     lowerCamelCase__					:   Tuple      				=	jnp.ones((1,)    ,  dtype=jnp.intaa			)
     lowerCamelCase__					:   Dict      				=	jnp.zeros((1, 1, self.cross_attention_dim)    ,  dtype=jnp.floataa			)

     lowerCamelCase__				,							lowerCamelCase__					:   Union[str, Any]      				=	jax.random.split(UpperCAmelCase			)
     lowerCamelCase__					:   Dict      				=	{'params': params_rng, 'dropout': dropout_rng}

     return self.init(UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase			)["params"]
 def   A_			(   self						:   Tuple			)    ->			Optional[int]:
     lowerCamelCase__					:   Any      				=	self.block_out_channels
     lowerCamelCase__					:   int      				=	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.
     lowerCamelCase__					:   Tuple      				=	self.num_attention_heads or self.attention_head_dim

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

     # time
     lowerCamelCase__					:   Optional[int]      				=	FlaxTimesteps(
         block_out_channels[0]    ,  flip_sin_to_cos=self.flip_sin_to_cos    ,  freq_shift=self.config.freq_shift			)
     lowerCamelCase__					:   int      				=	FlaxTimestepEmbedding(UpperCAmelCase    ,  dtype=self.dtype			)

     lowerCamelCase__					:   Optional[int]      				=	self.only_cross_attention
     if isinstance(UpperCAmelCase    ,  UpperCAmelCase			):
         lowerCamelCase__					:   str      				=	(only_cross_attention,) * len(self.down_block_types			)

     if isinstance(UpperCAmelCase    ,  UpperCAmelCase			):
         lowerCamelCase__					:   List[Any]      				=	(num_attention_heads,) * len(self.down_block_types			)

     # down
     lowerCamelCase__					:   List[Any]      				=	[]
     lowerCamelCase__					:   Dict      				=	block_out_channels[0]
     for i, down_block_type in enumerate(self.down_block_types			):
         lowerCamelCase__					:   Dict      				=	output_channel
         lowerCamelCase__					:   Optional[int]      				=	block_out_channels[i]
         lowerCamelCase__					:   List[Any]      				=	i == len(UpperCAmelCase			) - 1

         if down_block_type == "CrossAttnDownBlock2D":
             lowerCamelCase__					:   Tuple      				=	FlaxCrossAttnDownBlockaD(
                 in_channels=UpperCAmelCase    ,  out_channels=UpperCAmelCase    ,  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:
             lowerCamelCase__					:   str      				=	FlaxDownBlockaD(
                 in_channels=UpperCAmelCase    ,  out_channels=UpperCAmelCase    ,  dropout=self.dropout    ,  num_layers=self.layers_per_block    ,  add_downsample=not is_final_block    ,  dtype=self.dtype    ,  )

         down_blocks.append(UpperCAmelCase			)
     lowerCamelCase__					:   List[Any]      				=	down_blocks

     # mid
     lowerCamelCase__					:   Dict      				=	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
     lowerCamelCase__					:   Any      				=	[]
     lowerCamelCase__					:   Optional[int]      				=	list(reversed(UpperCAmelCase			)			)
     lowerCamelCase__					:   Any      				=	list(reversed(UpperCAmelCase			)			)
     lowerCamelCase__					:   int      				=	list(reversed(UpperCAmelCase			)			)
     lowerCamelCase__					:   Tuple      				=	reversed_block_out_channels[0]
     for i, up_block_type in enumerate(self.up_block_types			):
         lowerCamelCase__					:   str      				=	output_channel
         lowerCamelCase__					:   int      				=	reversed_block_out_channels[i]
         lowerCamelCase__					:   int      				=	reversed_block_out_channels[min(i + 1    ,  len(UpperCAmelCase			) - 1			)]

         lowerCamelCase__					:   Optional[Any]      				=	i == len(UpperCAmelCase			) - 1

         if up_block_type == "CrossAttnUpBlock2D":
             lowerCamelCase__					:   Tuple      				=	FlaxCrossAttnUpBlockaD(
                 in_channels=UpperCAmelCase    ,  out_channels=UpperCAmelCase    ,  prev_output_channel=UpperCAmelCase    ,  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:
             lowerCamelCase__					:   Optional[Any]      				=	FlaxUpBlockaD(
                 in_channels=UpperCAmelCase    ,  out_channels=UpperCAmelCase    ,  prev_output_channel=UpperCAmelCase    ,  num_layers=self.layers_per_block + 1    ,  add_upsample=not is_final_block    ,  dropout=self.dropout    ,  dtype=self.dtype    ,  )

         up_blocks.append(UpperCAmelCase			)
         lowerCamelCase__					:   Tuple      				=	output_channel
     lowerCamelCase__					:   Tuple      				=	up_blocks

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







 def __call__(   self						:   int    ,  UpperCAmelCase						:   Any    ,  UpperCAmelCase						:   Tuple    ,  UpperCAmelCase						:   Optional[Any]    ,  UpperCAmelCase						:   Optional[Any]=None    ,  UpperCAmelCase						:   Dict=None    ,  UpperCAmelCase						:   bool = True    ,  UpperCAmelCase						:   bool = False    ,  )    ->			Union[FlaxUNetaDConditionOutput, Tuple]:
     # 1. time
     if not isinstance(UpperCAmelCase    ,  jnp.ndarray			):
         lowerCamelCase__					:   List[str]      				=	jnp.array([timesteps]    ,  dtype=jnp.intaa			)
     elif isinstance(UpperCAmelCase    ,  jnp.ndarray			) and len(timesteps.shape			) == 0:
         lowerCamelCase__					:   List[Any]      				=	timesteps.astype(dtype=jnp.floataa			)
         lowerCamelCase__					:   Any      				=	jnp.expand_dims(UpperCAmelCase    ,  0			)

     lowerCamelCase__					:   List[str]      				=	self.time_proj(UpperCAmelCase			)
     lowerCamelCase__					:   Optional[Any]      				=	self.time_embedding(UpperCAmelCase			)

     # 2. pre-process
     lowerCamelCase__					:   Dict      				=	jnp.transpose(UpperCAmelCase    ,  (0, 2, 3, 1)			)
     lowerCamelCase__					:   Optional[Any]      				=	self.conv_in(UpperCAmelCase			)

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

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

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

         lowerCamelCase__					:   str      				=	new_down_block_res_samples

     # 4. mid
     lowerCamelCase__					:   List[Any]      				=	self.mid_block(UpperCAmelCase    ,  UpperCAmelCase    ,  UpperCAmelCase    ,  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:
         lowerCamelCase__					:   str      				=	down_block_res_samples[-(self.layers_per_block + 1) :]
         lowerCamelCase__					:   List[str]      				=	down_block_res_samples[: -(self.layers_per_block + 1)]
         if isinstance(UpperCAmelCase    ,  UpperCAmelCase			):
             lowerCamelCase__					:   List[Any]      				=	up_block(
                 UpperCAmelCase    ,  temb=UpperCAmelCase    ,  encoder_hidden_states=UpperCAmelCase    ,  res_hidden_states_tuple=UpperCAmelCase    ,  deterministic=not train    ,  )
         else:
             lowerCamelCase__					:   int      				=	up_block(UpperCAmelCase    ,  temb=UpperCAmelCase    ,  res_hidden_states_tuple=UpperCAmelCase    ,  deterministic=not train			)

        # 6. post-process
     lowerCamelCase__					:   str      				=	self.conv_norm_out(UpperCAmelCase			)
     lowerCamelCase__					:   Optional[Any]      				=	nn.silu(UpperCAmelCase			)
     lowerCamelCase__					:   Any      				=	self.conv_out(UpperCAmelCase			)
     lowerCamelCase__					:   Optional[Any]      				=	jnp.transpose(UpperCAmelCase    ,  (0, 3, 1, 2)			)

     if not return_dict:
         return (sample,)

     return FlaxUNetaDConditionOutput(sample=UpperCAmelCase			)





