





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_    (    _UpperCAmelCase							:Any  ,						_UpperCAmelCase							:List[str]				)	->				Union[str, Any]:
     '''simple docstring'''


     A_          =       []
     for part_id in partition_order:
          A_          =       df.where(f'SPARK_PARTITION_ID() = {part_id}'				).collect()
          for row_idx, row in enumerate(_UpperCAmelCase				):
               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_    (    )	->				List[Any]:
     '''simple docstring'''


     A_          =       pyspark.sql.SparkSession.builder.master('''local[*]'''				).appName('''pyspark'''				).getOrCreate()
     A_          =       spark.range(100				).repartition(1				)
     A_          =       Spark(_UpperCAmelCase				)
     # 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_    (    )	->				Any:
     '''simple docstring'''


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

     for i, (row_id, row_dict) in enumerate(generate_fn()				):
          A_   ,    A_          =       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_    (    )	->				Union[str, Any]:
     '''simple docstring'''


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

@require_not_windows
@require_dill_gt_0_3_2
def      UpperCAmelCase_    (    )	->				Tuple:
     '''simple docstring'''


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

          A_          =       SparkExamplesIterable(_UpperCAmelCase				).shuffle_data_sources(_UpperCAmelCase				)
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(_UpperCAmelCase				):
               A_   ,    A_          =       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_    (    )	->				Any:
     '''simple docstring'''


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

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

     # Partitions 1 and 3
     A_          =       SparkExamplesIterable(_UpperCAmelCase				).shard_data_sources(worker_id=1  ,						num_workers=2				)
     assert shard_it_a.n_shards == 2
     A_          =       _get_expected_row_ids_and_row_dicts_for_partition_order(_UpperCAmelCase  ,						[1, 3]				)
     for i, (row_id, row_dict) in enumerate(_UpperCAmelCase				):
          A_   ,    A_          =       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_    (    )	->				Union[str, Any]:
     '''simple docstring'''


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










import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest





class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE ):
  __lowerCAmelCase   :						int							=   (KDPMaDiscreteScheduler,)
  __lowerCAmelCase   :						Union[str, Any]							=   10


  def 						__UpperCAmelCase     (							self ,**__snake_case ):



       """simple docstring"""





       A_          =       {
           '''num_train_timesteps''': 1_1_0_0,
           '''beta_start''': 0.0001,
           '''beta_end''': 0.02,
           '''beta_schedule''': '''linear''',
       }

       config.update(**__snake_case )
       return config


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for timesteps in [1_0, 5_0, 1_0_0, 1_0_0_0]:
            self.check_over_configs(num_train_timesteps=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001] ,[0.0002, 0.002, 0.02] ):
            self.check_over_configs(beta_start=__snake_case ,beta_end=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for schedule in ["linear", "scaled_linear"]:
            self.check_over_configs(beta_schedule=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for prediction_type in ["epsilon", "v_prediction"]:
            self.check_over_configs(prediction_type=__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.scheduler_classes[0]
       A_          =       self.get_scheduler_config(prediction_type='''v_prediction''' )
       A_          =       scheduler_class(**__snake_case )

       scheduler.set_timesteps(self.num_inference_steps )

       A_          =       self.dummy_model()
       A_          =       self.dummy_sample_deter * scheduler.init_noise_sigma
       A_          =       sample.to(__snake_case )

       for i, t in enumerate(scheduler.timesteps ):
            A_          =       scheduler.scale_model_input(__snake_case ,__snake_case )

            A_          =       model(__snake_case ,__snake_case )

            A_          =       scheduler.step(__snake_case ,__snake_case ,__snake_case )
            A_          =       output.prev_sample

       A_          =       torch.sum(torch.abs(__snake_case ) )
       A_          =       torch.mean(torch.abs(__snake_case ) )

       if torch_device in ["cpu", "mps"]:
            assert abs(result_sum.item() - 4.69_34E-07 ) < 1E-2
            assert abs(result_mean.item() - 6.11_12E-10 ) < 1E-3
       else:
            # CUDA
            assert abs(result_sum.item() - 4.6_93_42_86_50_17_09_72E-07 ) < 1E-2
            assert abs(result_mean.item() - 0.0002 ) < 1E-3


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       if torch_device == "mps":
            return
       A_          =       self.scheduler_classes[0]
       A_          =       self.get_scheduler_config()
       A_          =       scheduler_class(**__snake_case )

       scheduler.set_timesteps(self.num_inference_steps )

       A_          =       self.dummy_model()
       A_          =       self.dummy_sample_deter * scheduler.init_noise_sigma
       A_          =       sample.to(__snake_case )

       for i, t in enumerate(scheduler.timesteps ):
            A_          =       scheduler.scale_model_input(__snake_case ,__snake_case )

            A_          =       model(__snake_case ,__snake_case )

            A_          =       scheduler.step(__snake_case ,__snake_case ,__snake_case )
            A_          =       output.prev_sample

       A_          =       torch.sum(torch.abs(__snake_case ) )
       A_          =       torch.mean(torch.abs(__snake_case ) )

       if torch_device in ["cpu", "mps"]:
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3
       else:
            # CUDA
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       if torch_device == "mps":
            return
       A_          =       self.scheduler_classes[0]
       A_          =       self.get_scheduler_config()
       A_          =       scheduler_class(**__snake_case )

       scheduler.set_timesteps(self.num_inference_steps ,device=__snake_case )

       A_          =       self.dummy_model()
       A_          =       self.dummy_sample_deter.to(__snake_case ) * scheduler.init_noise_sigma

       for t in scheduler.timesteps:
            A_          =       scheduler.scale_model_input(__snake_case ,__snake_case )

            A_          =       model(__snake_case ,__snake_case )

            A_          =       scheduler.step(__snake_case ,__snake_case ,__snake_case )
            A_          =       output.prev_sample

       A_          =       torch.sum(torch.abs(__snake_case ) )
       A_          =       torch.mean(torch.abs(__snake_case ) )

       if str(__snake_case ).startswith('''cpu''' ):
            # The following sum varies between 148 and 156 on mps. Why?
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3
       else:
            # CUDA
            assert abs(result_sum.item() - 20.4125 ) < 1E-2
            assert abs(result_mean.item() - 0.0266 ) < 1E-3



