import inspect
import unittest
import warnings

from transformers import DeiTConfig
from transformers.models.auto import get_values
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import (
        MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
        MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        MODEL_MAPPING,
        DeiTForImageClassification,
        DeiTForImageClassificationWithTeacher,
        DeiTForMaskedImageModeling,
        DeiTModel,
    )
    from transformers.models.deit.modeling_deit import DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import DeiTImageProcessor


class        a				:


    def __init__(					self	:  List[Any]   ,				__lowerCAmelCase	:  Optional[Any]   ,				__lowerCAmelCase	:  List[Any]=13   ,				__lowerCAmelCase	:  List[Any]=30   ,				__lowerCAmelCase	:  Union[str, Any]=2   ,				__lowerCAmelCase	:  Union[str, Any]=3   ,				__lowerCAmelCase	:  int=True   ,				__lowerCAmelCase	:  Union[str, Any]=True   ,				__lowerCAmelCase	:  Dict=32   ,				__lowerCAmelCase	:  Optional[int]=5   ,				__lowerCAmelCase	:  Optional[int]=4   ,				__lowerCAmelCase	:  Union[str, Any]=37   ,				__lowerCAmelCase	:  Union[str, Any]="gelu"   ,				__lowerCAmelCase	:  Union[str, Any]=0.1   ,				__lowerCAmelCase	:  Optional[int]=0.1   ,				__lowerCAmelCase	:  List[Any]=10   ,				__lowerCAmelCase	:  Optional[Any]=0.02   ,				__lowerCAmelCase	:  Optional[Any]=3   ,				__lowerCAmelCase	:  str=None   ,				__lowerCAmelCase	:  List[Any]=2   ,				):
      _UpperCAmelCase						 =      parent
      _UpperCAmelCase						 =      batch_size
      _UpperCAmelCase						 =      image_size
      _UpperCAmelCase						 =      patch_size
      _UpperCAmelCase						 =      num_channels
      _UpperCAmelCase						 =      is_training
      _UpperCAmelCase						 =      use_labels
      _UpperCAmelCase						 =      hidden_size
      _UpperCAmelCase						 =      num_hidden_layers
      _UpperCAmelCase						 =      num_attention_heads
      _UpperCAmelCase						 =      intermediate_size
      _UpperCAmelCase						 =      hidden_act
      _UpperCAmelCase						 =      hidden_dropout_prob
      _UpperCAmelCase						 =      attention_probs_dropout_prob
      _UpperCAmelCase						 =      type_sequence_label_size
      _UpperCAmelCase						 =      initializer_range
      _UpperCAmelCase						 =      scope
      _UpperCAmelCase						 =      encoder_stride

      # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
      _UpperCAmelCase						 =      (image_size // patch_size) ** 2
      _UpperCAmelCase						 =      num_patches + 2


    def      lowerCAmelCase_  (					self	:  List[Any]					):
      _UpperCAmelCase						 =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

      _UpperCAmelCase						 =      None
      if self.use_labels:
        _UpperCAmelCase						 =      ids_tensor([self.batch_size]   ,				self.type_sequence_label_size					)

      _UpperCAmelCase						 =      self.get_config()

      return config, pixel_values, labels


    def      lowerCAmelCase_  (					self	:  List[Any]					):
      return DeiTConfig(
          image_size=self.image_size   ,				patch_size=self.patch_size   ,				num_channels=self.num_channels   ,				hidden_size=self.hidden_size   ,				num_hidden_layers=self.num_hidden_layers   ,				num_attention_heads=self.num_attention_heads   ,				intermediate_size=self.intermediate_size   ,				hidden_act=self.hidden_act   ,				hidden_dropout_prob=self.hidden_dropout_prob   ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,				is_decoder=__lowerCAmelCase   ,				initializer_range=self.initializer_range   ,				encoder_stride=self.encoder_stride   ,				)


    def      lowerCAmelCase_  (					self	:  Tuple   ,				__lowerCAmelCase	:  List[Any]   ,				__lowerCAmelCase	:  List[str]   ,				__lowerCAmelCase	:  Dict					):
      _UpperCAmelCase						 =      DeiTModel(config=__lowerCAmelCase					)
      model.to(__lowerCAmelCase					)
      model.eval()
      _UpperCAmelCase						 =      model(__lowerCAmelCase					)
      self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.hidden_size)					)


    def      lowerCAmelCase_  (					self	:  Dict   ,				__lowerCAmelCase	:  Any   ,				__lowerCAmelCase	:  Tuple   ,				__lowerCAmelCase	:  Optional[Any]					):
      _UpperCAmelCase						 =      DeiTForMaskedImageModeling(config=__lowerCAmelCase					)
      model.to(__lowerCAmelCase					)
      model.eval()
      _UpperCAmelCase						 =      model(__lowerCAmelCase					)
      self.parent.assertEqual(
          result.reconstruction.shape   ,				(self.batch_size, self.num_channels, self.image_size, self.image_size)					)

      # test greyscale images
      _UpperCAmelCase						 =      1
      _UpperCAmelCase						 =      DeiTForMaskedImageModeling(__lowerCAmelCase					)
      model.to(__lowerCAmelCase					)
      model.eval()

      _UpperCAmelCase						 =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
      _UpperCAmelCase						 =      model(__lowerCAmelCase					)
      self.parent.assertEqual(result.reconstruction.shape   ,				(self.batch_size, 1, self.image_size, self.image_size)					)


    def      lowerCAmelCase_  (					self	:  int   ,				__lowerCAmelCase	:  Optional[Any]   ,				__lowerCAmelCase	:  str   ,				__lowerCAmelCase	:  List[Any]					):
      _UpperCAmelCase						 =      self.type_sequence_label_size
      _UpperCAmelCase						 =      DeiTForImageClassification(__lowerCAmelCase					)
      model.to(__lowerCAmelCase					)
      model.eval()
      _UpperCAmelCase						 =      model(__lowerCAmelCase   ,				labels=__lowerCAmelCase					)
      self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.type_sequence_label_size)					)

      # test greyscale images
      _UpperCAmelCase						 =      1
      _UpperCAmelCase						 =      DeiTForImageClassification(__lowerCAmelCase					)
      model.to(__lowerCAmelCase					)
      model.eval()

      _UpperCAmelCase						 =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
      _UpperCAmelCase						 =      model(__lowerCAmelCase   ,				labels=__lowerCAmelCase					)
      self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.type_sequence_label_size)					)


    def      lowerCAmelCase_  (					self	:  Dict					):
      _UpperCAmelCase						 =      self.prepare_config_and_inputs()
      (
          _UpperCAmelCase		
      )						 =      config_and_inputs
      _UpperCAmelCase						 =      {"""pixel_values""": pixel_values}
      return config, inputs_dict



@require_torch
class        a				(		UpperCAmelCase__	,							UpperCAmelCase__	,							unittest.TestCase   ):
    _snake_case     :  Dict     							=							(
        (
            DeiTModel,
            DeiTForImageClassification,
            DeiTForImageClassificationWithTeacher,
            DeiTForMaskedImageModeling,
        )
        if is_torch_available()
        else ()
    )
    _snake_case     :  List[str]     							=							(
        {
            'feature-extraction': DeiTModel,
            'image-classification': (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
        }
        if is_torch_available()
        else {}
    )

    _snake_case     :  List[Any]     							=							False
    _snake_case     :  List[str]     							=							False
    _snake_case     :  Union[str, Any]     							=							False


    def      lowerCAmelCase_  (					self	:  Any					):
      _UpperCAmelCase						 =      DeiTModelTester(self					)
      _UpperCAmelCase						 =      ConfigTester(self   ,				config_class=__lowerCAmelCase   ,				has_text_modality=__lowerCAmelCase   ,				hidden_size=37					)


    def      lowerCAmelCase_  (					self	:  Tuple					):
      self.config_tester.run_common_tests()


    @unittest.skip(reason="""DeiT does not use inputs_embeds"""					)
    def      lowerCAmelCase_  (					self	:  int					):
      pass


    def      lowerCAmelCase_  (					self	:  Optional[int]					):
      _UpperCAmelCase						 =      self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        _UpperCAmelCase						 =      model_class(__lowerCAmelCase					)
        self.assertIsInstance(model.get_input_embeddings()   ,				(nn.Module)					)
        _UpperCAmelCase						 =      model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(__lowerCAmelCase   ,				nn.Linear					)					)


    def      lowerCAmelCase_  (					self	:  Tuple					):
      _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.forward					)
        # 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      lowerCAmelCase_  (					self	:  Optional[Any]					):
      _UpperCAmelCase						 =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__lowerCAmelCase					)


    def      lowerCAmelCase_  (					self	:  Dict					):
      _UpperCAmelCase						 =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_image_modeling(*__lowerCAmelCase					)


    def      lowerCAmelCase_  (					self	:  int					):
      _UpperCAmelCase						 =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*__lowerCAmelCase					)


    def      lowerCAmelCase_  (					self	:  Optional[int]   ,				__lowerCAmelCase	:  Any   ,				__lowerCAmelCase	:  Any   ,				__lowerCAmelCase	:  Dict=False					):
      _UpperCAmelCase						 =      super()._prepare_for_class(__lowerCAmelCase   ,				__lowerCAmelCase   ,				return_labels=__lowerCAmelCase					)

      if return_labels:
        if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
          del inputs_dict["labels"]

      return inputs_dict


    def      lowerCAmelCase_  (					self	:  Optional[Any]					):




      if not self.model_tester.is_training:
        return

      _UpperCAmelCase						 =      self.model_tester.prepare_config_and_inputs_for_common()
      _UpperCAmelCase						 =      True

      for model_class in self.all_model_classes:
        # DeiTForImageClassificationWithTeacher supports inference-only
        if (
            model_class in get_values(__lowerCAmelCase					)
            or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
        ):
          continue
        _UpperCAmelCase						 =      model_class(__lowerCAmelCase					)
        model.to(__lowerCAmelCase					)
        model.train()
        _UpperCAmelCase						 =      self._prepare_for_class(__lowerCAmelCase   ,				__lowerCAmelCase   ,				return_labels=__lowerCAmelCase					)
        _UpperCAmelCase						 =      model(**__lowerCAmelCase					).loss
        loss.backward()


    def      lowerCAmelCase_  (					self	:  Tuple					):
      _UpperCAmelCase						 =      self.model_tester.prepare_config_and_inputs_for_common()
      if not self.model_tester.is_training:
        return

      _UpperCAmelCase						 =      False
      _UpperCAmelCase						 =      True

      for model_class in self.all_model_classes:
        if model_class in get_values(__lowerCAmelCase					) or not model_class.supports_gradient_checkpointing:
          continue
        # DeiTForImageClassificationWithTeacher supports inference-only
        if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
          continue
        _UpperCAmelCase						 =      model_class(__lowerCAmelCase					)
        model.gradient_checkpointing_enable()
        model.to(__lowerCAmelCase					)
        model.train()
        _UpperCAmelCase						 =      self._prepare_for_class(__lowerCAmelCase   ,				__lowerCAmelCase   ,				return_labels=__lowerCAmelCase					)
        _UpperCAmelCase						 =      model(**__lowerCAmelCase					).loss
        loss.backward()


    def      lowerCAmelCase_  (					self	:  Union[str, Any]					):
      _UpperCAmelCase						 =      self.model_tester.prepare_config_and_inputs_for_common()

      _UpperCAmelCase						 =      [
          {"""title""": """multi_label_classification""", """num_labels""": 2, """dtype""": torch.float},
          {"""title""": """single_label_classification""", """num_labels""": 1, """dtype""": torch.long},
          {"""title""": """regression""", """num_labels""": 1, """dtype""": torch.float},
      ]

      for model_class in self.all_model_classes:
        if (
            model_class
            not in [
                *get_values(__lowerCAmelCase					),
                *get_values(__lowerCAmelCase					),
            ]
            or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
        ):
          continue

        for problem_type in problem_types:
          with self.subTest(msg=f'''Testing {model_class} with {problem_type["title"]}'''					):
            _UpperCAmelCase						 =      problem_type["""title"""]
            _UpperCAmelCase						 =      problem_type["""num_labels"""]

            _UpperCAmelCase						 =      model_class(__lowerCAmelCase					)
            model.to(__lowerCAmelCase					)
            model.train()

            _UpperCAmelCase						 =      self._prepare_for_class(__lowerCAmelCase   ,				__lowerCAmelCase   ,				return_labels=__lowerCAmelCase					)

            if problem_type["num_labels"] > 1:
              _UpperCAmelCase						 =      inputs["""labels"""].unsqueeze(1					).repeat(1   ,				problem_type["""num_labels"""]					)

            _UpperCAmelCase						 =      inputs["""labels"""].to(problem_type["""dtype"""]					)

            # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
            # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
            # they have the same size." which is a symptom something in wrong for the regression problem.
            # See https://github.com/huggingface/transformers/issues/11780
            with warnings.catch_warnings(record=__lowerCAmelCase					) as warning_list:
              _UpperCAmelCase						 =      model(**__lowerCAmelCase					).loss
            for w in warning_list:
              if "Using a target size that is different to the input size" in str(w.message					):
                raise ValueError(
                    f'''Something is going wrong in the regression problem: intercepted {w.message}'''					)

            loss.backward()


    @slow
    def      lowerCAmelCase_  (					self	:  Union[str, Any]					):




      for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _UpperCAmelCase						 =      DeiTModel.from_pretrained(__lowerCAmelCase					)
        self.assertIsNotNone(__lowerCAmelCase					)






def     __UpperCAmelCase					(			):







  """simple docstring"""
  _UpperCAmelCase						 =      Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""		)
  return image


@require_torch
@require_vision
class        a				(		unittest.TestCase   ):


    @cached_property
    def      lowerCAmelCase_  (					self	:  Any					):
      return (
          DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224"""					)
          if is_vision_available()
          else None
      )


    @slow
    def      lowerCAmelCase_  (					self	:  int					):
      _UpperCAmelCase						 =      DeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""					).to(
          __lowerCAmelCase					)

      _UpperCAmelCase						 =      self.default_image_processor
      _UpperCAmelCase						 =      prepare_img()
      _UpperCAmelCase						 =      image_processor(images=__lowerCAmelCase   ,				return_tensors="""pt"""					).to(__lowerCAmelCase					)

      # forward pass
      with torch.no_grad():
        _UpperCAmelCase						 =      model(**__lowerCAmelCase					)

      # verify the logits
      _UpperCAmelCase						 =      torch.Size((1, 1000)					)
      self.assertEqual(outputs.logits.shape   ,				__lowerCAmelCase					)

      _UpperCAmelCase						 =      torch.tensor([-1.0_266, 0.1_912, -1.2_861]					).to(__lowerCAmelCase					)

      self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,				__lowerCAmelCase   ,				atol=1e-4					)					)


    @slow
    @require_accelerate
    @require_torch_gpu
    def      lowerCAmelCase_  (					self	:  int					):
      _UpperCAmelCase						 =      DeiTModel.from_pretrained(
          """facebook/deit-base-distilled-patch16-224"""   ,				torch_dtype=torch.floataa   ,				device_map="""auto"""					)
      _UpperCAmelCase						 =      self.default_image_processor

      _UpperCAmelCase						 =      prepare_img()
      _UpperCAmelCase						 =      image_processor(images=__lowerCAmelCase   ,				return_tensors="""pt"""					)
      _UpperCAmelCase						 =      inputs.pixel_values.to(__lowerCAmelCase					)

      # forward pass to make sure inference works in fp16
      with torch.no_grad():
        _UpperCAmelCase						 =      model(__lowerCAmelCase					)

"""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					(			lowercase   ,lowercase		):







  """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(lowercase		):
      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					(			):







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







  """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(lowercase   ,lowercase		)  # Reverse the partitions.
  _UpperCAmelCase						 =      _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase   ,lowercase		)

  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     __UpperCAmelCase					(			):







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






@require_not_windows
@require_dill_gt_0_3_2
def     __UpperCAmelCase					(			):







  """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 lowercase		: x.reverse()
    _UpperCAmelCase						 =      _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase   ,[2, 1, 0]		)

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







  """simple docstring"""
  _UpperCAmelCase						 =      pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
  _UpperCAmelCase						 =      spark.range(20		).repartition(4		)

  # Partitions 0 and 2
  _UpperCAmelCase						 =      SparkExamplesIterable(lowercase		).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(lowercase   ,[0, 2]		)
  for i, (row_id, row_dict) in enumerate(lowercase		):
    _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(lowercase		).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(lowercase   ,[1, 3]		)
  for i, (row_id, row_dict) in enumerate(lowercase		):
    _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     __UpperCAmelCase					(			):







  """simple docstring"""
  _UpperCAmelCase						 =      pyspark.sql.SparkSession.builder.master("""local[*]"""		).appName("""pyspark"""		).getOrCreate()
  _UpperCAmelCase						 =      spark.range(1_00		).repartition(1		)
  _UpperCAmelCase						 =      Spark(lowercase		)
  # 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
