


"""simple docstring"""

import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

from transformers import (
    BitConfig,
    ViTHybridConfig,
    ViTHybridForImageClassification,
    ViTHybridImageProcessor,
    ViTHybridModel,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase__					=logging.get_logger(__name__)


def lowerCAmelCase_   (					UpperCamelCase__						:			Tuple					, UpperCamelCase__						:			int=False			):






 """simple docstring"""







 __lowercase						=						[]

 # fmt: off
 # stem:
 rename_keys.append(("""cls_token""", """vit.embeddings.cls_token""")			)
 rename_keys.append(("""pos_embed""", """vit.embeddings.position_embeddings""")			)

 rename_keys.append(("""patch_embed.proj.weight""", """vit.embeddings.patch_embeddings.projection.weight""")			)
 rename_keys.append(("""patch_embed.proj.bias""", """vit.embeddings.patch_embeddings.projection.bias""")			)

 # backbone
 rename_keys.append(("""patch_embed.backbone.stem.conv.weight""", """vit.embeddings.patch_embeddings.backbone.bit.embedder.convolution.weight""")			)
 rename_keys.append(("""patch_embed.backbone.stem.norm.weight""", """vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.weight""")			)
 rename_keys.append(("""patch_embed.backbone.stem.norm.bias""", """vit.embeddings.patch_embeddings.backbone.bit.embedder.norm.bias""")			)

 for stage_idx in range(len(config.backbone_config.depths			)			):
  for layer_idx in range(config.backbone_config.depths[stage_idx]			):
   rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv1.weight''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv1.weight''')			)
   rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.weight''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.weight''')			)
   rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm1.bias''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm1.bias''')			)
   rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv2.weight''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv2.weight''')			)
   rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.weight''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.weight''')			)
   rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm2.bias''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm2.bias''')			)
   rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.conv3.weight''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.conv3.weight''')			)
   rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.weight''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.weight''')			)
   rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.{layer_idx}.norm3.bias''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.{layer_idx}.norm3.bias''')			)

  rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.conv.weight''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.conv.weight''')			)
  rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.weight''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.weight''')			)
  rename_keys.append((f'''patch_embed.backbone.stages.{stage_idx}.blocks.0.downsample.norm.bias''', f'''vit.embeddings.patch_embeddings.backbone.bit.encoder.stages.{stage_idx}.layers.0.downsample.norm.bias''')			)

 # transformer encoder
 for i in range(config.num_hidden_layers			):
  # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
  rename_keys.append((f'''blocks.{i}.norm1.weight''', f'''vit.encoder.layer.{i}.layernorm_before.weight''')			)
  rename_keys.append((f'''blocks.{i}.norm1.bias''', f'''vit.encoder.layer.{i}.layernorm_before.bias''')			)
  rename_keys.append((f'''blocks.{i}.attn.proj.weight''', f'''vit.encoder.layer.{i}.attention.output.dense.weight''')			)
  rename_keys.append((f'''blocks.{i}.attn.proj.bias''', f'''vit.encoder.layer.{i}.attention.output.dense.bias''')			)
  rename_keys.append((f'''blocks.{i}.norm2.weight''', f'''vit.encoder.layer.{i}.layernorm_after.weight''')			)
  rename_keys.append((f'''blocks.{i}.norm2.bias''', f'''vit.encoder.layer.{i}.layernorm_after.bias''')			)
  rename_keys.append((f'''blocks.{i}.mlp.fc1.weight''', f'''vit.encoder.layer.{i}.intermediate.dense.weight''')			)
  rename_keys.append((f'''blocks.{i}.mlp.fc1.bias''', f'''vit.encoder.layer.{i}.intermediate.dense.bias''')			)
  rename_keys.append((f'''blocks.{i}.mlp.fc2.weight''', f'''vit.encoder.layer.{i}.output.dense.weight''')			)
  rename_keys.append((f'''blocks.{i}.mlp.fc2.bias''', f'''vit.encoder.layer.{i}.output.dense.bias''')			)

 if base_model:
  # layernorm + pooler
  rename_keys.extend(
      [
          ("""norm.weight""", """layernorm.weight"""),
          ("""norm.bias""", """layernorm.bias"""),
          ("""pre_logits.fc.weight""", """pooler.dense.weight"""),
          ("""pre_logits.fc.bias""", """pooler.dense.bias"""),
      ]			)

  # if just the base model, we should remove "vit" from all keys that start with "vit"
  __lowercase						=						[(pair[0], pair[1][4:]) if pair[1].startswith("""vit"""			) else pair for pair in rename_keys]
 else:
  # layernorm + classification head
  rename_keys.extend(
      [
          ("""norm.weight""", """vit.layernorm.weight"""),
          ("""norm.bias""", """vit.layernorm.bias"""),
          ("""head.weight""", """classifier.weight"""),
          ("""head.bias""", """classifier.bias"""),
      ]			)
 # fmt: on

 return rename_keys


def lowerCAmelCase_   (					UpperCamelCase__						:			Optional[Any]					, UpperCamelCase__						:			Tuple					, UpperCamelCase__						:			Optional[Any]=False			):






 """simple docstring"""







 for i in range(config.num_hidden_layers			):
  if base_model:
   __lowercase						=						""""""
  else:
   __lowercase						=						"""vit."""
  # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
  __lowercase						=						state_dict.pop(f'''blocks.{i}.attn.qkv.weight'''			)
  __lowercase						=						state_dict.pop(f'''blocks.{i}.attn.qkv.bias'''			)
  # next, add query, keys and values (in that order) to the state dict
  __lowercase						=						in_proj_weight[
      : config.hidden_size, :
  ]
  __lowercase						=						in_proj_bias[: config.hidden_size]
  __lowercase						=						in_proj_weight[
      config.hidden_size : config.hidden_size * 2, :
  ]
  __lowercase						=						in_proj_bias[
      config.hidden_size : config.hidden_size * 2
  ]
  __lowercase						=						in_proj_weight[
      -config.hidden_size :, :
  ]
  __lowercase						=						in_proj_bias[-config.hidden_size :]


def lowerCAmelCase_   (					UpperCamelCase__						:			Dict			):






 """simple docstring"""







 __lowercase						=						["""head.weight""", """head.bias"""]
 for k in ignore_keys:
  state_dict.pop(UpperCamelCase__					, UpperCamelCase__			)


def lowerCAmelCase_   (					UpperCamelCase__						:			Tuple					, UpperCamelCase__						:			Optional[int]					, UpperCamelCase__						:			Any			):






 """simple docstring"""







 __lowercase						=						dct.pop(UpperCamelCase__			)
 __lowercase						=						val


def lowerCAmelCase_   (					):






 """simple docstring"""







 __lowercase						=						"""http://images.cocodataset.org/val2017/000000039769.jpg"""
 __lowercase						=						Image.open(requests.get(UpperCamelCase__					, stream=UpperCamelCase__			).raw			)
 return im


@torch.no_grad()
def lowerCAmelCase_   (					UpperCamelCase__						:			List[Any]					, UpperCamelCase__						:			Union[str, Any]					, UpperCamelCase__						:			Tuple=False			):






 """simple docstring"""







 __lowercase						=						BitConfig(
     global_padding="""same"""					, layer_type="""bottleneck"""					, depths=(3, 4, 9)					, out_features=["""stage3"""]					, embedding_dynamic_padding=UpperCamelCase__					, )
 __lowercase						=						ViTHybridConfig(backbone_config=UpperCamelCase__					, image_size=384					, num_labels=1000			)
 __lowercase						=						False

 # load original model from timm
 __lowercase						=						timm.create_model(UpperCamelCase__					, pretrained=UpperCamelCase__			)
 timm_model.eval()

 # load state_dict of original model, remove and rename some keys
 __lowercase						=						timm_model.state_dict()
 if base_model:
  remove_classification_head_(UpperCamelCase__			)
 __lowercase						=						create_rename_keys(UpperCamelCase__					, UpperCamelCase__			)
 for src, dest in rename_keys:
  rename_key(UpperCamelCase__					, UpperCamelCase__					, UpperCamelCase__			)
 read_in_q_k_v(UpperCamelCase__					, UpperCamelCase__					, UpperCamelCase__			)

 __lowercase						=						"""huggingface/label-files"""
 __lowercase						=						"""imagenet-1k-id2label.json"""
 __lowercase						=						json.load(open(hf_hub_download(UpperCamelCase__					, UpperCamelCase__					, repo_type="""dataset"""			)					, """r"""			)			)
 __lowercase						=						{int(UpperCamelCase__			): v for k, v in idalabel.items()}
 __lowercase						=						idalabel
 __lowercase						=						{v: k for k, v in idalabel.items()}

 # load HuggingFace model
 if vit_name[-5:] == "in21k":
  __lowercase						=						ViTHybridModel(UpperCamelCase__			).eval()
 else:
  __lowercase						=						ViTHybridForImageClassification(UpperCamelCase__			).eval()
 model.load_state_dict(UpperCamelCase__			)

 # create image processor
 __lowercase						=						create_transform(**resolve_data_config({}					, model=UpperCamelCase__			)			)
 __lowercase						=						transform.transforms

 __lowercase						=						{
     """bilinear""": PILImageResampling.BILINEAR,
     """bicubic""": PILImageResampling.BICUBIC,
     """nearest""": PILImageResampling.NEAREST,
 }

 __lowercase						=						ViTHybridImageProcessor(
     do_resize=UpperCamelCase__					, size={"""shortest_edge""": timm_transforms[0].size}					, resample=pillow_resamplings[timm_transforms[0].interpolation.value]					, do_center_crop=UpperCamelCase__					, crop_size={"""height""": timm_transforms[1].size[0], """width""": timm_transforms[1].size[1]}					, do_normalize=UpperCamelCase__					, image_mean=timm_transforms[-1].mean.tolist()					, image_std=timm_transforms[-1].std.tolist()					, )

 __lowercase						=						prepare_img()
 __lowercase						=						transform(UpperCamelCase__			).unsqueeze(0			)
 __lowercase						=						processor(UpperCamelCase__					, return_tensors="""pt"""			).pixel_values

 # verify pixel values
 assert torch.allclose(UpperCamelCase__					, UpperCamelCase__			)

 # verify logits
 with torch.no_grad():
  __lowercase						=						model(UpperCamelCase__			)
  __lowercase						=						outputs.logits

 print("""Predicted class:"""					, logits.argmax(-1			).item()			)
 if base_model:
  __lowercase						=						timm_model.forward_features(UpperCamelCase__			)
  assert timm_pooled_output.shape == outputs.pooler_output.shape
  assert torch.allclose(UpperCamelCase__					, outputs.pooler_output					, atol=1E-3			)
 else:
  __lowercase						=						timm_model(UpperCamelCase__			)
  assert timm_logits.shape == outputs.logits.shape
  assert torch.allclose(UpperCamelCase__					, outputs.logits					, atol=1E-3			)
 print("""Looks ok!"""			)

 if pytorch_dump_folder_path is not None:
  Path(UpperCamelCase__			).mkdir(exist_ok=UpperCamelCase__			)
  print(f'''Saving model {vit_name} to {pytorch_dump_folder_path}'''			)
  model.save_pretrained(UpperCamelCase__			)
  print(f'''Saving processor to {pytorch_dump_folder_path}'''			)
  processor.save_pretrained(UpperCamelCase__			)

 if push_to_hub:
  print(f'''Pushing model and processor to the hub {vit_name}'''			)
  model.push_to_hub(f'''ybelkada/{vit_name}'''			)
  processor.push_to_hub(f'''ybelkada/{vit_name}'''			)


if __name__ == "__main__":
       UpperCAmelCase__					=argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           "--vit_name",
           default="vit_base_r50_s16_384",
           type=str,
           help="Name of the hybrid ViT timm model you'd like to convert.",
       )
       parser.add_argument(
           "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
       )
       parser.add_argument(
           "--push_to_hub", action="store_true", help="Whether to upload the model to the HuggingFace hub."
       )

       UpperCAmelCase__					=parser.parse_args()
       convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path, args.push_to_hub)




"""simple docstring"""

import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
       import jax.numpy as jnp

if is_tf_available():
       import tensorflow as tf

if is_torch_available():
       import torch





class  lowerCamelCase__    (		unittest.TestCase     ):
      def SCREAMING_SNAKE_CASE_						(		self						:						Any				):






       '''simple docstring'''


       __lowercase						=						{
           """task_specific_params""": {
               """summarization""": {"""length_penalty""": 1.0, """max_length""": 1_2_8, """min_length""": 1_2, """num_beams""": 4},
               """summarization_cnn""": {"""length_penalty""": 2.0, """max_length""": 1_4_2, """min_length""": 5_6, """num_beams""": 4},
               """summarization_xsum""": {"""length_penalty""": 1.0, """max_length""": 6_2, """min_length""": 1_1, """num_beams""": 6},
           }
       }
       __lowercase						=						{
           """task_specific_params.summarization.length_penalty""": 1.0,
           """task_specific_params.summarization.max_length""": 1_2_8,
           """task_specific_params.summarization.min_length""": 1_2,
           """task_specific_params.summarization.num_beams""": 4,
           """task_specific_params.summarization_cnn.length_penalty""": 2.0,
           """task_specific_params.summarization_cnn.max_length""": 1_4_2,
           """task_specific_params.summarization_cnn.min_length""": 5_6,
           """task_specific_params.summarization_cnn.num_beams""": 4,
           """task_specific_params.summarization_xsum.length_penalty""": 1.0,
           """task_specific_params.summarization_xsum.max_length""": 6_2,
           """task_specific_params.summarization_xsum.min_length""": 1_1,
           """task_specific_params.summarization_xsum.num_beams""": 6,
       }

       self.assertEqual(flatten_dict(A_				)   ,							A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       self.assertTrue(np.allclose(transpose(A_				)   ,							x.transpose()				)				)

       __lowercase						=						np.random.randn(3   ,							4   ,							5				)
       self.assertTrue(np.allclose(transpose(A_   ,							axes=(1, 2, 0)				)   ,							x.transpose((1, 2, 0)				)				)				)
      @require_torch
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       __lowercase						=						torch.tensor(A_				)
       self.assertTrue(np.allclose(transpose(A_				)   ,							transpose(A_				).numpy()				)				)

       __lowercase						=						np.random.randn(3   ,							4   ,							5				)
       __lowercase						=						torch.tensor(A_				)
       self.assertTrue(np.allclose(transpose(A_   ,							axes=(1, 2, 0)				)   ,							transpose(A_   ,							axes=(1, 2, 0)				).numpy()				)				)
      @require_tf
      def SCREAMING_SNAKE_CASE_						(		self						:						int				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       __lowercase						=						tf.constant(A_				)
       self.assertTrue(np.allclose(transpose(A_				)   ,							transpose(A_				).numpy()				)				)

       __lowercase						=						np.random.randn(3   ,							4   ,							5				)
       __lowercase						=						tf.constant(A_				)
       self.assertTrue(np.allclose(transpose(A_   ,							axes=(1, 2, 0)				)   ,							transpose(A_   ,							axes=(1, 2, 0)				).numpy()				)				)
      @require_flax
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       __lowercase						=						jnp.array(A_				)
       self.assertTrue(np.allclose(transpose(A_				)   ,							np.asarray(transpose(A_				)				)				)				)

       __lowercase						=						np.random.randn(3   ,							4   ,							5				)
       __lowercase						=						jnp.array(A_				)
       self.assertTrue(np.allclose(transpose(A_   ,							axes=(1, 2, 0)				)   ,							np.asarray(transpose(A_   ,							axes=(1, 2, 0)				)				)				)				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       self.assertTrue(np.allclose(reshape(A_   ,							(4, 3)				)   ,							np.reshape(A_   ,							(4, 3)				)				)				)

       __lowercase						=						np.random.randn(3   ,							4   ,							5				)
       self.assertTrue(np.allclose(reshape(A_   ,							(1_2, 5)				)   ,							np.reshape(A_   ,							(1_2, 5)				)				)				)
      @require_torch
      def SCREAMING_SNAKE_CASE_						(		self						:						Dict				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       __lowercase						=						torch.tensor(A_				)
       self.assertTrue(np.allclose(reshape(A_   ,							(4, 3)				)   ,							reshape(A_   ,							(4, 3)				).numpy()				)				)

       __lowercase						=						np.random.randn(3   ,							4   ,							5				)
       __lowercase						=						torch.tensor(A_				)
       self.assertTrue(np.allclose(reshape(A_   ,							(1_2, 5)				)   ,							reshape(A_   ,							(1_2, 5)				).numpy()				)				)
      @require_tf
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       __lowercase						=						tf.constant(A_				)
       self.assertTrue(np.allclose(reshape(A_   ,							(4, 3)				)   ,							reshape(A_   ,							(4, 3)				).numpy()				)				)

       __lowercase						=						np.random.randn(3   ,							4   ,							5				)
       __lowercase						=						tf.constant(A_				)
       self.assertTrue(np.allclose(reshape(A_   ,							(1_2, 5)				)   ,							reshape(A_   ,							(1_2, 5)				).numpy()				)				)
      @require_flax
      def SCREAMING_SNAKE_CASE_						(		self						:						str				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       __lowercase						=						jnp.array(A_				)
       self.assertTrue(np.allclose(reshape(A_   ,							(4, 3)				)   ,							np.asarray(reshape(A_   ,							(4, 3)				)				)				)				)

       __lowercase						=						np.random.randn(3   ,							4   ,							5				)
       __lowercase						=						jnp.array(A_				)
       self.assertTrue(np.allclose(reshape(A_   ,							(1_2, 5)				)   ,							np.asarray(reshape(A_   ,							(1_2, 5)				)				)				)				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Dict				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(1   ,							3   ,							4				)
       self.assertTrue(np.allclose(squeeze(A_				)   ,							np.squeeze(A_				)				)				)

       __lowercase						=						np.random.randn(1   ,							4   ,							1   ,							5				)
       self.assertTrue(np.allclose(squeeze(A_   ,							axis=2				)   ,							np.squeeze(A_   ,							axis=2				)				)				)
      @require_torch
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(1   ,							3   ,							4				)
       __lowercase						=						torch.tensor(A_				)
       self.assertTrue(np.allclose(squeeze(A_				)   ,							squeeze(A_				).numpy()				)				)

       __lowercase						=						np.random.randn(1   ,							4   ,							1   ,							5				)
       __lowercase						=						torch.tensor(A_				)
       self.assertTrue(np.allclose(squeeze(A_   ,							axis=2				)   ,							squeeze(A_   ,							axis=2				).numpy()				)				)
      @require_tf
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(1   ,							3   ,							4				)
       __lowercase						=						tf.constant(A_				)
       self.assertTrue(np.allclose(squeeze(A_				)   ,							squeeze(A_				).numpy()				)				)

       __lowercase						=						np.random.randn(1   ,							4   ,							1   ,							5				)
       __lowercase						=						tf.constant(A_				)
       self.assertTrue(np.allclose(squeeze(A_   ,							axis=2				)   ,							squeeze(A_   ,							axis=2				).numpy()				)				)
      @require_flax
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(1   ,							3   ,							4				)
       __lowercase						=						jnp.array(A_				)
       self.assertTrue(np.allclose(squeeze(A_				)   ,							np.asarray(squeeze(A_				)				)				)				)

       __lowercase						=						np.random.randn(1   ,							4   ,							1   ,							5				)
       __lowercase						=						jnp.array(A_				)
       self.assertTrue(np.allclose(squeeze(A_   ,							axis=2				)   ,							np.asarray(squeeze(A_   ,							axis=2				)				)				)				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       self.assertTrue(np.allclose(expand_dims(A_   ,							axis=1				)   ,							np.expand_dims(A_   ,							axis=1				)				)				)
      @require_torch
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       __lowercase						=						torch.tensor(A_				)
       self.assertTrue(np.allclose(expand_dims(A_   ,							axis=1				)   ,							expand_dims(A_   ,							axis=1				).numpy()				)				)
      @require_tf
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       __lowercase						=						tf.constant(A_				)
       self.assertTrue(np.allclose(expand_dims(A_   ,							axis=1				)   ,							expand_dims(A_   ,							axis=1				).numpy()				)				)






      @require_flax
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]				):






       '''simple docstring'''


       __lowercase						=						np.random.randn(3   ,							4				)
       __lowercase						=						jnp.array(A_				)
       self.assertTrue(np.allclose(expand_dims(A_   ,							axis=1				)   ,							np.asarray(expand_dims(A_   ,							axis=1				)				)				)				)
