






"""simple docstring"""







import pytest

import datasets


# Import fixture modules as plugins
__a       :	Tuple             = ['tests.fixtures.files', 'tests.fixtures.hub', 'tests.fixtures.fsspec']





def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_ ,     lowerCamelCase_):
					# Mark tests as "unit" by default if not marked as "integration" (or already marked as "unit")
					for item in items:
										if any(marker in item.keywords for marker in ['''integration''', '''unit''']):
															continue
										item.add_marker(pytest.mark.unit)





def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_):
					config.addinivalue_line('''markers''' ,     '''torchaudio_latest: mark test to run with torchaudio>=0.12''')





@pytest.fixture(autouse=_SCREAMING_SNAKE_CASE)
def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_ ,     lowerCamelCase_):
					# test_hf_cache_home = tmp_path_factory.mktemp("cache")  # TODO: why a cache dir per test function does not work?
					a__          =  tmp_path_factory.getbasetemp() / '''cache'''
					a__          =  test_hf_cache_home / '''datasets'''
					a__          =  test_hf_cache_home / '''metrics'''
					a__          =  test_hf_cache_home / '''modules'''
					monkeypatch.setattr('''datasets.config.HF_DATASETS_CACHE''' ,     str(_SCREAMING_SNAKE_CASE))
					monkeypatch.setattr('''datasets.config.HF_METRICS_CACHE''' ,     str(_SCREAMING_SNAKE_CASE))
					monkeypatch.setattr('''datasets.config.HF_MODULES_CACHE''' ,     str(_SCREAMING_SNAKE_CASE))
					a__          =  test_hf_datasets_cache / '''downloads'''
					monkeypatch.setattr('''datasets.config.DOWNLOADED_DATASETS_PATH''' ,     str(_SCREAMING_SNAKE_CASE))
					a__          =  test_hf_datasets_cache / '''downloads''' / '''extracted'''
					monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_PATH''' ,     str(_SCREAMING_SNAKE_CASE))





@pytest.fixture(autouse=_SCREAMING_SNAKE_CASE ,     scope='''session''')
def        SCREAMING_SNAKE_CASE       (      ):
					datasets.disable_progress_bar()





@pytest.fixture(autouse=_SCREAMING_SNAKE_CASE)
def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_):
					# don't take tests into account when counting downloads
					monkeypatch.setattr('''datasets.config.HF_UPDATE_DOWNLOAD_COUNTS''' ,     _SCREAMING_SNAKE_CASE)





@pytest.fixture
def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_):
					# Required to suppress RemovedIn20Warning when feature(s) are not compatible with SQLAlchemy 2.0
					# To be removed once SQLAlchemy 2.0 supported
					monkeypatch.setattr('''sqlalchemy.util.deprecations.SILENCE_UBER_WARNING''' ,     _SCREAMING_SNAKE_CASE)










"""simple docstring"""







import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.out_channels				,		n_heads=self.num_attention_heads				,		d_head=self.out_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		only_cross_attention=self.only_cross_attention				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

												a__          =  resnets
												a__          =  attentions

												if self.add_downsample:
																	a__          =  FlaxDownsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Optional[int]				,		__A:     Union[str, Any]				,		__A:     str				,		__A:     Optional[Any]				,		__A:     Any=True    ):

												'''simple docstring'''

												a__          =  ()

												for resnet, attn in zip(self.resnets				,		self.attentions    ):
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )
																	output_states += (hidden_states,)

												if self.add_downsample:
																	a__          =  self.downsamplers_a(__A    )
																	output_states += (hidden_states,)

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )
												a__          =  resnets

												if self.add_downsample:
																	a__          =  FlaxDownsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Dict				,		__A:     int				,		__A:     Dict				,		__A:     Optional[Any]=True    ):

												'''simple docstring'''

												a__          =  ()

												for resnet in self.resnets:
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	output_states += (hidden_states,)

												if self.add_downsample:
																	a__          =  self.downsamplers_a(__A    )
																	output_states += (hidden_states,)

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Optional[int]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if (i == self.num_layers - 1) else self.out_channels
																	a__          =  self.prev_output_channel if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=resnet_in_channels + res_skip_channels				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.out_channels				,		n_heads=self.num_attention_heads				,		d_head=self.out_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		only_cross_attention=self.only_cross_attention				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

												a__          =  resnets
												a__          =  attentions

												if self.add_upsample:
																	a__          =  FlaxUpsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Any				,		__A:     Optional[int]				,		__A:     List[Any]				,		__A:     List[str]				,		__A:     Optional[Any]				,		__A:     Any=True    ):

												'''simple docstring'''

												for resnet, attn in zip(self.resnets				,		self.attentions    ):
																	# pop res hidden states
																	a__          =  res_hidden_states_tuple[-1]
																	a__          =  res_hidden_states_tuple[:-1]
																	a__          =  jnp.concatenate((hidden_states, res_hidden_states)				,		axis=-1    )

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     str    ):

												'''simple docstring'''

												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if (i == self.num_layers - 1) else self.out_channels
																	a__          =  self.prev_output_channel if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=resnet_in_channels + res_skip_channels				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

												a__          =  resnets

												if self.add_upsample:
																	a__          =  FlaxUpsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Tuple				,		__A:     Optional[Any]				,		__A:     Optional[Any]				,		__A:     Union[str, Any]				,		__A:     Dict=True    ):

												'''simple docstring'''

												for resnet in self.resnets:
																	# pop res hidden states
																	a__          =  res_hidden_states_tuple[-1]
																	a__          =  res_hidden_states_tuple[:-1]
																	a__          =  jnp.concatenate((hidden_states, res_hidden_states)				,		axis=-1    )

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Tuple    ):

												'''simple docstring'''

												a__          =  [
												    FlaxResnetBlockaD(
												        in_channels=self.in_channels				,		out_channels=self.in_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
												]

												a__          =  []

												for _ in range(self.num_layers    ):
																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.in_channels				,		n_heads=self.num_attention_heads				,		d_head=self.in_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=self.in_channels				,		out_channels=self.in_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

												a__          =  resnets
												a__          =  attentions





							def __call__(   self:     Any				,		__A:     Optional[int]				,		__A:     int				,		__A:     Tuple				,		__A:     str=True    ):

												'''simple docstring'''

												a__          =  self.resnets[0](__A				,		__A    )
												for attn, resnet in zip(self.attentions				,		self.resnets[1:]    ):
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )

												return hidden_states


