


'''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    lowerCamelCase__   ( A     :    Dict   ,     A     :    Any  ):




      '''simple docstring'''






      UpperCAmelCase         =			[]
      for part_id in partition_order:
            UpperCAmelCase         =			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    lowerCamelCase__   ( ):




      '''simple docstring'''






      UpperCAmelCase         =			pyspark.sql.SparkSession.builder.master('''local[*]'''  ).appName('''pyspark'''  ).getOrCreate()
      UpperCAmelCase         =			spark.range(1_00  ).repartition(1  )
      UpperCAmelCase         =			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    lowerCamelCase__   ( ):




      '''simple docstring'''






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

      for i, (row_id, row_dict) in enumerate(generate_fn()  ):
            UpperCAmelCase       ,       UpperCAmelCase         =			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    lowerCamelCase__   ( ):




      '''simple docstring'''






      UpperCAmelCase         =			pyspark.sql.SparkSession.builder.master('''local[*]'''  ).appName('''pyspark'''  ).getOrCreate()
      UpperCAmelCase         =			spark.range(10  ).repartition(1  )
      UpperCAmelCase         =			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    lowerCamelCase__   ( ):




      '''simple docstring'''






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

            UpperCAmelCase         =			SparkExamplesIterable(A  ).shuffle_data_sources(A  )
            assert shuffled_it.n_shards == 3
            for i, (row_id, row_dict) in enumerate(A  ):
                  UpperCAmelCase       ,       UpperCAmelCase         =			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    lowerCamelCase__   ( ):




      '''simple docstring'''






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

      # Partitions 0 and 2
      UpperCAmelCase         =			SparkExamplesIterable(A  ).shard_data_sources(worker_id=0   ,     num_workers=2  )
      assert shard_it_a.n_shards == 2
      UpperCAmelCase         =			_get_expected_row_ids_and_row_dicts_for_partition_order(A   ,     [0, 2]  )
      for i, (row_id, row_dict) in enumerate(A  ):
            UpperCAmelCase       ,       UpperCAmelCase         =			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         =			SparkExamplesIterable(A  ).shard_data_sources(worker_id=1   ,     num_workers=2  )
      assert shard_it_a.n_shards == 2
      UpperCAmelCase         =			_get_expected_row_ids_and_row_dicts_for_partition_order(A   ,     [1, 3]  )
      for i, (row_id, row_dict) in enumerate(A  ):
            UpperCAmelCase       ,       UpperCAmelCase         =			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    lowerCamelCase__   ( ):




      '''simple docstring'''






      UpperCAmelCase         =			pyspark.sql.SparkSession.builder.master('''local[*]'''  ).appName('''pyspark'''  ).getOrCreate()
      UpperCAmelCase         =			spark.range(1_00  ).repartition(1  )
      UpperCAmelCase         =			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() == 1_00









'''simple docstring'''



import inspect
import unittest

from transformers import RegNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
	import jax
	import jax.numpy as jnp

	from transformers.models.regnet.modeling_flax_regnet import FlaxRegNetForImageClassification, FlaxRegNetModel

if is_vision_available():
	from PIL import Image

	from transformers import AutoImageProcessor






class  UpperCamelCase__(     unittest.TestCase ):




				def __init__(	self    :				Dict   ,					lowerCAmelCase    :				List[str]   ,					lowerCAmelCase    :				List[Any]=3   ,					lowerCAmelCase    :				Tuple=32   ,					lowerCAmelCase    :				Optional[int]=3   ,					lowerCAmelCase    :				Optional[Any]=10   ,					lowerCAmelCase    :				Optional[Any]=[10, 20, 30, 40]   ,					lowerCAmelCase    :				Optional[int]=[1, 1, 2, 1]   ,					lowerCAmelCase    :				Union[str, Any]=True   ,					lowerCAmelCase    :				Tuple=True   ,					lowerCAmelCase    :				List[Any]="relu"   ,					lowerCAmelCase    :				List[Any]=3   ,					lowerCAmelCase    :				Union[str, Any]=None   ,					)->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase         =			parent
										UpperCAmelCase         =			batch_size
										UpperCAmelCase         =			image_size
										UpperCAmelCase         =			num_channels
										UpperCAmelCase         =			embeddings_size
										UpperCAmelCase         =			hidden_sizes
										UpperCAmelCase         =			depths
										UpperCAmelCase         =			is_training
										UpperCAmelCase         =			use_labels
										UpperCAmelCase         =			hidden_act
										UpperCAmelCase         =			num_labels
										UpperCAmelCase         =			scope
										UpperCAmelCase         =			len(lowerCAmelCase   )




				def   a__(	self    :				Optional[Any]   )->  Tuple:

										"""simple docstring"""




										UpperCAmelCase         =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

										UpperCAmelCase         =			self.get_config()

										return config, pixel_values




				def   a__(	self    :				str   )->  Optional[Any]:

										"""simple docstring"""




										return RegNetConfig(
										    num_channels=self.num_channels   ,					embeddings_size=self.embeddings_size   ,					hidden_sizes=self.hidden_sizes   ,					depths=self.depths   ,					hidden_act=self.hidden_act   ,					num_labels=self.num_labels   ,					image_size=self.image_size   ,					)




				def   a__(	self    :				Union[str, Any]   ,					lowerCAmelCase    :				List[Any]   ,					lowerCAmelCase    :				Any   )->  Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase         =			FlaxRegNetModel(config=lowerCAmelCase   )
										UpperCAmelCase         =			model(lowerCAmelCase   )

										# Output shape (b, c, h, w)
										self.parent.assertEqual(
										    result.last_hidden_state.shape   ,					(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)   ,					)




				def   a__(	self    :				List[Any]   ,					lowerCAmelCase    :				List[str]   ,					lowerCAmelCase    :				Optional[Any]   )->  List[str]:

										"""simple docstring"""




										UpperCAmelCase         =			self.num_labels
										UpperCAmelCase         =			FlaxRegNetForImageClassification(config=lowerCAmelCase   )
										UpperCAmelCase         =			model(lowerCAmelCase   )
										self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.num_labels)   )





				def   a__(	self    :				List[str]   )->  Any:

										"""simple docstring"""




										UpperCAmelCase         =			self.prepare_config_and_inputs()
										UpperCAmelCase       ,       UpperCAmelCase         =			config_and_inputs
										UpperCAmelCase         =			{'''pixel_values''': pixel_values}
										return config, inputs_dict







@require_flax
class  UpperCamelCase__(     lowerCAmelCase		,       unittest.TestCase ):
				__magic_name__       :   Optional[int]							= (FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

				__magic_name__       :   Optional[int]							= False
				__magic_name__       :   List[str]							= False
				__magic_name__       :   Dict							= False




				def   a__(	self    :				Union[str, Any]   )->  None:

										"""simple docstring"""




										UpperCAmelCase         =			FlaxRegNetModelTester(self   )
										UpperCAmelCase         =			ConfigTester(self   ,					config_class=lowerCAmelCase   ,					has_text_modality=lowerCAmelCase   )




				def   a__(	self    :				List[str]   )->  List[str]:

										"""simple docstring"""




										self.create_and_test_config_common_properties()
										self.config_tester.create_and_test_config_to_json_string()
										self.config_tester.create_and_test_config_to_json_file()
										self.config_tester.create_and_test_config_from_and_save_pretrained()
										self.config_tester.create_and_test_config_with_num_labels()
										self.config_tester.check_config_can_be_init_without_params()
										self.config_tester.check_config_arguments_init()




				def   a__(	self    :				Tuple   )->  Tuple:

										"""simple docstring"""




										return




				def   a__(	self    :				Optional[Any]   )->  Tuple:

										"""simple docstring"""




										UpperCAmelCase         =			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowerCAmelCase   )




				def   a__(	self    :				Any   )->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase         =			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase   )




				@unittest.skip(reason='''RegNet does not use inputs_embeds'''   )
				def   a__(	self    :				str   )->  Optional[Any]:

										"""simple docstring"""




										pass




				@unittest.skip(reason='''RegNet does not support input and output embeddings'''   )
				def   a__(	self    :				Any   )->  List[str]:

										"""simple docstring"""




										pass




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

										"""simple docstring"""




										UpperCAmelCase       ,       UpperCAmelCase         =			self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																UpperCAmelCase         =			model_class(lowerCAmelCase   )
																UpperCAmelCase         =			inspect.signature(model.__call__   )
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																UpperCAmelCase         =			[*signature.parameters.keys()]

																UpperCAmelCase         =			['''pixel_values''']
																self.assertListEqual(arg_names[:1]   ,					lowerCAmelCase   )




				def   a__(	self    :				Tuple   )->  int:

										"""simple docstring"""




										def check_hidden_states_output(lowerCAmelCase    :				Optional[Any]   ,					lowerCAmelCase    :				Union[str, Any]   ,					lowerCAmelCase    :				int   ):
																UpperCAmelCase         =			model_class(lowerCAmelCase   )

																UpperCAmelCase         =			model(**self._prepare_for_class(lowerCAmelCase   ,					lowerCAmelCase   )   )

																UpperCAmelCase         =			outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																UpperCAmelCase         =			self.model_tester.num_stages
																self.assertEqual(len(lowerCAmelCase   )   ,					expected_num_stages + 1   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																UpperCAmelCase         =			True
																check_hidden_states_output(lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   )

																# check that output_hidden_states also work using config
																del inputs_dict["output_hidden_states"]
																UpperCAmelCase         =			True

																check_hidden_states_output(lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   )





				def   a__(	self    :				Union[str, Any]   )->  List[str]:

										"""simple docstring"""




										UpperCAmelCase       ,       UpperCAmelCase         =			self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__   ):
																						UpperCAmelCase         =			self._prepare_for_class(lowerCAmelCase   ,					lowerCAmelCase   )
																						UpperCAmelCase         =			model_class(lowerCAmelCase   )

																						@jax.jit
																						def model_jitted(lowerCAmelCase    :				Tuple   ,					**lowerCAmelCase    :				Tuple   ):
																												return model(pixel_values=lowerCAmelCase   ,					**lowerCAmelCase   )

																						with self.subTest('''JIT Enabled'''   ):
																												UpperCAmelCase         =			model_jitted(**lowerCAmelCase   ).to_tuple()

																						with self.subTest('''JIT Disabled'''   ):
																												with jax.disable_jit():
																																		UpperCAmelCase         =			model_jitted(**lowerCAmelCase   ).to_tuple()

																						self.assertEqual(len(lowerCAmelCase   )   ,					len(lowerCAmelCase   )   )
																						for jitted_output, output in zip(lowerCAmelCase   ,					lowerCAmelCase   ):
																												self.assertEqual(jitted_output.shape   ,					output.shape   )







def    lowerCamelCase__   ( ):




						'''simple docstring'''






						UpperCAmelCase         =			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''  )
						return image






@require_flax
class  UpperCamelCase__(     unittest.TestCase ):




				@cached_property
				def   a__(	self    :				Dict   )->  int:

										"""simple docstring"""




										return AutoImageProcessor.from_pretrained('''facebook/regnet-y-040'''   ) if is_vision_available() else None





				@slow
				def   a__(	self    :				Union[str, Any]   )->  Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase         =			FlaxRegNetForImageClassification.from_pretrained('''facebook/regnet-y-040'''   )

										UpperCAmelCase         =			self.default_image_processor
										UpperCAmelCase         =			prepare_img()
										UpperCAmelCase         =			image_processor(images=lowerCAmelCase   ,					return_tensors='''np'''   )

										UpperCAmelCase         =			model(**lowerCAmelCase   )

										# verify the logits
										UpperCAmelCase         =			(1, 1000)
										self.assertEqual(outputs.logits.shape   ,					lowerCAmelCase   )

										UpperCAmelCase         =			jnp.array([-0.4180, -1.5051, -3.4836]   )

										self.assertTrue(jnp.allclose(outputs.logits[0, :3]   ,					lowerCAmelCase   ,					atol=1E-4   )   )





