




'''simple docstring'''

import json
import os
import shutil
import tempfile
import unittest

from transformers import BatchEncoding, CanineTokenizer
from transformers.testing_utils import require_tokenizers, require_torch
from transformers.tokenization_utils import AddedToken
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin






class 						a				(   _lowerCamelCase   ,     unittest.TestCase     ):
				snake_case_      =		CanineTokenizer
				snake_case_      =		False


				def   A_ (		self    :					List[Any]					):
									super().setUp()
									snake_case_									=				CanineTokenizer()
									tokenizer.save_pretrained(self.tmpdirname					)


				@cached_property
				def   A_ (		self    :					Tuple					):
									return CanineTokenizer.from_pretrained('''google/canine-s'''					)


				def   A_ (		self    :					Dict       ,				**lowercase_    :					List[Any]					):
									snake_case_									=				self.tokenizer_class.from_pretrained(self.tmpdirname       ,				**lowercase_					)
									snake_case_									=				1024
									return tokenizer


				@require_torch
				def   A_ (		self    :					Union[str, Any]					):
									snake_case_									=				self.canine_tokenizer
									snake_case_									=				['''Life is like a box of chocolates.''', '''You never know what you\'re gonna get.''']
									# fmt: off
									snake_case_									=				[5_7344, 76, 105, 102, 101, 32, 105, 115, 32, 108, 105, 107, 101, 32, 97, 32, 98, 111, 120, 32, 111, 102, 32, 99, 104, 111, 99, 111, 108, 97, 116, 101, 115, 46, 5_7345, 0, 0, 0, 0]
									# fmt: on
									snake_case_									=				tokenizer(lowercase_       ,				padding=lowercase_       ,				return_tensors='''pt'''					)
									self.assertIsInstance(lowercase_       ,				lowercase_					)

									snake_case_									=				list(batch.input_ids.numpy()[0]					)

									self.assertListEqual(lowercase_       ,				lowercase_					)

									self.assertEqual((2, 39)       ,				batch.input_ids.shape					)
									self.assertEqual((2, 39)       ,				batch.attention_mask.shape					)


				@require_torch
				def   A_ (		self    :					int					):
									snake_case_									=				self.canine_tokenizer
									snake_case_									=				['''Once there was a man.''', '''He wrote a test in HuggingFace Tranformers.''']
									snake_case_									=				tokenizer(lowercase_       ,				padding=lowercase_       ,				return_tensors='''pt'''					)
									# check if input_ids, attention_mask and token_type_ids are returned
									self.assertIn('''input_ids'''       ,				lowercase_					)
									self.assertIn('''attention_mask'''       ,				lowercase_					)
									self.assertIn('''token_type_ids'''       ,				lowercase_					)


				@require_torch
				def   A_ (		self    :					Dict					):
									snake_case_									=				self.canine_tokenizer
									snake_case_									=				[
									    '''What\'s the weater?''',
									    '''It\'s about 25 degrees.''',
									]
									snake_case_									=				tokenizer(
									    text_target=lowercase_       ,				max_length=32       ,				padding='''max_length'''       ,				truncation=lowercase_       ,				return_tensors='''pt'''					)
									self.assertEqual(32       ,				targets['''input_ids'''].shape[1]					)


				def   A_ (		self    :					Dict					):
									# safety check on max_len default value so we are sure the test works
									snake_case_									=				self.get_tokenizers()
									for tokenizer in tokenizers:
														with self.subTest(F"{tokenizer.__class__.__name__}"					):
																			self.assertNotEqual(tokenizer.model_max_length       ,				42					)

        # Now let's start the test
									snake_case_									=				self.get_tokenizers()
									for tokenizer in tokenizers:
														with self.subTest(F"{tokenizer.__class__.__name__}"					):
																			# Isolate this from the other tests because we save additional tokens/etc
																			snake_case_									=				tempfile.mkdtemp()

																			snake_case_									=				''' He is very happy, UNwant\u00E9d,running'''
																			snake_case_									=				tokenizer.encode(lowercase_       ,				add_special_tokens=lowercase_					)
																			tokenizer.save_pretrained(lowercase_					)

																			snake_case_									=				tokenizer.__class__.from_pretrained(lowercase_					)
																			snake_case_									=				after_tokenizer.encode(lowercase_       ,				add_special_tokens=lowercase_					)
																			self.assertListEqual(lowercase_       ,				lowercase_					)

																			shutil.rmtree(lowercase_					)

									snake_case_									=				self.get_tokenizers(model_max_length=42					)
									for tokenizer in tokenizers:
														with self.subTest(F"{tokenizer.__class__.__name__}"					):
																			# Isolate this from the other tests because we save additional tokens/etc
																			snake_case_									=				tempfile.mkdtemp()

																			snake_case_									=				''' He is very happy, UNwant\u00E9d,running'''

																			snake_case_									=				tokenizer.additional_special_tokens

																			# We can add a new special token for Canine as follows:
																			snake_case_									=				chr(0XE_0_0_7					)
																			additional_special_tokens.append(lowercase_					)
																			tokenizer.add_special_tokens({'''additional_special_tokens''': additional_special_tokens}					)
																			snake_case_									=				tokenizer.encode(lowercase_       ,				add_special_tokens=lowercase_					)
																			tokenizer.save_pretrained(lowercase_					)

																			snake_case_									=				tokenizer.__class__.from_pretrained(lowercase_					)
																			snake_case_									=				after_tokenizer.encode(lowercase_       ,				add_special_tokens=lowercase_					)
																			self.assertListEqual(lowercase_       ,				lowercase_					)
																			self.assertIn(lowercase_       ,				after_tokenizer.additional_special_tokens					)
																			self.assertEqual(after_tokenizer.model_max_length       ,				42					)

																			snake_case_									=				tokenizer.__class__.from_pretrained(lowercase_       ,				model_max_length=43					)
																			self.assertEqual(tokenizer.model_max_length       ,				43					)

																			shutil.rmtree(lowercase_					)


				def   A_ (		self    :					str					):
									snake_case_									=				self.get_tokenizers(do_lower_case=lowercase_					)
									for tokenizer in tokenizers:
														with self.subTest(F"{tokenizer.__class__.__name__}"					):
																			snake_case_			,snake_case_									=				self.get_clean_sequence(lowercase_					)

																			# a special token for Canine can be defined as follows:
																			snake_case_									=				0XE_0_0_5
																			snake_case_									=				chr(lowercase_					)

																			tokenizer.add_special_tokens({'''cls_token''': special_token}					)
																			snake_case_									=				tokenizer.encode(lowercase_       ,				add_special_tokens=lowercase_					)
																			self.assertEqual(len(lowercase_					)       ,				1					)

																			snake_case_									=				tokenizer.decode(ids + encoded_special_token       ,				clean_up_tokenization_spaces=lowercase_					)
																			snake_case_									=				tokenizer.encode(lowercase_       ,				add_special_tokens=lowercase_					)

																			snake_case_									=				tokenizer.encode(lowercase_       ,				add_special_tokens=lowercase_					)
																			snake_case_									=				tokenizer.encode(lowercase_       ,				add_special_tokens=lowercase_					)
																			self.assertEqual(lowercase_       ,				input_encoded + special_token_id					)

																			snake_case_									=				tokenizer.decode(lowercase_       ,				skip_special_tokens=lowercase_					)
																			self.assertTrue(special_token not in decoded					)


				def   A_ (		self    :					Dict					):
									snake_case_									=				self.get_tokenizers(do_lower_case=lowercase_					)
									for tokenizer in tokenizers:
														with self.subTest(F"{tokenizer.__class__.__name__}"					):
																			snake_case_									=				chr(0XE_0_0_5					)
																			snake_case_									=				chr(0XE_0_0_6					)

																			# `add_tokens` method stores special tokens only in `tokenizer.unique_no_split_tokens`. (in tokenization_utils.py)
																			tokenizer.add_tokens([SPECIAL_TOKEN_1]       ,				special_tokens=lowercase_					)
																			# `add_special_tokens` method stores special tokens in `tokenizer.additional_special_tokens`,
																			# which also occur in `tokenizer.all_special_tokens`. (in tokenization_utils_base.py)
																			tokenizer.add_special_tokens({'''additional_special_tokens''': [SPECIAL_TOKEN_2]}					)

																			snake_case_									=				tokenizer.tokenize(lowercase_					)
																			snake_case_									=				tokenizer.tokenize(lowercase_					)

																			self.assertEqual(len(lowercase_					)       ,				1					)
																			self.assertEqual(len(lowercase_					)       ,				1					)
																			self.assertEqual(token_a[0]       ,				lowercase_					)
																			self.assertEqual(token_a[0]       ,				lowercase_					)


				@require_tokenizers
				def   A_ (		self    :					List[Any]					):
									snake_case_									=				self.get_tokenizers(do_lower_case=lowercase_					)
									for tokenizer in tokenizers:
														with self.subTest(F"{tokenizer.__class__.__name__}"					):
																			# a special token for Canine can be defined as follows:
																			snake_case_									=				0XE_0_0_6
																			snake_case_									=				chr(lowercase_					)

																			snake_case_									=				AddedToken(lowercase_       ,				lstrip=lowercase_					)
																			tokenizer.add_special_tokens({'''additional_special_tokens''': [new_token]}					)

																			with tempfile.TemporaryDirectory() as tmp_dir_name:
																								tokenizer.save_pretrained(lowercase_					)
																								tokenizer.from_pretrained(lowercase_					)


				def   A_ (		self    :					List[Any]					):
									snake_case_									=				[]
									if self.test_slow_tokenizer:
														tokenizer_list.append((self.tokenizer_class, self.get_tokenizer())					)

									if self.test_rust_tokenizer:
														tokenizer_list.append((self.rust_tokenizer_class, self.get_rust_tokenizer())					)

									for tokenizer_class, tokenizer_utils in tokenizer_list:
														with tempfile.TemporaryDirectory() as tmp_dir:
																			tokenizer_utils.save_pretrained(lowercase_					)

																			with open(os.path.join(lowercase_       ,				'''special_tokens_map.json'''					)       ,				encoding='''utf-8'''					) as json_file:
																								snake_case_									=				json.load(lowercase_					)

																			with open(os.path.join(lowercase_       ,				'''tokenizer_config.json'''					)       ,				encoding='''utf-8'''					) as json_file:
																								snake_case_									=				json.load(lowercase_					)

																			# a special token for Canine can be defined as follows:
																			snake_case_									=				0XE_0_0_6
																			snake_case_									=				chr(lowercase_					)

																			snake_case_									=				[new_token_a]
																			snake_case_									=				[new_token_a]

																			with open(os.path.join(lowercase_       ,				'''special_tokens_map.json'''					)       ,				'''w'''       ,				encoding='''utf-8'''					) as outfile:
																								json.dump(lowercase_       ,				lowercase_					)
																			with open(os.path.join(lowercase_       ,				'''tokenizer_config.json'''					)       ,				'''w'''       ,				encoding='''utf-8'''					) as outfile:
																								json.dump(lowercase_       ,				lowercase_					)

																			# the following checks allow us to verify that our test works as expected, i.e. that the tokenizer takes
																			# into account the new value of additional_special_tokens given in the "tokenizer_config.json" and
																			# "special_tokens_map.json" files
																			snake_case_									=				tokenizer_class.from_pretrained(lowercase_       ,				extra_ids=0					)
																			self.assertIn(lowercase_       ,				tokenizer_without_change_in_init.additional_special_tokens					)
																			# self.assertIn("an_additional_special_token",tokenizer_without_change_in_init.get_vocab()) # ByT5Tokenization no vocab
																			self.assertEqual(
																			    [new_token_a]       ,				tokenizer_without_change_in_init.convert_ids_to_tokens(
																			        tokenizer_without_change_in_init.convert_tokens_to_ids([new_token_a]					)					)       ,				)

																			snake_case_									=				0XE_0_0_7
																			snake_case_									=				chr(lowercase_					)
																			# Now we test that we can change the value of additional_special_tokens in the from_pretrained
																			snake_case_									=				[AddedToken(lowercase_       ,				lstrip=lowercase_					)]
																			snake_case_									=				tokenizer_class.from_pretrained(
																			    lowercase_       ,				additional_special_tokens=lowercase_       ,				extra_ids=0					)

																			self.assertIn(lowercase_       ,				tokenizer.additional_special_tokens					)
																			# self.assertIn(new_token_2,tokenizer.get_vocab()) # ByT5Tokenization no vocab
																			self.assertEqual(
																			    [new_token_a]       ,				tokenizer.convert_ids_to_tokens(tokenizer.convert_tokens_to_ids([new_token_a]					)					)					)


				@require_tokenizers
				def   A_ (		self    :					Optional[Any]					):
									snake_case_									=				self.get_tokenizers(do_lower_case=lowercase_					)
									for tokenizer in tokenizers:
														with self.subTest(F"{tokenizer.__class__.__name__}"					):
																			snake_case_									=				'''hello world'''
																			if self.space_between_special_tokens:
																								snake_case_									=				'''[CLS] hello world [SEP]'''
																			else:
																								snake_case_									=				input
																			snake_case_									=				tokenizer.encode(lowercase_       ,				add_special_tokens=lowercase_					)
																			snake_case_									=				tokenizer.decode(lowercase_       ,				spaces_between_special_tokens=self.space_between_special_tokens					)
																			self.assertIn(lowercase_       ,				[output, output.lower()]					)


				def   A_ (		self    :					Any					):
									snake_case_									=				self.get_tokenizers()
									for tokenizer in tokenizers:
														with self.subTest(F"{tokenizer.__class__.__name__}"					):
																			snake_case_									=				[
																			    '''bos_token''',
																			    '''eos_token''',
																			    '''unk_token''',
																			    '''sep_token''',
																			    '''pad_token''',
																			    '''cls_token''',
																			    '''mask_token''',
																			]

									snake_case_									=				'''a'''
									snake_case_									=				ord(lowercase_					)

									for attr in attributes_list:
														setattr(lowercase_       ,				attr + '''_id'''       ,				lowercase_					)
														self.assertEqual(getattr(lowercase_       ,				lowercase_					)       ,				lowercase_					)
														self.assertEqual(getattr(lowercase_       ,				attr + '''_id'''					)       ,				lowercase_					)

														setattr(lowercase_       ,				attr + '''_id'''       ,				lowercase_					)
														self.assertEqual(getattr(lowercase_       ,				lowercase_					)       ,				lowercase_					)
														self.assertEqual(getattr(lowercase_       ,				attr + '''_id'''					)       ,				lowercase_					)

									setattr(lowercase_       ,				'''additional_special_tokens_ids'''       ,				[]					)
									self.assertListEqual(getattr(lowercase_       ,				'''additional_special_tokens'''					)       ,				[]					)
									self.assertListEqual(getattr(lowercase_       ,				'''additional_special_tokens_ids'''					)       ,				[]					)

									snake_case_									=				0XE_0_0_6
									snake_case_									=				chr(lowercase_					)
									setattr(lowercase_       ,				'''additional_special_tokens_ids'''       ,				[additional_special_token_id]					)
									self.assertListEqual(getattr(lowercase_       ,				'''additional_special_tokens'''					)       ,				[additional_special_token]					)
									self.assertListEqual(getattr(lowercase_       ,				'''additional_special_tokens_ids'''					)       ,				[additional_special_token_id]					)


				def   A_ (		self    :					Dict					):
									pass


				def   A_ (		self    :					Tuple					):
									pass


				def   A_ (		self    :					Optional[Any]					):
									pass


				def   A_ (		self    :					Optional[int]					):
									pass


				def   A_ (		self    :					Optional[Any]					):
									pass


				def   A_ (		self    :					Dict					):
									pass


				def   A_ (		self    :					Tuple					):
									pass


				def   A_ (		self    :					Tuple					):
									pass







'''simple docstring'''

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				(   _lowerCamelCase     ):
				snake_case_      =		42







@flax_register_to_config
class 						a				(   nn.Module   ,     _lowerCamelCase   ,     _lowerCamelCase     ):
				snake_case_      =		32
				snake_case_      =		4
				snake_case_      =		4
				snake_case_      =		(
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "DownBlock2D",
				)
				snake_case_      =		("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
				snake_case_      =		False
				snake_case_      =		(320, 640, 1_280, 1_280)
				snake_case_      =		2
				snake_case_      =		8
				snake_case_      =		None
				snake_case_      =		1_280
				snake_case_      =		0.0
				snake_case_      =		False
				snake_case_      =		jnp.floataa
				snake_case_      =		True
				snake_case_      =		0
				snake_case_      =		False


				def   A_ (		self    :					Optional[int]       ,				lowercase_    :					jax.random.KeyArray					):
									# init input tensors
									snake_case_									=				(1, self.in_channels, self.sample_size, self.sample_size)
									snake_case_									=				jnp.zeros(lowercase_       ,				dtype=jnp.floataa					)
									snake_case_									=				jnp.ones((1,)       ,				dtype=jnp.intaa					)
									snake_case_									=				jnp.zeros((1, 1, self.cross_attention_dim)       ,				dtype=jnp.floataa					)

									snake_case_			,snake_case_									=				jax.random.split(lowercase_					)
									snake_case_									=				{'''params''': params_rng, '''dropout''': dropout_rng}

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


				def   A_ (		self    :					List[str]					):
									snake_case_									=				self.block_out_channels
									snake_case_									=				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.
									snake_case_									=				self.num_attention_heads or self.attention_head_dim

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

									# time
									snake_case_									=				FlaxTimesteps(
									    block_out_channels[0]       ,				flip_sin_to_cos=self.flip_sin_to_cos       ,				freq_shift=self.config.freq_shift					)
									snake_case_									=				FlaxTimestepEmbedding(lowercase_       ,				dtype=self.dtype					)

									snake_case_									=				self.only_cross_attention
									if isinstance(lowercase_       ,				lowercase_					):
														snake_case_									=				(only_cross_attention,) * len(self.down_block_types					)

									if isinstance(lowercase_       ,				lowercase_					):
														snake_case_									=				(num_attention_heads,) * len(self.down_block_types					)

									# down
									snake_case_									=				[]
									snake_case_									=				block_out_channels[0]
									for i, down_block_type in enumerate(self.down_block_types					):
														snake_case_									=				output_channel
														snake_case_									=				block_out_channels[i]
														snake_case_									=				i == len(lowercase_					) - 1

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

														down_blocks.append(lowercase_					)
									snake_case_									=				down_blocks

									# mid
									snake_case_									=				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
									snake_case_									=				[]
									snake_case_									=				list(reversed(lowercase_					)					)
									snake_case_									=				list(reversed(lowercase_					)					)
									snake_case_									=				list(reversed(lowercase_					)					)
									snake_case_									=				reversed_block_out_channels[0]
									for i, up_block_type in enumerate(self.up_block_types					):
														snake_case_									=				output_channel
														snake_case_									=				reversed_block_out_channels[i]
														snake_case_									=				reversed_block_out_channels[min(i + 1       ,				len(lowercase_					) - 1					)]

														snake_case_									=				i == len(lowercase_					) - 1

														if up_block_type == "CrossAttnUpBlock2D":
																			snake_case_									=				FlaxCrossAttnUpBlockaD(
																			    in_channels=lowercase_       ,				out_channels=lowercase_       ,				prev_output_channel=lowercase_       ,				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:
																			snake_case_									=				FlaxUpBlockaD(
																			    in_channels=lowercase_       ,				out_channels=lowercase_       ,				prev_output_channel=lowercase_       ,				num_layers=self.layers_per_block + 1       ,				add_upsample=not is_final_block       ,				dropout=self.dropout       ,				dtype=self.dtype       ,				)

														up_blocks.append(lowercase_					)
														snake_case_									=				output_channel
									snake_case_									=				up_blocks

									# out
									snake_case_									=				nn.GroupNorm(num_groups=32       ,				epsilon=1e-5					)
									snake_case_									=				nn.Conv(
									    self.out_channels       ,				kernel_size=(3, 3)       ,				strides=(1, 1)       ,				padding=((1, 1), (1, 1))       ,				dtype=self.dtype       ,				)


				def __call__(		self    :					Union[str, Any]       ,				lowercase_    :					Union[str, Any]       ,				lowercase_    :					int       ,				lowercase_    :					Any       ,				lowercase_    :					int=None       ,				lowercase_    :					Any=None       ,				lowercase_    :					bool = True       ,				lowercase_    :					bool = False       ,				):
									# 1. time
									if not isinstance(lowercase_       ,				jnp.ndarray					):
														snake_case_									=				jnp.array([timesteps]       ,				dtype=jnp.intaa					)
									elif isinstance(lowercase_       ,				jnp.ndarray					) and len(timesteps.shape					) == 0:
														snake_case_									=				timesteps.astype(dtype=jnp.floataa					)
														snake_case_									=				jnp.expand_dims(lowercase_       ,				0					)

									snake_case_									=				self.time_proj(lowercase_					)
									snake_case_									=				self.time_embedding(lowercase_					)

									# 2. pre-process
									snake_case_									=				jnp.transpose(lowercase_       ,				(0, 2, 3, 1)					)
									snake_case_									=				self.conv_in(lowercase_					)

									# 3. down
									snake_case_									=				(sample,)
									for down_block in self.down_blocks:
														if isinstance(lowercase_       ,				lowercase_					):
																			snake_case_			,snake_case_									=				down_block(lowercase_       ,				lowercase_       ,				lowercase_       ,				deterministic=not train					)
														else:
																			snake_case_			,snake_case_									=				down_block(lowercase_       ,				lowercase_       ,				deterministic=not train					)
														down_block_res_samples += res_samples

									if down_block_additional_residuals is not None:
														snake_case_									=				()

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

														snake_case_									=				new_down_block_res_samples

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

        # 6. post-process
									snake_case_									=				self.conv_norm_out(lowercase_					)
									snake_case_									=				nn.silu(lowercase_					)
									snake_case_									=				self.conv_out(lowercase_					)
									snake_case_									=				jnp.transpose(lowercase_       ,				(0, 3, 1, 2)					)

									if not return_dict:
														return (sample,)

									return FlaxUNetaDConditionOutput(sample=lowercase_					)

