

'''simple docstring'''

import collections
import inspect
import unittest

from transformers import DinatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
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 DinatBackbone, DinatForImageClassification, DinatModel
						from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
						from PIL import Image

						from transformers import AutoImageProcessor







class lowerCamelCase      :

				def __init__(					self,	lowercase_,	lowercase_=13,	lowercase_=64,	lowercase_=4,	lowercase_=3,	lowercase_=16,	lowercase_=[1, 2, 1],	lowercase_=[2, 4, 8],	lowercase_=3,	lowercase_=[[3], [1, 2], [1]],	lowercase_=2.0,	lowercase_=True,	lowercase_=0.0,	lowercase_=0.0,	lowercase_=0.1,	lowercase_="gelu",	lowercase_=True,	lowercase_=0.02,	lowercase_=1E-5,	lowercase_=True,	lowercase_=None,	lowercase_=True,	lowercase_=10,	lowercase_=["stage1", "stage2"],	lowercase_=[1, 2],	)						->       Any:
					snake_case							 =				parent
					snake_case							 =				batch_size
					snake_case							 =				image_size
					snake_case							 =				patch_size
					snake_case							 =				num_channels
					snake_case							 =				embed_dim
					snake_case							 =				depths
					snake_case							 =				num_heads
					snake_case							 =				kernel_size
					snake_case							 =				dilations
					snake_case							 =				mlp_ratio
					snake_case							 =				qkv_bias
					snake_case							 =				hidden_dropout_prob
					snake_case							 =				attention_probs_dropout_prob
					snake_case							 =				drop_path_rate
					snake_case							 =				hidden_act
					snake_case							 =				patch_norm
					snake_case							 =				layer_norm_eps
					snake_case							 =				initializer_range
					snake_case							 =				is_training
					snake_case							 =				scope
					snake_case							 =				use_labels
					snake_case							 =				num_labels
					snake_case							 =				out_features
					snake_case							 =				out_indices

				def 				_lowerCamelCase						(					self   )						->       Tuple:
					snake_case							 =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

					snake_case							 =				None
					if self.use_labels:
						snake_case							 =				ids_tensor([self.batch_size],	self.num_labels   )

					snake_case							 =				self.get_config()

					return config, pixel_values, labels

				def 				_lowerCamelCase						(					self   )						->       Optional[int]:
					return DinatConfig(
					    num_labels=self.num_labels,	image_size=self.image_size,	patch_size=self.patch_size,	num_channels=self.num_channels,	embed_dim=self.embed_dim,	depths=self.depths,	num_heads=self.num_heads,	kernel_size=self.kernel_size,	dilations=self.dilations,	mlp_ratio=self.mlp_ratio,	qkv_bias=self.qkv_bias,	hidden_dropout_prob=self.hidden_dropout_prob,	attention_probs_dropout_prob=self.attention_probs_dropout_prob,	drop_path_rate=self.drop_path_rate,	hidden_act=self.hidden_act,	patch_norm=self.patch_norm,	layer_norm_eps=self.layer_norm_eps,	initializer_range=self.initializer_range,	out_features=self.out_features,	out_indices=self.out_indices,	)

				def 				_lowerCamelCase						(					self,	lowercase_,	lowercase_,	lowercase_   )						->       Optional[int]:
					snake_case							 =				DinatModel(config=lowercase_   )
					model.to(lowercase_   )
					model.eval()
					snake_case							 =				model(lowercase_   )

					snake_case							 =				snake_case							 =				(config.image_size // config.patch_size) // (2 ** (len(config.depths   ) - 1))
					snake_case							 =				int(config.embed_dim * 2 ** (len(config.depths   ) - 1)   )

					self.parent.assertEqual(
					    result.last_hidden_state.shape,	(self.batch_size, expected_height, expected_width, expected_dim)   )

				def 				_lowerCamelCase						(					self,	lowercase_,	lowercase_,	lowercase_   )						->       Tuple:
					snake_case							 =				DinatForImageClassification(lowercase_   )
					model.to(lowercase_   )
					model.eval()
					snake_case							 =				model(lowercase_,	labels=lowercase_   )
					self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.num_labels)   )

					# test greyscale images
					snake_case							 =				1
					snake_case							 =				DinatForImageClassification(lowercase_   )
					model.to(lowercase_   )
					model.eval()

					snake_case							 =				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
					snake_case							 =				model(lowercase_   )
					self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.num_labels)   )

				def 				_lowerCamelCase						(					self,	lowercase_,	lowercase_,	lowercase_   )						->       Tuple:
					snake_case							 =				DinatBackbone(config=lowercase_   )
					model.to(lowercase_   )
					model.eval()
					snake_case							 =				model(lowercase_   )

					# verify hidden states
					self.parent.assertEqual(len(result.feature_maps   ),	len(config.out_features   )   )
					self.parent.assertListEqual(list(result.feature_maps[0].shape   ),	[self.batch_size, model.channels[0], 16, 16]   )

					# verify channels
					self.parent.assertEqual(len(model.channels   ),	len(config.out_features   )   )

					# verify backbone works with out_features=None
					snake_case							 =				None
					snake_case							 =				DinatBackbone(config=lowercase_   )
					model.to(lowercase_   )
					model.eval()
					snake_case							 =				model(lowercase_   )

					# verify feature maps
					self.parent.assertEqual(len(result.feature_maps   ),	1   )
					self.parent.assertListEqual(list(result.feature_maps[0].shape   ),	[self.batch_size, model.channels[-1], 4, 4]   )

					# verify channels
					self.parent.assertEqual(len(model.channels   ),	1   )






				def 				_lowerCamelCase						(					self   )						->       List[str]:
					snake_case							 =				self.prepare_config_and_inputs()
					snake_case , snake_case , snake_case							 =				config_and_inputs
					snake_case							 =				{'pixel_values': pixel_values}
					return config, inputs_dict







@require_natten
@require_torch
class lowerCamelCase      (						__lowerCAmelCase							,							__lowerCAmelCase							,							unittest.TestCase							):
				snake_case_           =   (
				    (
				        DinatModel,
				        DinatForImageClassification,
				        DinatBackbone,
				    )
				    if is_torch_available()
				    else ()
				)
				snake_case_           =   (
				    {'''feature-extraction''': DinatModel, '''image-classification''': DinatForImageClassification}
				    if is_torch_available()
				    else {}
				)
				snake_case_           =   False

				snake_case_           =   False
				snake_case_           =   False
				snake_case_           =   False
				snake_case_           =   False

				def 				_lowerCamelCase						(					self   )						->       str:
					snake_case							 =				DinatModelTester(self   )
					snake_case							 =				ConfigTester(self,	config_class=lowercase_,	embed_dim=37   )

				def 				_lowerCamelCase						(					self   )						->       Dict:
					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 				_lowerCamelCase						(					self   )						->       List[str]:
					return

				def 				_lowerCamelCase						(					self   )						->       Optional[Any]:
					snake_case							 =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*lowercase_   )

				def 				_lowerCamelCase						(					self   )						->       int:
					snake_case							 =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*lowercase_   )

				def 				_lowerCamelCase						(					self   )						->       Any:
					snake_case							 =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_backbone(*lowercase_   )

				@unittest.skip(reason='Dinat does not use inputs_embeds'   )
				def 				_lowerCamelCase						(					self   )						->       int:
					pass

				@unittest.skip(reason='Dinat does not use feedforward chunking'   )
				def 				_lowerCamelCase						(					self   )						->       Tuple:
					pass

				def 				_lowerCamelCase						(					self   )						->       List[Any]:
					snake_case , snake_case							 =				self.model_tester.prepare_config_and_inputs_for_common()

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

				def 				_lowerCamelCase						(					self   )						->       int:
					snake_case , snake_case							 =				self.model_tester.prepare_config_and_inputs_for_common()

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

						snake_case							 =				['pixel_values']
						self.assertListEqual(arg_names[:1],	lowercase_   )

				def 				_lowerCamelCase						(					self   )						->       List[Any]:
					self.skipTest('Dinat\'s attention operation is handled entirely by NATTEN.'   )

				def 				_lowerCamelCase						(					self,	lowercase_,	lowercase_,	lowercase_,	lowercase_   )						->       int:
					snake_case							 =				model_class(lowercase_   )
					model.to(lowercase_   )
					model.eval()

					with torch.no_grad():
						snake_case							 =				model(**self._prepare_for_class(lowercase_,	lowercase_   )   )

					snake_case							 =				outputs.hidden_states

					snake_case							 =				getattr(
					    self.model_tester,	'expected_num_hidden_layers',	len(self.model_tester.depths   ) + 1   )
					self.assertEqual(len(lowercase_   ),	lowercase_   )

					# Dinat has a different seq_length
					snake_case							 =				(
					    config.patch_size
					    if isinstance(config.patch_size,	collections.abc.Iterable   )
					    else (config.patch_size, config.patch_size)
					)

					snake_case							 =				image_size[0] // patch_size[0]
					snake_case							 =				image_size[1] // patch_size[1]

					self.assertListEqual(
					    list(hidden_states[0].shape[-3:]   ),	[height, width, self.model_tester.embed_dim],	)

					if model_class.__name__ != "DinatBackbone":
						snake_case							 =				outputs.reshaped_hidden_states
						self.assertEqual(len(lowercase_   ),	lowercase_   )

						snake_case , snake_case , snake_case , snake_case							 =				reshaped_hidden_states[0].shape
						snake_case							 =				(
						    reshaped_hidden_states[0].view(lowercase_,	lowercase_,	lowercase_,	lowercase_   ).permute(0,	2,	3,	1   )
						)
						self.assertListEqual(
						    list(reshaped_hidden_states.shape[-3:]   ),	[height, width, self.model_tester.embed_dim],	)

				def 				_lowerCamelCase						(					self   )						->       Optional[Any]:
					snake_case , snake_case							 =				self.model_tester.prepare_config_and_inputs_for_common()

					snake_case							 =				(
					    self.model_tester.image_size
					    if isinstance(self.model_tester.image_size,	collections.abc.Iterable   )
					    else (self.model_tester.image_size, self.model_tester.image_size)
					)

					for model_class in self.all_model_classes:
						snake_case							 =				True
						self.check_hidden_states_output(lowercase_,	lowercase_,	lowercase_,	lowercase_   )

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

						self.check_hidden_states_output(lowercase_,	lowercase_,	lowercase_,	lowercase_   )

				@slow
				def 				_lowerCamelCase						(					self   )						->       int:
					for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						snake_case							 =				DinatModel.from_pretrained(lowercase_   )
						self.assertIsNotNone(lowercase_   )






				def 				_lowerCamelCase						(					self   )						->       Dict:
					snake_case , snake_case							 =				self.model_tester.prepare_config_and_inputs_for_common()

					snake_case							 =				_config_zero_init(lowercase_   )
					for model_class in self.all_model_classes:
						snake_case							 =				model_class(config=lowercase_   )
						for name, param in model.named_parameters():
							if "embeddings" not in name and 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''',	)







@require_natten
@require_vision
@require_torch
class lowerCamelCase      (						unittest.TestCase							):

				@cached_property
				def 				_lowerCamelCase						(					self   )						->       str:
					return AutoImageProcessor.from_pretrained('shi-labs/dinat-mini-in1k-224'   ) if is_vision_available() else None






				@slow
				def 				_lowerCamelCase						(					self   )						->       int:
					snake_case							 =				DinatForImageClassification.from_pretrained('shi-labs/dinat-mini-in1k-224'   ).to(lowercase_   )
					snake_case							 =				self.default_image_processor

					snake_case							 =				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
					snake_case							 =				image_processor(images=lowercase_,	return_tensors='pt'   ).to(lowercase_   )

					# forward pass
					with torch.no_grad():
						snake_case							 =				model(**lowercase_   )

					# verify the logits
					snake_case							 =				torch.Size((1, 1000)   )
					self.assertEqual(outputs.logits.shape,	lowercase_   )
					snake_case							 =				torch.tensor([-0.1_545, -0.7_667, 0.4_642]   ).to(lowercase_   )
					self.assertTrue(torch.allclose(outputs.logits[0, :3],	lowercase_,	atol=1E-4   )   )







@require_torch
@require_natten
class lowerCamelCase      (						unittest.TestCase							,							__lowerCAmelCase							):
				snake_case_           =   (DinatBackbone,) if is_torch_available() else ()
				snake_case_           =   DinatConfig

				def 				_lowerCamelCase						(					self   )						->       Optional[int]:
					snake_case							 =				DinatModelTester(self   )








'''simple docstring'''

from itertools import compress, repeat
from math import ceil, sqrt

def     __magic_name__       ( A							)			->		list[int]:

	if num <= 2:
		return []
	if num == 3:
		return [2]

	# Odd sieve for numbers in range [3, num - 1]
	snake_case							 =				bytearray(b'\x01'							) * ((num >> 1) - 1)

	for i in range(3					,   int(sqrt(A							)							) + 1					,   2							):
		if sieve[(i >> 1) - 1]:
			snake_case							 =				i**2
			snake_case							 =				repeat(
			    0					,   ceil((num - i_squared) / (i << 1)							)							)

	return [2] + list(compress(range(3					,   A					,   2							)					,   A							)							)


if __name__ == "__main__":
						import doctest

						doctest.testmod()





