





import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch






def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[Any]					)      ->    Tuple:  # picklable for multiprocessing

						"""simple docstring"""





						return x.sum()






def 		lowerCAmelCase  (	UpperCamelCase__   :   List[Any]					)      ->    Optional[int]:  # picklable for multiprocessing

						"""simple docstring"""





						return i + 1





@dataclass
class 				a			:
				SCREAMING_SNAKE_CASE__     :   int
				SCREAMING_SNAKE_CASE__     :   str





class 				a			(      __lowercase			):






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Tuple           =   {}
										__SCREAMING_SNAKE_CASE:				Dict           =   []
										__SCREAMING_SNAKE_CASE:				Dict           =   1
										__SCREAMING_SNAKE_CASE:				List[Any]           =   [1, 2]
										__SCREAMING_SNAKE_CASE:				List[Any]           =   {'''a''': 1, '''b''': 2}
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   {'''a''': [1, 2], '''b''': [3, 4]}
										__SCREAMING_SNAKE_CASE:				int           =   {'''a''': {'''1''': 1}, '''b''': 2}
										__SCREAMING_SNAKE_CASE:				List[Any]           =   {'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   {}
										__SCREAMING_SNAKE_CASE:				Tuple           =   []
										__SCREAMING_SNAKE_CASE:				str           =   2
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   [2, 3]
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   {'''a''': 2, '''b''': 3}
										__SCREAMING_SNAKE_CASE:				List[Any]           =   {'''a''': [2, 3], '''b''': [4, 5]}
										__SCREAMING_SNAKE_CASE:				List[Any]           =   {'''a''': {'''1''': 2}, '''b''': 3}
										__SCREAMING_SNAKE_CASE:				str           =   {'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase					)			, _lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   2
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, num_proc=_lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, num_proc=_lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, num_proc=_lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, num_proc=_lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, num_proc=_lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, num_proc=_lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, num_proc=_lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, num_proc=_lowerCAmelCase					)			, _lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				str           =   {'''a''': np.eye(2					), '''b''': np.zeros(3					), '''c''': np.ones(2					)}
										__SCREAMING_SNAKE_CASE:				List[Any]           =   {'''a''': 2, '''b''': 0, '''c''': 2}
										__SCREAMING_SNAKE_CASE:				Tuple           =   {
										    '''a''': np.eye(2					).astype(_lowerCAmelCase					),
										    '''b''': np.zeros(3					).astype(_lowerCAmelCase					),
										    '''c''': np.ones(2					).astype(_lowerCAmelCase					),
										}
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, map_numpy=_lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(
										    {k: v.tolist() for k, v in map_nested(_lowerCAmelCase			, _lowerCAmelCase			, map_numpy=_lowerCAmelCase					).items()}			, {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}			, )
										self.assertEqual(map_nested(_lowerCAmelCase			, _lowerCAmelCase			, map_numpy=_lowerCAmelCase			, num_proc=_lowerCAmelCase					)			, _lowerCAmelCase					)
										self.assertEqual(
										    {k: v.tolist() for k, v in map_nested(_lowerCAmelCase			, _lowerCAmelCase			, map_numpy=_lowerCAmelCase			, num_proc=_lowerCAmelCase					).items()}			, {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}			, )
										with self.assertRaises(_lowerCAmelCase					):  # can't pickle a local lambda
																map_nested(lambda _lowerCAmelCase					: x + 1			, _lowerCAmelCase			, num_proc=_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   {'''a''': 1, '''b''': 2}
										__SCREAMING_SNAKE_CASE:				int           =   {'''a''': 3, '''b''': 4}
										__SCREAMING_SNAKE_CASE:				List[Any]           =   {'''a''': 5, '''b''': 6}
										__SCREAMING_SNAKE_CASE:				List[Any]           =   sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))]					)
										self.assertEqual(sorted(zip_dict(_lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					)					)			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										class 				a			:
														SCREAMING_SNAKE_CASE__     :   Optional[Any]		      =     '''bar'''




										__SCREAMING_SNAKE_CASE:				Any           =   Foo()
										self.assertEqual(foo.my_attr			, '''bar'''					)
										with temporary_assignment(_lowerCAmelCase			, '''my_attr'''			, '''BAR'''					):
																self.assertEqual(foo.my_attr			, '''BAR'''					)
										self.assertEqual(foo.my_attr			, '''bar'''					)






@pytest.mark.parametrize(
    '''iterable_length, num_proc, expected_num_proc'''       ,			[
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]       ,			)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[Any]       ,			UpperCamelCase__   :   Any       ,			UpperCamelCase__   :   Union[str, Any]					)      ->    str:

						"""simple docstring"""





						with patch('''datasets.utils.py_utils._single_map_nested'''					) as mock_single_map_nested, patch(
						    '''datasets.parallel.parallel.Pool'''					) as mock_multiprocessing_pool:
												__SCREAMING_SNAKE_CASE:				List[str]           =   {F"""{i}""": i for i in range(UpperCamelCase__					)}
												__SCREAMING_SNAKE_CASE:				Dict           =   map_nested(lambda UpperCamelCase__					: x + 10       ,			UpperCamelCase__       ,			num_proc=UpperCamelCase__       ,			parallel_min_length=16					)
												if expected_num_proc == 1:
																		assert mock_single_map_nested.called
																		assert not mock_multiprocessing_pool.called
												else:
																		assert not mock_single_map_nested.called
																		assert mock_multiprocessing_pool.called
																		assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc





class 				a			(      __lowercase			):






				@require_tf
				def 			snake_case_    (  self					):

										"""simple docstring"""




										import tensorflow as tf
										from tensorflow.keras import layers

										__SCREAMING_SNAKE_CASE:				Optional[int]           =   layers.Dense(2					)

										def gen_random_output():
																__SCREAMING_SNAKE_CASE:				List[Any]           =   tf.random.uniform((1, 3)					)
																return model(_lowerCAmelCase					).numpy()

										with temp_seed(42			, set_tensorflow=_lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Any           =   gen_random_output()
										with temp_seed(42			, set_tensorflow=_lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   gen_random_output()
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   gen_random_output()

										np.testing.assert_equal(_lowerCAmelCase			, _lowerCAmelCase					)
										self.assertGreater(np.abs(outa - outa					).sum()			, 0					)






				@require_torch
				def 			snake_case_    (  self					):

										"""simple docstring"""




										import torch

										def gen_random_output():
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   torch.nn.Linear(3			, 2					)
																__SCREAMING_SNAKE_CASE:				List[Any]           =   torch.rand(1			, 3					)
																return model(_lowerCAmelCase					).detach().numpy()

										with temp_seed(42			, set_pytorch=_lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   gen_random_output()
										with temp_seed(42			, set_pytorch=_lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   gen_random_output()
										__SCREAMING_SNAKE_CASE:				Dict           =   gen_random_output()

										np.testing.assert_equal(_lowerCAmelCase			, _lowerCAmelCase					)
										self.assertGreater(np.abs(outa - outa					).sum()			, 0					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										def gen_random_output():
																return np.random.rand(1			, 3					)

										with temp_seed(42					):
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   gen_random_output()
										with temp_seed(42					):
																__SCREAMING_SNAKE_CASE:				Any           =   gen_random_output()
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   gen_random_output()

										np.testing.assert_equal(_lowerCAmelCase			, _lowerCAmelCase					)
										self.assertGreater(np.abs(outa - outa					).sum()			, 0					)






@pytest.mark.parametrize('''input_data'''       ,			[{}]					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Union[str, Any]					)      ->    Optional[Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[int]           =   NestedDataStructure(UpperCamelCase__					).data
						assert output_data == input_data






@pytest.mark.parametrize(
    '''data, expected_output'''       ,			[
        ({}, []),
        ([], []),
        ('''foo''', ['''foo''']),
        (['''foo''', '''bar'''], ['''foo''', '''bar''']),
        ([['''foo''', '''bar''']], ['''foo''', '''bar''']),
        ([[['''foo'''], ['''bar''']]], ['''foo''', '''bar''']),
        ([[['''foo'''], '''bar''']], ['''foo''', '''bar''']),
        ({'''a''': 1, '''b''': 2}, [1, 2]),
        ({'''a''': [1, 2], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[1, 2]], '''b''': [[3, 4]]}, [1, 2, 3, 4]),
        ({'''a''': [[1, 2]], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [[[3], [4]]]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [[3, 4]]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [3, [4]]}, [1, 2, 3, 4]),
        ({'''a''': {'''1''': 1}, '''b''': 2}, [1, 2]),
        ({'''a''': {'''1''': [1]}, '''b''': 2}, [1, 2]),
        ({'''a''': {'''1''': [1]}, '''b''': [2]}, [1, 2]),
    ]       ,			)
def 		lowerCAmelCase  (	UpperCamelCase__   :   int       ,			UpperCamelCase__   :   List[Any]					)      ->    Optional[int]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Tuple           =   NestedDataStructure(UpperCamelCase__					).flatten()
						assert output == expected_output






def 		lowerCAmelCase  (	)      ->    Dict:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				str           =   A(x=1       ,			y='''foobar'''					)
						__SCREAMING_SNAKE_CASE:				int           =   {'''x''': 1, '''y''': '''foobar'''}
						assert asdict(UpperCamelCase__					) == expected_output

						__SCREAMING_SNAKE_CASE:				List[str]           =   {'''a''': {'''b''': A(x=10       ,			y='''foo'''					)}, '''c''': [A(x=20       ,			y='''bar'''					)]}
						__SCREAMING_SNAKE_CASE:				int           =   {'''a''': {'''b''': {'''x''': 10, '''y''': '''foo'''}}, '''c''': [{'''x''': 20, '''y''': '''bar'''}]}
						assert asdict(UpperCamelCase__					) == expected_output

						with pytest.raises(UpperCamelCase__					):
												asdict([1, A(x=10       ,			y='''foo'''					)]					)






def 		lowerCAmelCase  (	UpperCamelCase__   :   str					)      ->    str:

						"""simple docstring"""





						return text.split()






def 		lowerCAmelCase  (	UpperCamelCase__   :   List[Any]					)      ->    Dict:

						"""simple docstring"""





						yield (time.time(), content)
						time.sleep(2					)
						yield (time.time(), content)






def 		lowerCAmelCase  (	)      ->    str:

						"""simple docstring"""





						with Pool(2					) as pool:
												__SCREAMING_SNAKE_CASE:				List[Any]           =   list(iflatmap_unordered(UpperCamelCase__       ,			_split_text       ,			kwargs_iterable=[{'''text''': '''hello there'''}] * 10					)					)
												assert out.count('''hello'''					) == 10
												assert out.count('''there'''					) == 10
												assert len(UpperCamelCase__					) == 20

						# check multiprocess from pathos (uses dill for pickling)
						with multiprocess.Pool(2					) as pool:
												__SCREAMING_SNAKE_CASE:				List[Any]           =   list(iflatmap_unordered(UpperCamelCase__       ,			_split_text       ,			kwargs_iterable=[{'''text''': '''hello there'''}] * 10					)					)
												assert out.count('''hello'''					) == 10
												assert out.count('''there'''					) == 10
												assert len(UpperCamelCase__					) == 20

						# check that we get items as fast as possible
						with Pool(2					) as pool:
												__SCREAMING_SNAKE_CASE:				Any           =   []
												for yield_time, content in iflatmap_unordered(
												    UpperCamelCase__       ,			_aseconds_generator_of_aitems_with_timing       ,			kwargs_iterable=[{'''content''': '''a'''}, {'''content''': '''b'''}]					):
																		assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
																		out.append(UpperCamelCase__					)
												assert out.count('''a'''					) == 2
												assert out.count('''b'''					) == 2
												assert len(UpperCamelCase__					) == 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					)



