


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        __lowercase       (							__lowerCAmelCase  ):



		"""simple docstring"""

		UpperCamelCase     :					jnp.ndarray







@flax_register_to_config
class        __lowercase       (							nn.Module			,							__lowerCAmelCase			,							__lowerCAmelCase  ):



		"""simple docstring"""

		UpperCamelCase     :					int			       =		3_2
		UpperCamelCase     :					int			       =		4
		UpperCamelCase     :					int			       =		4
		UpperCamelCase     :					Tuple[str]			       =		(
		    "CrossAttnDownBlock2D",
		    "CrossAttnDownBlock2D",
		    "CrossAttnDownBlock2D",
		    "DownBlock2D",
		)
		UpperCamelCase     :					Tuple[str]			       =		("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
		UpperCamelCase     :					Union[bool, Tuple[bool]]			       =		False
		UpperCamelCase     :					Tuple[int]			       =		(3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
		UpperCamelCase     :					int			       =		2
		UpperCamelCase     :					Union[int, Tuple[int]]			       =		8
		UpperCamelCase     :					Optional[Union[int, Tuple[int]]]			       =		None
		UpperCamelCase     :					int			       =		1_2_8_0
		UpperCamelCase     :					float			       =		0.0
		UpperCamelCase     :					bool			       =		False
		UpperCamelCase     :					jnp.dtype			       =		jnp.floataa
		UpperCamelCase     :					bool			       =		True
		UpperCamelCase     :					int			       =		0
		UpperCamelCase     :					bool			       =		False







		def 					__A   ( self		,      A	)  -> FrozenDict:







					'''simple docstring'''







					lowerCamelCase    		=  (1, self.in_channels, self.sample_size, self.sample_size)
					lowerCamelCase    		=  jnp.zeros(__SCREAMING_SNAKE_CASE		,      dtype=jnp.floataa	)
					lowerCamelCase    		=  jnp.ones((1,)		,      dtype=jnp.intaa	)
					lowerCamelCase    		=  jnp.zeros((1, 1, self.cross_attention_dim)		,      dtype=jnp.floataa	)

					lowerCamelCase		,		lowerCamelCase    		=  jax.random.split(__SCREAMING_SNAKE_CASE	)
					lowerCamelCase    		=  {"""params""": params_rng, """dropout""": dropout_rng}

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







		def 					__A   ( self	)  -> Optional[Any]:







					'''simple docstring'''







					lowerCamelCase    		=  self.block_out_channels
					lowerCamelCase    		=  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    		=  self.num_attention_heads or self.attention_head_dim

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

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

					lowerCamelCase    		=  self.only_cross_attention
					if isinstance(__SCREAMING_SNAKE_CASE		,      __SCREAMING_SNAKE_CASE	):
								lowerCamelCase    		=  (only_cross_attention,) * len(self.down_block_types	)

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

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

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

								down_blocks.append(__SCREAMING_SNAKE_CASE	)
					lowerCamelCase    		=  down_blocks

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

								lowerCamelCase    		=  i == len(__SCREAMING_SNAKE_CASE	) - 1

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

								up_blocks.append(__SCREAMING_SNAKE_CASE	)
								lowerCamelCase    		=  output_channel
					lowerCamelCase    		=  up_blocks

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







		def __call__( self		,      A		,      A		,      A		,      A=None		,      A=None		,      A = True		,      A = False		,      )  -> Union[FlaxUNetaDConditionOutput, Tuple]:







					'''simple docstring'''







					if not isinstance(__SCREAMING_SNAKE_CASE		,      jnp.ndarray	):
								lowerCamelCase    		=  jnp.array([timesteps]		,      dtype=jnp.intaa	)
					elif isinstance(__SCREAMING_SNAKE_CASE		,      jnp.ndarray	) and len(timesteps.shape	) == 0:
								lowerCamelCase    		=  timesteps.astype(dtype=jnp.floataa	)
								lowerCamelCase    		=  jnp.expand_dims(__SCREAMING_SNAKE_CASE		,      0	)

					lowerCamelCase    		=  self.time_proj(__SCREAMING_SNAKE_CASE	)
					lowerCamelCase    		=  self.time_embedding(__SCREAMING_SNAKE_CASE	)

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

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

					if down_block_additional_residuals is not None:
								lowerCamelCase    		=  ()

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

								lowerCamelCase    		=  new_down_block_res_samples

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

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

					if not return_dict:
								return (sample,)

					return FlaxUNetaDConditionOutput(sample=__SCREAMING_SNAKE_CASE	)







'''simple docstring'''



from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)







def       _UpperCamelCase			(_lowerCamelCase				:  Union[str, Any] ,						_lowerCamelCase				:  List[Any]						)-> Optional[Any]:




     '''simple docstring'''





     __snake_case            =	[]
     for part_id in partition_order:
          __snake_case            =	df.where(f'''SPARK_PARTITION_ID() = {part_id}'''						).collect()
          for row_idx, row in enumerate(_lowerCamelCase						):
               expected_row_ids_and_row_dicts.append((f'''{part_id}_{row_idx}''', row.asDict())						)
     return expected_row_ids_and_row_dicts







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Any:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(1_00						).repartition(1						)
     __snake_case            =	Spark(_lowerCamelCase						)
     # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
     # that each partition can hold 2 rows.
     spark_builder._repartition_df_if_needed(max_shard_size=16						)
     # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
     assert spark_builder.df.rdd.getNumPartitions() == 50







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Tuple:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(10						).repartition(2						)
     __snake_case            =	[1, 0]
     __snake_case            =	_generate_iterable_examples(_lowerCamelCase ,						_lowerCamelCase						)  # Reverse the partitions.
     __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						_lowerCamelCase						)

     for i, (row_id, row_dict) in enumerate(generate_fn()						):
          __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> int:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(10						).repartition(1						)
     __snake_case            =	SparkExamplesIterable(_lowerCamelCase						)
     assert it.n_shards == 1
     for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
          assert row_id == f'''0_{i}'''
          assert row_dict == {"id": i}







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Union[str, Any]:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(30						).repartition(3						)
     # Mock the generator so that shuffle reverses the partition indices.
     with patch('''numpy.random.Generator'''						) as generator_mock:
          __snake_case            =	lambda _lowerCamelCase						: x.reverse()
          __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						[2, 1, 0]						)

          __snake_case            =	SparkExamplesIterable(_lowerCamelCase						).shuffle_data_sources(_lowerCamelCase						)
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
               __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts[i]
               assert row_id == expected_row_id
               assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Tuple:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(20						).repartition(4						)

     # Partitions 0 and 2
     __snake_case            =	SparkExamplesIterable(_lowerCamelCase						).shard_data_sources(worker_id=0 ,						num_workers=2						)
     assert shard_it_a.n_shards == 2
     __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						[0, 2]						)
     for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
          __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict

     # Partitions 1 and 3
     __snake_case            =	SparkExamplesIterable(_lowerCamelCase						).shard_data_sources(worker_id=1 ,						num_workers=2						)
     assert shard_it_a.n_shards == 2
     __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						[1, 3]						)
     for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
          __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Optional[int]:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(1_00						).repartition(1						)
     __snake_case            =	Spark(_lowerCamelCase						)
     # Choose a small max_shard_size for maximum partitioning.
     spark_builder._repartition_df_if_needed(max_shard_size=1						)
     # The new number of partitions should not be greater than the number of rows.
     assert spark_builder.df.rdd.getNumPartitions() == 1_00






