


'''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     a   (    __a				,      __a     )						->			Optional[int]:



					'''simple docstring'''






					UpperCamelCase__     :str							      =				[]
					for part_id in partition_order:
										UpperCamelCase__     :Optional[int]							      =				df.where(f'''SPARK_PARTITION_ID() = {part_id}'''     ).collect()
										for row_idx, row in enumerate(__a     ):
															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     a   (    )						->			Union[str, Any]:



					'''simple docstring'''






					UpperCamelCase__     :Optional[int]							      =				pyspark.sql.SparkSession.builder.master('''local[*]'''     ).appName('''pyspark'''     ).getOrCreate()
					UpperCamelCase__     :List[Any]							      =				spark.range(100     ).repartition(1     )
					UpperCamelCase__     :int							      =				Spark(__a     )
					# 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     a   (    )						->			str:



					'''simple docstring'''






					UpperCamelCase__     :Any							      =				pyspark.sql.SparkSession.builder.master('''local[*]'''     ).appName('''pyspark'''     ).getOrCreate()
					UpperCamelCase__     :Union[str, Any]							      =				spark.range(10     ).repartition(2     )
					UpperCamelCase__     :str							      =				[1, 0]
					UpperCamelCase__     :Optional[int]							      =				_generate_iterable_examples(__a				,      __a     )  # Reverse the partitions.
					UpperCamelCase__     :str							      =				_get_expected_row_ids_and_row_dicts_for_partition_order(__a				,      __a     )

					for i, (row_id, row_dict) in enumerate(generate_fn()     ):
										UpperCamelCase__      ,       UpperCamelCase__     :List[Any]							      =				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     a   (    )						->			Dict:



					'''simple docstring'''






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


@require_not_windows
@require_dill_gt_0_3_2
def     a   (    )						->			List[Any]:



					'''simple docstring'''






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

										UpperCamelCase__     :Dict							      =				SparkExamplesIterable(__a     ).shuffle_data_sources(__a     )
										assert shuffled_it.n_shards == 3
										for i, (row_id, row_dict) in enumerate(__a     ):
															UpperCamelCase__      ,       UpperCamelCase__     :Optional[Any]							      =				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     a   (    )						->			Dict:



					'''simple docstring'''






					UpperCamelCase__     :str							      =				pyspark.sql.SparkSession.builder.master('''local[*]'''     ).appName('''pyspark'''     ).getOrCreate()
					UpperCamelCase__     :Union[str, Any]							      =				spark.range(20     ).repartition(4     )

					# Partitions 0 and 2
					UpperCamelCase__     :Dict							      =				SparkExamplesIterable(__a     ).shard_data_sources(worker_id=0				,      num_workers=2     )
					assert shard_it_a.n_shards == 2
					UpperCamelCase__     :Tuple							      =				_get_expected_row_ids_and_row_dicts_for_partition_order(__a				,      [0, 2]     )
					for i, (row_id, row_dict) in enumerate(__a     ):
										UpperCamelCase__      ,       UpperCamelCase__     :Optional[Any]							      =				expected_row_ids_and_row_dicts_a[i]
										assert row_id == expected_row_id
										assert row_dict == expected_row_dict

					# Partitions 1 and 3
					UpperCamelCase__     :Optional[int]							      =				SparkExamplesIterable(__a     ).shard_data_sources(worker_id=1				,      num_workers=2     )
					assert shard_it_a.n_shards == 2
					UpperCamelCase__     :List[Any]							      =				_get_expected_row_ids_and_row_dicts_for_partition_order(__a				,      [1, 3]     )
					for i, (row_id, row_dict) in enumerate(__a     ):
										UpperCamelCase__      ,       UpperCamelCase__     :List[Any]							      =				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     a   (    )						->			Optional[int]:



					'''simple docstring'''






					UpperCamelCase__     :List[str]							      =				pyspark.sql.SparkSession.builder.master('''local[*]'''     ).appName('''pyspark'''     ).getOrCreate()
					UpperCamelCase__     :Optional[int]							      =				spark.range(100     ).repartition(1     )
					UpperCamelCase__     :str							      =				Spark(__a     )
					# 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() == 100



'''simple docstring'''
import math
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    is_torch_available,
    is_torch_tensor,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_torch_available():
		import torch

if is_vision_available():
		import PIL


__snake_case							   =  logging.get_logger(__name__)


def     a   (    __a				,      __a				,      __a				,      __a     )						->			Tuple[int, int]:



					'''simple docstring'''






					def constraint_to_multiple_of(__a				,      __a				,      __a=0				,      __a=None     ):
										UpperCamelCase__     :Union[str, Any]							      =				round(val / multiple     ) * multiple

										if max_val is not None and x > max_val:
															UpperCamelCase__     :List[str]							      =				math.floor(val / multiple     ) * multiple

										if x < min_val:
															UpperCamelCase__     :Union[str, Any]							      =				math.ceil(val / multiple     ) * multiple

										return x

					UpperCamelCase__     :str							      =				(output_size, output_size) if isinstance(__a				,      __a     ) else output_size

					UpperCamelCase__      ,       UpperCamelCase__     :int							      =				get_image_size(__a     )
					UpperCamelCase__      ,       UpperCamelCase__     :Optional[int]							      =				output_size

					# determine new height and width
					UpperCamelCase__     :List[str]							      =				output_height / input_height
					UpperCamelCase__     :str							      =				output_width / input_width

					if keep_aspect_ratio:
										# scale as little as possible
										if abs(1 - scale_width     ) < abs(1 - scale_height     ):
															# fit width
															UpperCamelCase__     :str							      =				scale_width
										else:
															# fit height
															UpperCamelCase__     :int							      =				scale_height

					UpperCamelCase__     :Optional[int]							      =				constraint_to_multiple_of(scale_height * input_height				,      multiple=__a     )
					UpperCamelCase__     :List[Any]							      =				constraint_to_multiple_of(scale_width * input_width				,      multiple=__a     )

					return (new_height, new_width)







class      lowercase			(							A__      ):

		"""simple docstring"""

		_a									=       ['pixel_values']




		def __init__(	self					,     UpperCamelCase_ = True					,     UpperCamelCase_ = None					,     UpperCamelCase_ = PILImageResampling.BILINEAR					,     UpperCamelCase_ = False					,     UpperCamelCase_ = 1					,     UpperCamelCase_ = True					,     UpperCamelCase_ = 1 / 255					,     UpperCamelCase_ = True					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     **UpperCamelCase_					,     ):

							'''simple docstring'''


							super().__init__(**UpperCamelCase_							)
							UpperCamelCase__     :List[Any]							      =				size if size is not None else {'''height''': 384, '''width''': 384}
							UpperCamelCase__     :Dict							      =				get_size_dict(UpperCamelCase_							)
							UpperCamelCase__     :Dict							      =				do_resize
							UpperCamelCase__     :Union[str, Any]							      =				size
							UpperCamelCase__     :List[Any]							      =				keep_aspect_ratio
							UpperCamelCase__     :Optional[int]							      =				ensure_multiple_of
							UpperCamelCase__     :Union[str, Any]							      =				resample
							UpperCamelCase__     :str							      =				do_rescale
							UpperCamelCase__     :Union[str, Any]							      =				rescale_factor
							UpperCamelCase__     :List[str]							      =				do_normalize
							UpperCamelCase__     :List[str]							      =				image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
							UpperCamelCase__     :Dict							      =				image_std if image_std is not None else IMAGENET_STANDARD_STD




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_ = False					,     UpperCamelCase_ = 1					,     UpperCamelCase_ = PILImageResampling.BICUBIC					,     UpperCamelCase_ = None					,     **UpperCamelCase_					,     ):

							'''simple docstring'''


							UpperCamelCase__     :int							      =				get_size_dict(UpperCamelCase_							)
							if "height" not in size or "width" not in size:
												raise ValueError(F'''The size dictionary must contain the keys \'height\' and \'width\'. Got {size.keys()}'''							)
							UpperCamelCase__     :int							      =				get_resize_output_image_size(
							    UpperCamelCase_					,     output_size=(size['''height'''], size['''width'''])					,     keep_aspect_ratio=UpperCamelCase_					,     multiple=UpperCamelCase_					,     )
							return resize(UpperCamelCase_					,     size=UpperCamelCase_					,     resample=UpperCamelCase_					,     data_format=UpperCamelCase_					,     **UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_ = None					,     **UpperCamelCase_					,     ):

							'''simple docstring'''


							return rescale(UpperCamelCase_					,     scale=UpperCamelCase_					,     data_format=UpperCamelCase_					,     **UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_ = None					,     **UpperCamelCase_					,     ):

							'''simple docstring'''


							return normalize(UpperCamelCase_					,     mean=UpperCamelCase_					,     std=UpperCamelCase_					,     data_format=UpperCamelCase_					,     **UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = ChannelDimension.FIRST					,     **UpperCamelCase_					,     ):

							'''simple docstring'''


							UpperCamelCase__     :str							      =				do_resize if do_resize is not None else self.do_resize
							UpperCamelCase__     :Any							      =				size if size is not None else self.size
							UpperCamelCase__     :List[str]							      =				get_size_dict(UpperCamelCase_							)
							UpperCamelCase__     :Any							      =				keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
							UpperCamelCase__     :str							      =				ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
							UpperCamelCase__     :Union[str, Any]							      =				resample if resample is not None else self.resample
							UpperCamelCase__     :str							      =				do_rescale if do_rescale is not None else self.do_rescale
							UpperCamelCase__     :int							      =				rescale_factor if rescale_factor is not None else self.rescale_factor
							UpperCamelCase__     :Optional[Any]							      =				do_normalize if do_normalize is not None else self.do_normalize
							UpperCamelCase__     :Optional[int]							      =				image_mean if image_mean is not None else self.image_mean
							UpperCamelCase__     :List[str]							      =				image_std if image_std is not None else self.image_std

							UpperCamelCase__     :List[str]							      =				make_list_of_images(UpperCamelCase_							)

							if not valid_images(UpperCamelCase_							):
												raise ValueError(
												    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
												    '''torch.Tensor, tf.Tensor or jax.ndarray.'''							)

							if do_resize and size is None or resample is None:
												raise ValueError('''Size and resample must be specified if do_resize is True.'''							)

							if do_rescale and rescale_factor is None:
												raise ValueError('''Rescale factor must be specified if do_rescale is True.'''							)

							if do_normalize and (image_mean is None or image_std is None):
												raise ValueError('''Image mean and std must be specified if do_normalize is True.'''							)

							# All transformations expect numpy arrays.
							UpperCamelCase__     :str							      =				[to_numpy_array(UpperCamelCase_							) for image in images]

							if do_resize:
												UpperCamelCase__     :str							      =				[self.resize(image=UpperCamelCase_					,     size=UpperCamelCase_					,     resample=UpperCamelCase_							) for image in images]

							if do_rescale:
												UpperCamelCase__     :int							      =				[self.rescale(image=UpperCamelCase_					,     scale=UpperCamelCase_							) for image in images]

							if do_normalize:
												UpperCamelCase__     :str							      =				[self.normalize(image=UpperCamelCase_					,     mean=UpperCamelCase_					,     std=UpperCamelCase_							) for image in images]

							UpperCamelCase__     :Optional[int]							      =				[to_channel_dimension_format(UpperCamelCase_					,     UpperCamelCase_							) for image in images]

							UpperCamelCase__     :List[str]							      =				{'''pixel_values''': images}
							return BatchFeature(data=UpperCamelCase_					,     tensor_type=UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_ = None							):

							'''simple docstring'''


							UpperCamelCase__     :List[Any]							      =				outputs.logits

							# Resize logits and compute semantic segmentation maps
							if target_sizes is not None:
												if len(UpperCamelCase_							) != len(UpperCamelCase_							):
																	raise ValueError(
																	    '''Make sure that you pass in as many target sizes as the batch dimension of the logits'''							)

												if is_torch_tensor(UpperCamelCase_							):
																	UpperCamelCase__     :Union[str, Any]							      =				target_sizes.numpy()

												UpperCamelCase__     :int							      =				[]

												for idx in range(len(UpperCamelCase_							)							):
																	UpperCamelCase__     :str							      =				torch.nn.functional.interpolate(
																	    logits[idx].unsqueeze(dim=0							)					,     size=target_sizes[idx]					,     mode='''bilinear'''					,     align_corners=UpperCamelCase_							)
																	UpperCamelCase__     :List[Any]							      =				resized_logits[0].argmax(dim=0							)
																	semantic_segmentation.append(UpperCamelCase_							)
							else:
												UpperCamelCase__     :Any							      =				logits.argmax(dim=1							)
												UpperCamelCase__     :Optional[Any]							      =				[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]							)]

							return semantic_segmentation