





'''simple docstring'''






import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_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, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import (
        MODEL_MAPPING,
        BeitForImageClassification,
        BeitForMaskedImageModeling,
        BeitForSemanticSegmentation,
        BeitModel,
    )
    from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    import PIL
    from PIL import Image

    from transformers import BeitImageProcessor
class 					__snake_case  :






   '''simple docstring'''





   def __init__(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE=1_0_0	,		__SCREAMING_SNAKE_CASE=1_3	,		__SCREAMING_SNAKE_CASE=3_0	,		__SCREAMING_SNAKE_CASE=2	,		__SCREAMING_SNAKE_CASE=3	,		__SCREAMING_SNAKE_CASE=True	,		__SCREAMING_SNAKE_CASE=True	,		__SCREAMING_SNAKE_CASE=3_2	,		__SCREAMING_SNAKE_CASE=4	,		__SCREAMING_SNAKE_CASE=4	,		__SCREAMING_SNAKE_CASE=3_7	,		__SCREAMING_SNAKE_CASE="gelu"	,		__SCREAMING_SNAKE_CASE=0.1	,		__SCREAMING_SNAKE_CASE=0.1	,		__SCREAMING_SNAKE_CASE=1_0	,		__SCREAMING_SNAKE_CASE=0.02	,		__SCREAMING_SNAKE_CASE=3	,		__SCREAMING_SNAKE_CASE=None	,		__SCREAMING_SNAKE_CASE=[0, 1, 2, 3]	,		):
    snake_case__     :					Tuple    =  parent
    snake_case__     :					Optional[Any]    =  1_0_0
    snake_case__     :					List[Any]    =  batch_size
    snake_case__     :					List[str]    =  image_size
    snake_case__     :					int    =  patch_size
    snake_case__     :					Tuple    =  num_channels
    snake_case__     :					List[str]    =  is_training
    snake_case__     :					List[str]    =  use_labels
    snake_case__     :					int    =  hidden_size
    snake_case__     :					List[str]    =  num_hidden_layers
    snake_case__     :					Any    =  num_attention_heads
    snake_case__     :					Union[str, Any]    =  intermediate_size
    snake_case__     :					str    =  hidden_act
    snake_case__     :					Optional[Any]    =  hidden_dropout_prob
    snake_case__     :					Optional[int]    =  attention_probs_dropout_prob
    snake_case__     :					Optional[Any]    =  type_sequence_label_size
    snake_case__     :					Optional[Any]    =  initializer_range
    snake_case__     :					Union[str, Any]    =  scope
    snake_case__     :					Any    =  out_indices
    snake_case__     :					str    =  num_labels

    # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
    snake_case__     :					Optional[Any]    =  (image_size // patch_size) ** 2
    snake_case__     :					str    =  num_patches + 1

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Optional[Any]    =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

    snake_case__     :					int    =  None
    snake_case__     :					List[Any]    =  None
    if self.use_labels:
     snake_case__     :					List[str]    =  ids_tensor([self.batch_size]	,		self.type_sequence_label_size       )
     snake_case__     :					Optional[Any]    =  ids_tensor([self.batch_size, self.image_size, self.image_size]	,		self.num_labels       )

    snake_case__     :					Dict    =  self.get_config()

    return config, pixel_values, labels, pixel_labels

   def 		__UpperCamelCase						(							self       ):
    return BeitConfig(
        vocab_size=self.vocab_size	,		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=__SCREAMING_SNAKE_CASE	,		initializer_range=self.initializer_range	,		out_indices=self.out_indices	,		)

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Dict    =  BeitModel(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					Optional[int]    =  model(__SCREAMING_SNAKE_CASE       )
    self.parent.assertEqual(result.last_hidden_state.shape	,		(self.batch_size, self.seq_length, self.hidden_size)       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					int    =  BeitForMaskedImageModeling(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					List[Any]    =  model(__SCREAMING_SNAKE_CASE       )
    self.parent.assertEqual(result.logits.shape	,		(self.batch_size, self.seq_length - 1, self.vocab_size)       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Union[str, Any]    =  self.type_sequence_label_size
    snake_case__     :					Any    =  BeitForImageClassification(__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					Any    =  model(__SCREAMING_SNAKE_CASE	,		labels=__SCREAMING_SNAKE_CASE       )
    self.parent.assertEqual(result.logits.shape	,		(self.batch_size, self.type_sequence_label_size)       )

    # test greyscale images
    snake_case__     :					Any    =  1
    snake_case__     :					Any    =  BeitForImageClassification(__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()

    snake_case__     :					Any    =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
    snake_case__     :					List[str]    =  model(__SCREAMING_SNAKE_CASE	,		labels=__SCREAMING_SNAKE_CASE       )
    self.parent.assertEqual(result.logits.shape	,		(self.batch_size, self.type_sequence_label_size)       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Dict    =  self.num_labels
    snake_case__     :					str    =  BeitForSemanticSegmentation(__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					int    =  model(__SCREAMING_SNAKE_CASE       )
    self.parent.assertEqual(
        result.logits.shape	,		(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)       )
    snake_case__     :					Any    =  model(__SCREAMING_SNAKE_CASE	,		labels=__SCREAMING_SNAKE_CASE       )
    self.parent.assertEqual(
        result.logits.shape	,		(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)       )


   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Dict    =  self.prepare_config_and_inputs()
    snake_case__  ,							snake_case__  ,							snake_case__  ,							snake_case__     :					List[Any]    =  config_and_inputs
    snake_case__     :					List[str]    =  {"""pixel_values""": pixel_values}
    return config, inputs_dict



@require_torch
class 					__snake_case  (					__SCREAMING_SNAKE_CASE		,				__SCREAMING_SNAKE_CASE		,				unittest.TestCase							):






   '''simple docstring'''





   lowerCamelCase__    							=      (
       (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
       if is_torch_available()
       else ()
   )
   lowerCamelCase__    							=      (
       {
           '''feature-extraction''': BeitModel,
           '''image-classification''': BeitForImageClassification,
           '''image-segmentation''': BeitForSemanticSegmentation,
       }
       if is_torch_available()
       else {}
   )

   lowerCamelCase__    							=      False
   lowerCamelCase__    							=      False
   lowerCamelCase__    							=      False

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					List[str]    =  BeitModelTester(self       )
    snake_case__     :					List[Any]    =  ConfigTester(self	,		config_class=__SCREAMING_SNAKE_CASE	,		has_text_modality=__SCREAMING_SNAKE_CASE	,		hidden_size=3_7       )

   def 		__UpperCamelCase						(							self       ):
    self.config_tester.run_common_tests()

   @unittest.skip(reason="""BEiT does not use inputs_embeds"""       )
   def 		__UpperCamelCase						(							self       ):
    pass

   @require_torch_multi_gpu
   @unittest.skip(reason="""BEiT has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""       )
   def 		__UpperCamelCase						(							self       ):
    pass

   def 		__UpperCamelCase						(							self       ):
    snake_case__  ,							snake_case__     :					Any    =  self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
     snake_case__     :					List[Any]    =  model_class(__SCREAMING_SNAKE_CASE       )
     self.assertIsInstance(model.get_input_embeddings()	,		(nn.Module)       )
     snake_case__     :					int    =  model.get_output_embeddings()
     self.assertTrue(x is None or isinstance(__SCREAMING_SNAKE_CASE	,		nn.Linear       )       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__  ,							snake_case__     :					Tuple    =  self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
     snake_case__     :					Tuple    =  model_class(__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Optional[Any]    =  inspect.signature(model.forward       )
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     snake_case__     :					int    =  [*signature.parameters.keys()]

     snake_case__     :					Any    =  ["""pixel_values"""]
     self.assertListEqual(arg_names[:1]	,		__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					int    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Optional[int]    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Optional[int]    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_image_classification(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Dict    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_semantic_segmentation(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    if not self.model_tester.is_training:
     return

    snake_case__  ,							snake_case__     :					List[Any]    =  self.model_tester.prepare_config_and_inputs_for_common()
    snake_case__     :					Dict    =  True

    for model_class in self.all_model_classes:
     # we don't test BeitForMaskedImageModeling
     if model_class in [*get_values(__SCREAMING_SNAKE_CASE       ), BeitForMaskedImageModeling]:
      continue

     snake_case__     :					List[str]    =  model_class(__SCREAMING_SNAKE_CASE       )
     model.to(__SCREAMING_SNAKE_CASE       )
     model.train()
     snake_case__     :					List[Any]    =  self._prepare_for_class(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		return_labels=__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Dict    =  model(**__SCREAMING_SNAKE_CASE       ).loss
     loss.backward()

   def 		__UpperCamelCase						(							self       ):
    snake_case__  ,							snake_case__     :					Optional[Any]    =  self.model_tester.prepare_config_and_inputs_for_common()
    if not self.model_tester.is_training:
     return

    snake_case__     :					Optional[int]    =  False
    snake_case__     :					Any    =  True

    for model_class in self.all_model_classes:
     # we don't test BeitForMaskedImageModeling
     if (
         model_class in [*get_values(__SCREAMING_SNAKE_CASE       ), BeitForMaskedImageModeling]
         or not model_class.supports_gradient_checkpointing
     ):
      continue

     snake_case__     :					str    =  model_class(__SCREAMING_SNAKE_CASE       )
     model.gradient_checkpointing_enable()
     model.to(__SCREAMING_SNAKE_CASE       )
     model.train()
     snake_case__     :					str    =  self._prepare_for_class(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		return_labels=__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Dict    =  model(**__SCREAMING_SNAKE_CASE       ).loss
     loss.backward()

   def 		__UpperCamelCase						(							self       ):
    snake_case__  ,							snake_case__     :					Optional[Any]    =  self.model_tester.prepare_config_and_inputs_for_common()

    snake_case__     :					int    =  _config_zero_init(__SCREAMING_SNAKE_CASE       )
    for model_class in self.all_model_classes:
     snake_case__     :					Dict    =  model_class(config=__SCREAMING_SNAKE_CASE       )
     for name, param in model.named_parameters():
      # we skip lambda parameters as these require special initial values
      # determined by config.layer_scale_init_value
      if "lambda" in name:
       continue
      if param.requires_grad:
       self.assertIn(
           ((param.data.mean() * 1e9).round() / 1e9).item()	,		[0.0, 1.0]	,		msg=f"Parameter {name} of model {model_class} seems not properly initialized"	,		)


   @slow
   def 		__UpperCamelCase						(							self       ):
    for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     snake_case__     :					List[str]    =  BeitModel.from_pretrained(__SCREAMING_SNAKE_CASE       )
     self.assertIsNotNone(__SCREAMING_SNAKE_CASE       )



def        UpperCamelCase__		(						)				->			str:
 '''simple docstring'''






 snake_case__     :					Optional[int]    =  Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""      )
 return image





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






   '''simple docstring'''





   @cached_property
   def 		__UpperCamelCase						(							self       ):
    return BeitImageProcessor.from_pretrained("""microsoft/beit-base-patch16-224"""       ) if is_vision_available() else None

   @slow
   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Optional[Any]    =  BeitForMaskedImageModeling.from_pretrained("""microsoft/beit-base-patch16-224-pt22k"""       ).to(__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Union[str, Any]    =  self.default_image_processor
    snake_case__     :					Optional[int]    =  prepare_img()
    snake_case__     :					Dict    =  image_processor(images=__SCREAMING_SNAKE_CASE	,		return_tensors="""pt"""       ).pixel_values.to(__SCREAMING_SNAKE_CASE       )

    # prepare bool_masked_pos
    snake_case__     :					str    =  torch.ones((1, 1_9_6)	,		dtype=torch.bool       ).to(__SCREAMING_SNAKE_CASE       )

    # forward pass
    with torch.no_grad():
     snake_case__     :					Optional[Any]    =  model(pixel_values=__SCREAMING_SNAKE_CASE	,		bool_masked_pos=__SCREAMING_SNAKE_CASE       )
    snake_case__     :					int    =  outputs.logits

    # verify the logits
    snake_case__     :					Dict    =  torch.Size((1, 1_9_6, 8_1_9_2)       )
    self.assertEqual(logits.shape	,		__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Optional[int]    =  torch.tensor(
        [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]       ).to(__SCREAMING_SNAKE_CASE       )

    self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]	,		__SCREAMING_SNAKE_CASE	,		atol=1e-2       )       )

   @slow
   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Dict    =  BeitForImageClassification.from_pretrained("""microsoft/beit-base-patch16-224"""       ).to(__SCREAMING_SNAKE_CASE       )

    snake_case__     :					List[Any]    =  self.default_image_processor
    snake_case__     :					Tuple    =  prepare_img()
    snake_case__     :					int    =  image_processor(images=__SCREAMING_SNAKE_CASE	,		return_tensors="""pt"""       ).to(__SCREAMING_SNAKE_CASE       )

    # forward pass
    with torch.no_grad():
     snake_case__     :					str    =  model(**__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Optional[int]    =  outputs.logits

    # verify the logits
    snake_case__     :					List[str]    =  torch.Size((1, 1_0_0_0)       )
    self.assertEqual(logits.shape	,		__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Union[str, Any]    =  torch.tensor([-1.2385, -1.0987, -1.0108]       ).to(__SCREAMING_SNAKE_CASE       )

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

    snake_case__     :					List[str]    =  2_8_1
    self.assertEqual(logits.argmax(-1       ).item()	,		__SCREAMING_SNAKE_CASE       )

   @slow
   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					List[Any]    =  BeitForImageClassification.from_pretrained("""microsoft/beit-large-patch16-224-pt22k-ft22k"""       ).to(
        __SCREAMING_SNAKE_CASE       )

    snake_case__     :					Union[str, Any]    =  self.default_image_processor
    snake_case__     :					int    =  prepare_img()
    snake_case__     :					Any    =  image_processor(images=__SCREAMING_SNAKE_CASE	,		return_tensors="""pt"""       ).to(__SCREAMING_SNAKE_CASE       )

    # forward pass
    with torch.no_grad():
     snake_case__     :					Optional[Any]    =  model(**__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Dict    =  outputs.logits

    # verify the logits
    snake_case__     :					str    =  torch.Size((1, 2_1_8_4_1)       )
    self.assertEqual(logits.shape	,		__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Optional[int]    =  torch.tensor([1.6881, -0.2787, 0.5901]       ).to(__SCREAMING_SNAKE_CASE       )

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

    snake_case__     :					Union[str, Any]    =  2_3_9_6
    self.assertEqual(logits.argmax(-1       ).item()	,		__SCREAMING_SNAKE_CASE       )

   @slow
   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Optional[Any]    =  BeitForSemanticSegmentation.from_pretrained("""microsoft/beit-base-finetuned-ade-640-640"""       )
    snake_case__     :					List[Any]    =  model.to(__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Any    =  BeitImageProcessor(do_resize=__SCREAMING_SNAKE_CASE	,		size=6_4_0	,		do_center_crop=__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Any    =  load_dataset("""hf-internal-testing/fixtures_ade20k"""	,		split="""test"""       )
    snake_case__     :					int    =  Image.open(ds[0]["""file"""]       )
    snake_case__     :					Tuple    =  image_processor(images=__SCREAMING_SNAKE_CASE	,		return_tensors="""pt"""       ).to(__SCREAMING_SNAKE_CASE       )

    # forward pass
    with torch.no_grad():
     snake_case__     :					List[str]    =  model(**__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Optional[Any]    =  outputs.logits

    # verify the logits
    snake_case__     :					str    =  torch.Size((1, 1_5_0, 1_6_0, 1_6_0)       )
    self.assertEqual(logits.shape	,		__SCREAMING_SNAKE_CASE       )

    snake_case__     :					List[Any]    =  version.parse(PIL.__version__       ) < version.parse("""9.0.0"""       )

    if is_pillow_less_than_a:
     snake_case__     :					List[str]    =  torch.tensor(
         [
             [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
             [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
             [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
         ]	,		device=__SCREAMING_SNAKE_CASE	,		)
    else:
     snake_case__     :					Tuple    =  torch.tensor(
         [
             [[-4.8960, -2.3688, -3.0355], [-2.8478, -0.9836, -1.7418], [-2.9449, -1.3332, -2.1456]],
             [[-5.8081, -3.4124, -4.1006], [-3.8561, -2.2081, -3.0323], [-3.8365, -2.4601, -3.3669]],
             [[-0.0309, 3.9868, 4.0540], [2.9640, 4.6877, 4.9976], [3.2081, 4.7690, 4.9942]],
         ]	,		device=__SCREAMING_SNAKE_CASE	,		)

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


   @slow
   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Optional[int]    =  BeitForSemanticSegmentation.from_pretrained("""microsoft/beit-base-finetuned-ade-640-640"""       )
    snake_case__     :					str    =  model.to(__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Optional[int]    =  BeitImageProcessor(do_resize=__SCREAMING_SNAKE_CASE	,		size=6_4_0	,		do_center_crop=__SCREAMING_SNAKE_CASE       )

    snake_case__     :					List[str]    =  load_dataset("""hf-internal-testing/fixtures_ade20k"""	,		split="""test"""       )
    snake_case__     :					Optional[int]    =  Image.open(ds[0]["""file"""]       )
    snake_case__     :					List[str]    =  image_processor(images=__SCREAMING_SNAKE_CASE	,		return_tensors="""pt"""       ).to(__SCREAMING_SNAKE_CASE       )

    # forward pass
    with torch.no_grad():
     snake_case__     :					int    =  model(**__SCREAMING_SNAKE_CASE       )

    snake_case__     :					List[Any]    =  outputs.logits.detach().cpu()

    snake_case__     :					int    =  image_processor.post_process_semantic_segmentation(outputs=__SCREAMING_SNAKE_CASE	,		target_sizes=[(5_0_0, 3_0_0)]       )
    snake_case__     :					Tuple    =  torch.Size((5_0_0, 3_0_0)       )
    self.assertEqual(segmentation[0].shape	,		__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Optional[Any]    =  image_processor.post_process_semantic_segmentation(outputs=__SCREAMING_SNAKE_CASE       )
    snake_case__     :					List[Any]    =  torch.Size((1_6_0, 1_6_0)       )
    self.assertEqual(segmentation[0].shape	,		__SCREAMING_SNAKE_CASE       )







'''simple docstring'''



from math import isqrt







def   a__			(    a__		):




     """simple docstring"""


     return all(number % divisor != 0 for divisor in range(2 ,       isqrt(a__		) + 1		)		)







def   a__			(    a__ = 10**6		):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE      = 0
     __SCREAMING_SNAKE_CASE      = 1
     __SCREAMING_SNAKE_CASE      = 7
     while prime_candidate < max_prime:
          primes_count += is_prime(a__		)

          cube_index += 1
          prime_candidate += 6 * cube_index

     return primes_count


if __name__ == "__main__":
  print(f"""{solution() = }""")






