

import argparse
import collections

import jax
import jax.numpy as jnp
import torch
import torch.nn as nn
from clip.model import CLIP
from flax.training import checkpoints
from huggingface_hub import Repository

from transformers import (
    CLIPTokenizer,
    OwlViTConfig,
    OwlViTForObjectDetection,
    OwlViTImageProcessor,
    OwlViTModel,
    OwlViTProcessor,
)


_lowerCAmelCase :	Union[str, Any]     				=    {
    'vit_b32': {
        'embed_dim': 5_1_2,
        'image_resolution': 7_6_8,
        'context_length': 1_6,
        'vocab_size': 4_9_4_0_8,
        'vision_layers': 1_2,
        'vision_width': 7_6_8,
        'vision_patch_size': 3_2,
        'transformer_width': 5_1_2,
        'transformer_heads': 8,
        'transformer_layers': 1_2,
    },
    'vit_b16': {
        'embed_dim': 5_1_2,
        'image_resolution': 7_6_8,
        'context_length': 1_6,
        'vocab_size': 4_9_4_0_8,
        'vision_layers': 1_2,
        'vision_width': 7_6_8,
        'vision_patch_size': 1_6,
        'transformer_width': 5_1_2,
        'transformer_heads': 8,
        'transformer_layers': 1_2,
    },
    'vit_l14': {
        'embed_dim': 7_6_8,
        'image_resolution': 8_4_0,
        'context_length': 1_6,
        'vocab_size': 4_9_4_0_8,
        'vision_layers': 2_4,
        'vision_width': 1_0_2_4,
        'vision_patch_size': 1_4,
        'transformer_width': 7_6_8,
        'transformer_heads': 1_2,
        'transformer_layers': 1_2,
    },
}







def a_		(       UpperCamelCase_   :							Optional[Any]   ,					UpperCamelCase_   :							List[Any]=""   ,					UpperCamelCase_   :							List[str]="/"  )							->   List[str]:







 """simple docstring"""

 lowerCamelCase      						=    []

 for k, v in params.items():
  lowerCamelCase      						=    parent_key + sep + k if parent_key else k

  if isinstance(UpperCamelCase_   ,					collections.MutableMapping  ):
   items.extend(flatten_nested_dict(UpperCamelCase_   ,					UpperCamelCase_   ,					sep=UpperCamelCase_  ).items()  )
  else:
   items.append((new_key, v)  )
 return dict(UpperCamelCase_  )







def a_		(       UpperCamelCase_   :							Optional[int]  )							->   int:







 """simple docstring"""

 return jax.tree_util.tree_map(lambda UpperCamelCase_  : x.astype(jnp.floataa  ) if x.dtype == jnp.bfloataa else x   ,					UpperCamelCase_  )







def a_		(       UpperCamelCase_   :							Tuple   ,					UpperCamelCase_   :							Union[str, Any]  )							->   Any:







 """simple docstring"""

 lowerCamelCase						,	lowerCamelCase						,	lowerCamelCase      						=    pt_attn_layer.in_proj_weight.chunk(3   ,					dim=0  )
 lowerCamelCase						,	lowerCamelCase						,	lowerCamelCase      						=    pt_attn_layer.in_proj_bias.chunk(3   ,					dim=0  )

 lowerCamelCase      						=    pt_attn_layer.out_proj.weight
 lowerCamelCase      						=    pt_attn_layer.out_proj.bias

 lowerCamelCase      						=    q_proj
 lowerCamelCase      						=    q_proj_bias

 lowerCamelCase      						=    k_proj
 lowerCamelCase      						=    k_proj_bias

 lowerCamelCase      						=    v_proj
 lowerCamelCase      						=    v_proj_bias

 lowerCamelCase      						=    out_proj_weights
 lowerCamelCase      						=    out_proj_bias







def a_		(       UpperCamelCase_   :							Optional[Any]   ,					UpperCamelCase_   :							List[str]  )							->   Any:







 """simple docstring"""

 copy_linear(hf_mlp.fca   ,					pt_mlp.c_fc  )
 copy_linear(hf_mlp.fca   ,					pt_mlp.c_proj  )







def a_		(       UpperCamelCase_   :							List[Any]   ,					UpperCamelCase_   :							Any  )							->   Optional[int]:







 """simple docstring"""

 lowerCamelCase      						=    pt_linear.weight
 lowerCamelCase      						=    pt_linear.bias







def a_		(       UpperCamelCase_   :							Any   ,					UpperCamelCase_   :							Optional[Any]  )							->   Optional[int]:







 """simple docstring"""

 copy_linear(hf_layer.layer_norma   ,					pt_layer.ln_a  )
 copy_linear(hf_layer.layer_norma   ,					pt_layer.ln_a  )

 # copy MLP
 copy_mlp(hf_layer.mlp   ,					pt_layer.mlp  )

 # copy attn
 copy_attn_layer(hf_layer.self_attn   ,					pt_layer.attn  )







def a_		(       UpperCamelCase_   :							str   ,					UpperCamelCase_   :							List[str]  )							->   Optional[Any]:







 """simple docstring"""

 for hf_layer, pt_layer in zip(UpperCamelCase_   ,					UpperCamelCase_  ):
  copy_layer(UpperCamelCase_   ,					UpperCamelCase_  )







def a_		(       UpperCamelCase_   :							str   ,					UpperCamelCase_   :							List[Any]  )							->   Optional[Any]:







 """simple docstring"""

 lowerCamelCase      						=    pt_model.token_embedding.weight
 lowerCamelCase      						=    pt_model.positional_embedding

 # copy layer norm
 copy_linear(hf_encoder.final_layer_norm   ,					pt_model.ln_final  )

 # copy hidden layers
 copy_layers(hf_encoder.encoder.layers   ,					pt_model.transformer.resblocks  )







def a_		(       UpperCamelCase_   :							int   ,					UpperCamelCase_   :							List[str]  )							->   List[str]:







 """simple docstring"""

 lowerCamelCase      						=    pt_model.text_projection.data.T

 # copy text encoder
 copy_encoder(hf_model.text_model   ,					UpperCamelCase_  )







def a_		(       UpperCamelCase_   :							Any   ,					UpperCamelCase_   :							Tuple  )							->   Optional[int]:







 """simple docstring"""

 lowerCamelCase      						=    pt_model.visual.proj.data.T

 # copy layer norms
 copy_linear(hf_model.vision_model.pre_layernorm   ,					pt_model.visual.ln_pre  )
 copy_linear(hf_model.vision_model.post_layernorm   ,					pt_model.visual.ln_post  )

 # copy embeds
 lowerCamelCase      						=    pt_model.visual.conva.weight.data
 lowerCamelCase      						=    pt_model.visual.class_embedding
 lowerCamelCase      						=    pt_model.visual.positional_embedding.data

 # copy encoder
 copy_layers(hf_model.vision_model.encoder.layers   ,					pt_model.visual.transformer.resblocks  )







def a_		(       UpperCamelCase_   :							List[Any]   ,					UpperCamelCase_   :							Any  )							->   Optional[Any]:







 """simple docstring"""

 lowerCamelCase      						=    flatten_nested_dict(flax_params['backbone']['merged_class_token']  )

 lowerCamelCase      						=    torch.from_numpy(flax_class_token_params['scale']  )
 lowerCamelCase      						=    torch.from_numpy(flax_class_token_params['bias']  )
 lowerCamelCase      						=    nn.Parameter(UpperCamelCase_  )
 lowerCamelCase      						=    nn.Parameter(UpperCamelCase_  )







def a_		(       UpperCamelCase_   :							Tuple   ,					UpperCamelCase_   :							List[str]  )							->   List[Any]:







 """simple docstring"""

 lowerCamelCase      						=    hf_model.state_dict()
 lowerCamelCase      						=    {}

 # Rename class prediction head flax params to pytorch HF
 lowerCamelCase      						=    flatten_nested_dict(flax_params['class_head']  )

 for flax_key, v in flax_class_params.items():
  lowerCamelCase      						=    flax_key.replace('/'   ,					'.'  )
  lowerCamelCase      						=    torch_key.replace('.kernel'   ,					'.weight'  )
  lowerCamelCase      						=    torch_key.replace('Dense_0'   ,					'dense0'  )
  lowerCamelCase      						=    'class_head.' + torch_key

  if "weight" in torch_key and v.ndim == 2:
   lowerCamelCase      						=    v.T

  lowerCamelCase      						=    nn.Parameter(torch.from_numpy(UpperCamelCase_  )  )

 # Rename box prediction box flax params to pytorch HF
 lowerCamelCase      						=    flatten_nested_dict(flax_params['obj_box_head']  )

 for flax_key, v in flax_box_params.items():
  lowerCamelCase      						=    flax_key.replace('/'   ,					'.'  )
  lowerCamelCase      						=    torch_key.replace('.kernel'   ,					'.weight'  )
  lowerCamelCase      						=    torch_key.replace('_'   ,					''  ).lower()
  lowerCamelCase      						=    'box_head.' + torch_key

  if "weight" in torch_key and v.ndim == 2:
   lowerCamelCase      						=    v.T

  lowerCamelCase      						=    nn.Parameter(torch.from_numpy(UpperCamelCase_  )  )

 # Copy flax params to PyTorch params
 for name, param in new_params.items():
  if name in pt_params.keys():
   pt_params[name].copy_(UpperCamelCase_  )







def a_		(       UpperCamelCase_   :							Optional[Any]   ,					UpperCamelCase_   :							Optional[int]  )							->   str:







 """simple docstring"""

 for k, v in flax_attn_params.items():
  if k.startswith('transformer'  ):
   lowerCamelCase      						=    k.replace('transformer.resblocks'   ,					'text_model.encoder.layers'  )
  else:
   lowerCamelCase      						=    k.replace('visual.transformer.resblocks'   ,					'vision_model.encoder.layers'  )

  lowerCamelCase      						=    torch_key.replace('attn'   ,					'self_attn'  )
  lowerCamelCase      						=    torch_key.replace('key'   ,					'k_proj'  )
  lowerCamelCase      						=    torch_key.replace('value'   ,					'v_proj'  )
  lowerCamelCase      						=    torch_key.replace('query'   ,					'q_proj'  )
  lowerCamelCase      						=    torch_key.replace('out'   ,					'out_proj'  )

  if "bias" in torch_key and v.ndim == 2:
   lowerCamelCase      						=    v.shape[0] * v.shape[1]
   lowerCamelCase      						=    v.reshape(UpperCamelCase_  )

  if "weight" in torch_key and "out" in torch_key:
   lowerCamelCase      						=    (v.shape[0] * v.shape[1], v.shape[2])
   lowerCamelCase      						=    v.reshape(UpperCamelCase_  ).T

  if "weight" in torch_key and "out" not in torch_key:
   lowerCamelCase      						=    (v.shape[0], v.shape[1] * v.shape[2])
   lowerCamelCase      						=    v.reshape(UpperCamelCase_  ).T

  # Copy flax CLIP attn params to HF PyTorch params
  lowerCamelCase      						=    torch.from_numpy(UpperCamelCase_  )
  hf_backbone.state_dict()[torch_key].copy_(UpperCamelCase_  )







def a_		(       UpperCamelCase_   :							str  )							->   int:







 """simple docstring"""

 lowerCamelCase      						=    {}
 lowerCamelCase      						=    []

 for k, v in params.items():
  if "attn." in k:
   lowerCamelCase      						=    k[: k.rindex('attn.'  ) + 5]
   if base in processed_attn_layers:
    continue

   processed_attn_layers.append(UpperCamelCase_  )
   lowerCamelCase      						=    params[base + 'out.weight'].shape[-1]
   lowerCamelCase      						=    params[base + 'out.weight'].reshape(UpperCamelCase_   ,					UpperCamelCase_  ).T
   lowerCamelCase      						=    params[base + 'out.bias']
  else:
   lowerCamelCase      						=    v
 return new_params







def a_		(       UpperCamelCase_   :							Union[str, Any]   ,					UpperCamelCase_   :							Optional[int]  )							->   Any:







 """simple docstring"""

 lowerCamelCase      						=    CLIP(**UpperCamelCase_  )
 torch_model.eval()
 lowerCamelCase      						=    torch_model.state_dict()

 lowerCamelCase      						=    flatten_nested_dict(flax_params['backbone']['clip']  )
 lowerCamelCase      						=    {}

 for flax_key, v in flax_clip_params.items():
  lowerCamelCase      						=    flax_key.replace('/'   ,					'.'  )
  lowerCamelCase      						=    torch_key.replace('text.token_embedding.embedding'   ,					'token_embedding.kernel'  )

  if (
      torch_key.startswith('text.transformer'  )
      or torch_key.startswith('text.text_projection'  )
      or torch_key.startswith('text.ln_final'  )
      or torch_key.startswith('text.positional_embedding'  )
  ):
   lowerCamelCase      						=    torch_key[5:]

  lowerCamelCase      						=    torch_key.replace('text_projection.kernel'   ,					'text_projection'  )
  lowerCamelCase      						=    torch_key.replace('visual.proj.kernel'   ,					'visual.proj'  )
  lowerCamelCase      						=    torch_key.replace('.scale'   ,					'.weight'  )
  lowerCamelCase      						=    torch_key.replace('.kernel'   ,					'.weight'  )

  if "conv" in torch_key or "downsample.0.weight" in torch_key:
   lowerCamelCase      						=    v.transpose(3   ,					2   ,					0   ,					1  )

  elif "weight" in torch_key and v.ndim == 2 and "embedding" not in torch_key:
   # Fully connected layers are transposed, embeddings are not
   lowerCamelCase      						=    v.T

  lowerCamelCase      						=    v

 lowerCamelCase      						=    _convert_attn_layers(UpperCamelCase_  )
 new_torch_params.update(UpperCamelCase_  )
 lowerCamelCase      						=    {}

 # Copy flax CLIP backbone params to PyTorch params
 for name, param in new_torch_params.items():
  if name in torch_clip_params.keys():
   lowerCamelCase      						=    torch.from_numpy(new_torch_params[name]  )
   torch_clip_params[name].copy_(UpperCamelCase_  )
  else:
   lowerCamelCase      						=    param

 return torch_clip_params, torch_model, attn_params







@torch.no_grad()
def a_		(       UpperCamelCase_   :							List[Any]   ,					UpperCamelCase_   :							List[Any]   ,					UpperCamelCase_   :							Union[str, Any]   ,					UpperCamelCase_   :							int   ,					UpperCamelCase_   :							List[Any]=None  )							->   Optional[int]:







 """simple docstring"""

 lowerCamelCase      						=    Repository(UpperCamelCase_   ,					clone_from=f'''google/{pytorch_dump_folder_path}'''  )
 repo.git_pull()

 if config_path is not None:
  lowerCamelCase      						=    OwlViTConfig.from_pretrained(UpperCamelCase_  )
 else:
  lowerCamelCase      						=    OwlViTConfig()

 lowerCamelCase      						=    OwlViTModel(UpperCamelCase_  ).eval()
 lowerCamelCase      						=    OwlViTForObjectDetection(UpperCamelCase_  ).eval()

 copy_text_model_and_projection(UpperCamelCase_   ,					UpperCamelCase_  )
 copy_vision_model_and_projection(UpperCamelCase_   ,					UpperCamelCase_  )
 lowerCamelCase      						=    pt_backbone.logit_scale
 copy_flax_attn_params(UpperCamelCase_   ,					UpperCamelCase_  )

 lowerCamelCase      						=    hf_backbone
 copy_class_merge_token(UpperCamelCase_   ,					UpperCamelCase_  )
 copy_class_box_heads(UpperCamelCase_   ,					UpperCamelCase_  )

 # Save HF model
 hf_model.save_pretrained(repo.local_dir  )

 # Initialize image processor
 lowerCamelCase      						=    OwlViTImageProcessor(
     size=config.vision_config.image_size   ,					crop_size=config.vision_config.image_size  )
 # Initialize tokenizer
 lowerCamelCase      						=    CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32'   ,					pad_token='!'   ,					model_max_length=1_6  )

 # Initialize processor
 lowerCamelCase      						=    OwlViTProcessor(image_processor=UpperCamelCase_   ,					tokenizer=UpperCamelCase_  )
 image_processor.save_pretrained(repo.local_dir  )
 processor.save_pretrained(repo.local_dir  )

 repo.git_add()
 repo.git_commit('Upload model and processor'  )
 repo.git_push()


if __name__ == "__main__":
     _lowerCAmelCase :	Union[str, Any]     				=    argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--owlvit_version',
         default=None,
         type=str,
         required=True,
         help='OWL-ViT model name [clip_b16, clip_b32, clip_l14].',
     )
     parser.add_argument(
         '--owlvit_checkpoint', default=None, type=str, required=True, help='Path to flax model checkpoint.'
     )
     parser.add_argument('--hf_config', default=None, type=str, required=True, help='Path to HF model config.')
     parser.add_argument(
         '--pytorch_dump_folder_path', default='hf_model', type=str, help='Path to the output PyTorch model.'
     )
     _lowerCAmelCase :	Any     				=    parser.parse_args()

     # Initialize PyToch clip model
     _lowerCAmelCase :	Tuple     				=    args.owlvit_version
     if model_name == "clip_b16":
          _lowerCAmelCase :	str     				=    CONFIGS['vit_b16']
     elif model_name == "clip_b32":
          _lowerCAmelCase :	Any     				=    CONFIGS['vit_b32']
     elif model_name == "clip_l14":
          _lowerCAmelCase :	Dict     				=    CONFIGS['vit_l14']

     # Load from checkpoint and convert params to float-32
     _lowerCAmelCase :	Union[str, Any]     				=    checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)['optimizer']['target']
     _lowerCAmelCase :	List[str]     				=    jax.tree_util.tree_map(lambda x: x.astype(jnp.floataa) if x.dtype == jnp.bfloataa else x, variables)
     del variables

     # Convert CLIP backbone
     _lowerCAmelCase				,   _lowerCAmelCase				,   _lowerCAmelCase :	str     				=    convert_clip_backbone(flax_params, torch_config)

     convert_owlvit_checkpoint(clip_pt, flax_params, attn_params, args.pytorch_dump_folder_path, args.hf_config)







import collections
import inspect
import unittest

from transformers import NatConfig
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 NatBackbone, NatForImageClassification, NatModel
     from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor


class     lowerCAmelCase				:


       '''simple docstring'''







       def __init__(	self					:  Union[str, Any]						,	__snake_case					:  Optional[Any]						,	__snake_case					:  Tuple=13						,	__snake_case					:  Optional[Any]=64						,	__snake_case					:  str=4						,	__snake_case					:  List[Any]=3						,	__snake_case					:  Dict=16						,	__snake_case					:  str=[1, 2, 1]						,	__snake_case					:  Any=[2, 4, 8]						,	__snake_case					:  Union[str, Any]=3						,	__snake_case					:  List[str]=2.0						,	__snake_case					:  Optional[Any]=True						,	__snake_case					:  str=0.0						,	__snake_case					:  Optional[Any]=0.0						,	__snake_case					:  Any=0.1						,	__snake_case					:  List[Any]="gelu"						,	__snake_case					:  str=True						,	__snake_case					:  str=0.02						,	__snake_case					:  Optional[Any]=1e-5						,	__snake_case					:  int=True						,	__snake_case					:  int=None						,	__snake_case					:  str=True						,	__snake_case					:  Any=10						,	__snake_case					:  List[str]=["stage1", "stage2"]						,	__snake_case					:  int=[1, 2]						,	)   -> int:







        '''simple docstring'''

        lowerCamelCase      						=    parent
        lowerCamelCase      						=    batch_size
        lowerCamelCase      						=    image_size
        lowerCamelCase      						=    patch_size
        lowerCamelCase      						=    num_channels
        lowerCamelCase      						=    embed_dim
        lowerCamelCase      						=    depths
        lowerCamelCase      						=    num_heads
        lowerCamelCase      						=    kernel_size
        lowerCamelCase      						=    mlp_ratio
        lowerCamelCase      						=    qkv_bias
        lowerCamelCase      						=    hidden_dropout_prob
        lowerCamelCase      						=    attention_probs_dropout_prob
        lowerCamelCase      						=    drop_path_rate
        lowerCamelCase      						=    hidden_act
        lowerCamelCase      						=    patch_norm
        lowerCamelCase      						=    layer_norm_eps
        lowerCamelCase      						=    initializer_range
        lowerCamelCase      						=    is_training
        lowerCamelCase      						=    scope
        lowerCamelCase      						=    use_labels
        lowerCamelCase      						=    num_labels
        lowerCamelCase      						=    out_features
        lowerCamelCase      						=    out_indices

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







        '''simple docstring'''

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

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

        lowerCamelCase      						=    self.get_config()

        return config, pixel_values, labels

       def   lowerCamelCase__						(	self					:  Tuple     )   -> int:







        '''simple docstring'''

        return NatConfig(
            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						,	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					:  str						,	__snake_case					:  Optional[Any]						,	__snake_case					:  List[Any]						,	__snake_case					:  List[str]     )   -> int:







        '''simple docstring'''

        lowerCamelCase      						=    NatModel(config=__snake_case     )
        model.to(__snake_case     )
        model.eval()
        lowerCamelCase      						=    model(__snake_case     )

        lowerCamelCase      						=    lowerCamelCase      						=    (config.image_size // config.patch_size) // (2 ** (len(config.depths     ) - 1))
        lowerCamelCase      						=    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					:  Tuple						,	__snake_case					:  List[Any]						,	__snake_case					:  Optional[int]						,	__snake_case					:  Union[str, Any]     )   -> str:







        '''simple docstring'''

        lowerCamelCase      						=    NatForImageClassification(__snake_case     )
        model.to(__snake_case     )
        model.eval()
        lowerCamelCase      						=    model(__snake_case						,	labels=__snake_case     )
        self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_labels)     )

        # test greyscale images
        lowerCamelCase      						=    1
        lowerCamelCase      						=    NatForImageClassification(__snake_case     )
        model.to(__snake_case     )
        model.eval()

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

       def   lowerCamelCase__						(	self					:  Union[str, Any]						,	__snake_case					:  Any						,	__snake_case					:  Any						,	__snake_case					:  Dict     )   -> List[Any]:







        '''simple docstring'''

        lowerCamelCase      						=    NatBackbone(config=__snake_case     )
        model.to(__snake_case     )
        model.eval()
        lowerCamelCase      						=    model(__snake_case     )

        # 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
        lowerCamelCase      						=    None
        lowerCamelCase      						=    NatBackbone(config=__snake_case     )
        model.to(__snake_case     )
        model.eval()
        lowerCamelCase      						=    model(__snake_case     )

        # 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					:  Tuple     )   -> Union[str, Any]:







        '''simple docstring'''

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


@require_natten
@require_torch
class     lowerCAmelCase				(    __UpperCamelCase   ,				__UpperCamelCase   ,				unittest.TestCase				):


       '''simple docstring'''







       snake_case    							=	(
           (
               NatModel,
               NatForImageClassification,
               NatBackbone,
           )
           if is_torch_available()
           else ()
       )
       snake_case    							=	(
           {'feature-extraction': NatModel, 'image-classification': NatForImageClassification}
           if is_torch_available()
           else {}
       )
       snake_case    							=	False

       snake_case    							=	False
       snake_case    							=	False
       snake_case    							=	False
       snake_case    							=	False

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







        '''simple docstring'''

        lowerCamelCase      						=    NatModelTester(self     )
        lowerCamelCase      						=    ConfigTester(self						,	config_class=__snake_case						,	embed_dim=37     )

       def   lowerCamelCase__						(	self					:  Dict     )   -> 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   lowerCamelCase__						(	self					:  Optional[Any]     )   -> Union[str, Any]:







        '''simple docstring'''

        return

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







        '''simple docstring'''

        lowerCamelCase      						=    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__snake_case     )

       def   lowerCamelCase__						(	self					:  Optional[int]     )   -> Optional[int]:







        '''simple docstring'''

        lowerCamelCase      						=    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*__snake_case     )

       def   lowerCamelCase__						(	self					:  List[str]     )   -> Dict:







        '''simple docstring'''

        lowerCamelCase      						=    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_backbone(*__snake_case     )

       @unittest.skip(reason='Nat does not use inputs_embeds'     )
       def   lowerCamelCase__						(	self					:  Tuple     )   -> Optional[int]:







        '''simple docstring'''

        pass

       @unittest.skip(reason='Nat does not use feedforward chunking'     )
       def   lowerCamelCase__						(	self					:  Optional[Any]     )   -> List[str]:







        '''simple docstring'''

        pass

       def   lowerCamelCase__						(	self					:  Tuple     )   -> str:







        '''simple docstring'''

        lowerCamelCase						,	lowerCamelCase      						=    self.model_tester.prepare_config_and_inputs_for_common()

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

       def   lowerCamelCase__						(	self					:  Any     )   -> List[Any]:







        '''simple docstring'''

        lowerCamelCase						,	lowerCamelCase      						=    self.model_tester.prepare_config_and_inputs_for_common()

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

         lowerCamelCase      						=    ['pixel_values']
         self.assertListEqual(arg_names[:1]						,	__snake_case     )

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







        '''simple docstring'''

        self.skipTest('Nat\'s attention operation is handled entirely by NATTEN.'     )

       def   lowerCamelCase__						(	self					:  str						,	__snake_case					:  List[str]						,	__snake_case					:  Optional[Any]						,	__snake_case					:  Tuple						,	__snake_case					:  Union[str, Any]     )   -> int:







        '''simple docstring'''

        lowerCamelCase      						=    model_class(__snake_case     )
        model.to(__snake_case     )
        model.eval()

        with torch.no_grad():
         lowerCamelCase      						=    model(**self._prepare_for_class(__snake_case						,	__snake_case     )     )

        lowerCamelCase      						=    outputs.hidden_states

        lowerCamelCase      						=    getattr(
            self.model_tester						,	'expected_num_hidden_layers'						,	len(self.model_tester.depths     ) + 1     )
        self.assertEqual(len(__snake_case     )						,	__snake_case     )

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

        lowerCamelCase      						=    image_size[0] // patch_size[0]
        lowerCamelCase      						=    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__ != "NatBackbone":
         lowerCamelCase      						=    outputs.reshaped_hidden_states
         self.assertEqual(len(__snake_case     )						,	__snake_case     )

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

       def   lowerCamelCase__						(	self					:  Dict     )   -> Tuple:







        '''simple docstring'''

        lowerCamelCase						,	lowerCamelCase      						=    self.model_tester.prepare_config_and_inputs_for_common()

        lowerCamelCase      						=    (
            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:
         lowerCamelCase      						=    True
         self.check_hidden_states_output(__snake_case						,	__snake_case						,	__snake_case						,	__snake_case     )

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

         self.check_hidden_states_output(__snake_case						,	__snake_case						,	__snake_case						,	__snake_case     )

       @slow
       def   lowerCamelCase__						(	self					:  Any     )   -> Optional[int]:







        '''simple docstring'''

        for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowerCamelCase      						=    NatModel.from_pretrained(__snake_case     )
         self.assertIsNotNone(__snake_case     )





       def   lowerCamelCase__						(	self					:  Any     )   -> Any:







        '''simple docstring'''

        lowerCamelCase						,	lowerCamelCase      						=    self.model_tester.prepare_config_and_inputs_for_common()

        lowerCamelCase      						=    _config_zero_init(__snake_case     )
        for model_class in self.all_model_classes:
         lowerCamelCase      						=    model_class(config=__snake_case     )
         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				):


       '''simple docstring'''







       @cached_property
       def   lowerCamelCase__						(	self					:  Union[str, Any]     )   -> Dict:







        '''simple docstring'''

        return AutoImageProcessor.from_pretrained('shi-labs/nat-mini-in1k-224'     ) if is_vision_available() else None





       @slow
       def   lowerCamelCase__						(	self					:  List[str]     )   -> Tuple:







        '''simple docstring'''

        lowerCamelCase      						=    NatForImageClassification.from_pretrained('shi-labs/nat-mini-in1k-224'     ).to(__snake_case     )
        lowerCamelCase      						=    self.default_image_processor

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

        # forward pass
        with torch.no_grad():
         lowerCamelCase      						=    model(**__snake_case     )

        # verify the logits
        lowerCamelCase      						=    torch.Size((1, 1000)     )
        self.assertEqual(outputs.logits.shape						,	__snake_case     )
        lowerCamelCase      						=    torch.tensor([0.3805, -0.8676, -0.3912]     ).to(__snake_case     )
        self.assertTrue(torch.allclose(outputs.logits[0, :3]						,	__snake_case						,	atol=1e-4     )     )


@require_torch
@require_natten
class     lowerCAmelCase				(    unittest.TestCase   ,				__UpperCamelCase				):


       '''simple docstring'''







       snake_case    							=	(NatBackbone,) if is_torch_available() else ()
       snake_case    							=	NatConfig

       def   lowerCamelCase__						(	self					:  Union[str, Any]     )   -> Dict:







        '''simple docstring'''

        lowerCamelCase      						=    NatModelTester(self     )




