
import argparse
import json
import os

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
a__				:      List[str]    		=     logging.get_logger(__name__)







def 				UpperCAmelCase_(    a__   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = ConvNextVaConfig()

    if "atto" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				Tuple				      = [2, 2, 6, 2]
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = [40, 80, 160, 320]
    if "femto" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				List[str]				      = [2, 2, 6, 2]
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = [48, 96, 192, 384]
    if "pico" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				List[Any]				      = [2, 2, 6, 2]
        SCREAMING_SNAKE_CASE				:				int				      = [64, 128, 256, 512]
    if "nano" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				Tuple				      = [2, 2, 8, 2]
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = [80, 160, 320, 640]
    if "tiny" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				str				      = [3, 3, 9, 3]
        SCREAMING_SNAKE_CASE				:				int				      = [96, 192, 384, 768]
    if "base" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = [3, 3, 27, 3]
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = [128, 256, 512, 1_024]
    if "large" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				int				      = [3, 3, 27, 3]
        SCREAMING_SNAKE_CASE				:				Any				      = [192, 384, 768, 1_536]
    if "huge" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = [3, 3, 27, 3]
        SCREAMING_SNAKE_CASE				:				List[str]				      = [352, 704, 1_408, 2_816]

    SCREAMING_SNAKE_CASE				:				int				      = 1_000
    SCREAMING_SNAKE_CASE				:				List[Any]				      = '''imagenet-1k-id2label.json'''
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = (1, 1_000)

    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = '''huggingface/label-files'''
    SCREAMING_SNAKE_CASE				:				Tuple				      = num_labels
    SCREAMING_SNAKE_CASE				:				int				      = json.load(open(hf_hub_download(a__ ,						a__ ,						repo_type='''dataset'''   ) ,						'''r'''   )   )
    SCREAMING_SNAKE_CASE				:				str				      = {int(a__   ): v for k, v in idalabel.items()}

    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = idalabel
    SCREAMING_SNAKE_CASE				:				str				      = {v: k for k, v in idalabel.items()}
    SCREAMING_SNAKE_CASE				:				Any				      = hidden_sizes
    SCREAMING_SNAKE_CASE				:				List[Any]				      = depths

    return config, expected_shape







def 				UpperCAmelCase_(    a__   ):



    """simple docstring"""



    if "downsample_layers.0.0" in name:
        SCREAMING_SNAKE_CASE				:				str				      = name.replace('''downsample_layers.0.0''' ,						'''embeddings.patch_embeddings'''   )
    if "downsample_layers.0.1" in name:
        SCREAMING_SNAKE_CASE				:				int				      = name.replace('''downsample_layers.0.1''' ,						'''embeddings.norm'''   )  # we rename to layernorm later on
    if "downsample_layers.1.0" in name:
        SCREAMING_SNAKE_CASE				:				str				      = name.replace('''downsample_layers.1.0''' ,						'''stages.1.downsampling_layer.0'''   )
    if "downsample_layers.1.1" in name:
        SCREAMING_SNAKE_CASE				:				int				      = name.replace('''downsample_layers.1.1''' ,						'''stages.1.downsampling_layer.1'''   )
    if "downsample_layers.2.0" in name:
        SCREAMING_SNAKE_CASE				:				Tuple				      = name.replace('''downsample_layers.2.0''' ,						'''stages.2.downsampling_layer.0'''   )
    if "downsample_layers.2.1" in name:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = name.replace('''downsample_layers.2.1''' ,						'''stages.2.downsampling_layer.1'''   )
    if "downsample_layers.3.0" in name:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = name.replace('''downsample_layers.3.0''' ,						'''stages.3.downsampling_layer.0'''   )
    if "downsample_layers.3.1" in name:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = name.replace('''downsample_layers.3.1''' ,						'''stages.3.downsampling_layer.1'''   )
    if "stages" in name and "downsampling_layer" not in name:
        # stages.0.0. for instance should be renamed to stages.0.layers.0.
        SCREAMING_SNAKE_CASE				:				Dict				      = name[: len('''stages.0'''   )] + '''.layers''' + name[len('''stages.0'''   ) :]
    if "gamma" in name:
        SCREAMING_SNAKE_CASE				:				Any				      = name.replace('''gamma''' ,						'''weight'''   )
    if "beta" in name:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = name.replace('''beta''' ,						'''bias'''   )
    if "stages" in name:
        SCREAMING_SNAKE_CASE				:				List[Any]				      = name.replace('''stages''' ,						'''encoder.stages'''   )
    if "norm" in name:
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = name.replace('''norm''' ,						'''layernorm'''   )
    if "head" in name:
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = name.replace('''head''' ,						'''classifier'''   )

    return name







def 				UpperCAmelCase_(    ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Tuple				      = '''http://images.cocodataset.org/val2017/000000039769.jpg'''
    SCREAMING_SNAKE_CASE				:				str				      = Image.open(requests.get(a__ ,						stream=a__   ).raw   )
    return im







def 				UpperCAmelCase_(    a__   ):



    """simple docstring"""



    if "224" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				Any				      = 224
        SCREAMING_SNAKE_CASE				:				int				      = 224 / 256
    elif "384" in checkpoint_url:
        SCREAMING_SNAKE_CASE				:				int				      = 384
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = None
    else:
        SCREAMING_SNAKE_CASE				:				List[Any]				      = 512
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = None

    return ConvNextImageProcessor(
        size=a__ ,						crop_pct=a__ ,						image_mean=[0.485, 0.456, 0.406] ,						image_std=[0.229, 0.224, 0.225] ,						resample=PILImageResampling.BICUBIC ,						)







@torch.no_grad()
def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						a__   ):



    """simple docstring"""



    print('''Downloading original model from checkpoint...'''   )
    # define ConvNeXTV2 configuration based on URL
    SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				List[Any]				      = get_convnextva_config(a__   )
    # load original state_dict from URL
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.hub.load_state_dict_from_url(a__   )['''model''']

    print('''Converting model parameters...'''   )
    # rename keys
    for key in state_dict.copy().keys():
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = state_dict.pop(a__   )
        SCREAMING_SNAKE_CASE				:				Tuple				      = val
    # add prefix to all keys expect classifier head
    for key in state_dict.copy().keys():
        SCREAMING_SNAKE_CASE				:				List[Any]				      = state_dict.pop(a__   )
        if not key.startswith('''classifier'''   ):
            SCREAMING_SNAKE_CASE				:				Dict				      = '''convnextv2.''' + key
        SCREAMING_SNAKE_CASE				:				Dict				      = val

    # load HuggingFace model
    SCREAMING_SNAKE_CASE				:				Dict				      = ConvNextVaForImageClassification(a__   )
    model.load_state_dict(a__   )
    model.eval()

    # Check outputs on an image, prepared by ConvNextImageProcessor
    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = convert_preprocessor(a__   )
    SCREAMING_SNAKE_CASE				:				Any				      = preprocessor(images=prepare_img() ,						return_tensors='''pt'''   )
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = model(**a__   ).logits

    # note: the logits below were obtained without center cropping
    if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.tensor([-0.3_930, 0.1_747, -0.5_246, 0.4_177, 0.4_295]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				str				      = torch.tensor([-0.1_727, -0.5_341, -0.7_818, -0.4_745, -0.6_566]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				List[str]				      = torch.tensor([-0.0_333, 0.1_563, -0.9_137, 0.1_054, 0.0_381]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				int				      = torch.tensor([-0.1_744, -0.1_555, -0.0_713, 0.0_950, -0.1_431]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				int				      = torch.tensor([0.9_996, 0.1_966, -0.4_386, -0.3_472, 0.6_661]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = torch.tensor([-0.2_553, -0.6_708, -0.1_359, 0.2_518, -0.2_488]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				Dict				      = torch.tensor([-0.0_673, -0.5_627, -0.3_753, -0.2_722, 0.0_178]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				Any				      = torch.tensor([-0.6_377, -0.7_458, -0.2_150, 0.1_184, -0.0_597]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				Tuple				      = torch.tensor([1.0_799, 0.2_322, -0.8_860, 1.0_219, 0.6_231]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.tensor([0.3_766, 0.4_917, -1.1_426, 0.9_942, 0.6_024]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				str				      = torch.tensor([0.4_220, -0.6_919, -0.4_317, -0.2_881, -0.6_609]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
        SCREAMING_SNAKE_CASE				:				int				      = torch.tensor([0.1_082, -0.8_286, -0.5_095, 0.4_681, -0.8_085]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				int				      = torch.tensor([-0.2_419, -0.6_221, 0.2_176, -0.0_980, -0.7_527]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = torch.tensor([0.0_391, -0.4_371, 0.3_786, 0.1_251, -0.2_784]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = torch.tensor([-0.0_504, 0.5_636, -0.1_729, -0.6_507, -0.3_949]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = torch.tensor([0.3_560, 0.9_486, 0.3_149, -0.2_667, -0.5_138]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
        SCREAMING_SNAKE_CASE				:				Tuple				      = torch.tensor([-0.2_469, -0.4_550, -0.5_853, -0.0_810, 0.0_309]   )
    elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.tensor([-0.3_090, 0.0_802, -0.0_682, -0.1_979, -0.2_826]   )
    else:
        raise ValueError(F"""Unknown URL: {checkpoint_url}"""   )

    assert torch.allclose(logits[0, :5] ,						a__ ,						atol=1e-3   )
    assert logits.shape == expected_shape
    print('''Model outputs match the original results!'''   )

    if save_model:
        print('''Saving model to local...'''   )
        # Create folder to save model
        if not os.path.isdir(a__   ):
            os.mkdir(a__   )

        model.save_pretrained(a__   )
        preprocessor.save_pretrained(a__   )

    SCREAMING_SNAKE_CASE				:				Tuple				      = '''convnextv2'''
    if "atto" in checkpoint_url:
        model_name += "-atto"
    if "femto" in checkpoint_url:
        model_name += "-femto"
    if "pico" in checkpoint_url:
        model_name += "-pico"
    if "nano" in checkpoint_url:
        model_name += "-nano"
    elif "tiny" in checkpoint_url:
        model_name += "-tiny"
    elif "base" in checkpoint_url:
        model_name += "-base"
    elif "large" in checkpoint_url:
        model_name += "-large"
    elif "huge" in checkpoint_url:
        model_name += "-huge"
    if "22k" in checkpoint_url and "1k" not in checkpoint_url:
        model_name += "-22k"
    elif "22k" in checkpoint_url and "1k" in checkpoint_url:
        model_name += "-22k-1k"
    elif "1k" in checkpoint_url:
        model_name += "-1k"
    if "224" in checkpoint_url:
        model_name += "-224"
    elif "384" in checkpoint_url:
        model_name += "-384"
    elif "512" in checkpoint_url:
        model_name += "-512"

    if push_to_hub:
        print(F"""Pushing {model_name} to the hub..."""   )
        model.push_to_hub(a__   )
        preprocessor.push_to_hub(a__   )


if __name__ == "__main__":
      a__				:      Optional[Any]    		=     argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          '''--checkpoint_url''',
          default='''https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt''',
          type=str,
          help='''URL of the original ConvNeXTV2 checkpoint you\'d like to convert.''',
      )
      parser.add_argument(
          '''--pytorch_dump_folder_path''',
          default='''model''',
          type=str,
          help='''Path to the output PyTorch model directory.''',
      )
      parser.add_argument('''--save_model''', action='''store_true''', help='''Save model to local''')
      parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image preprocessor to the hub''')

      a__				:      int    		=     parser.parse_args()
      convert_convnextva_checkpoint(
          args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
      )


import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


a__				:      Tuple    		=     [nn.LayerNorm]

a__				:      Tuple    		=     logging.get_logger(__name__)

a__				:      int    		=     version.parse(version.parse(torch.__version__).base_version)

a__				:      Dict    		=     parsed_torch_version_base >= version.parse('''2.0''')
a__				:      Any    		=     parsed_torch_version_base >= version.parse('''1.12''')
a__				:      Union[str, Any]    		=     parsed_torch_version_base >= version.parse('''1.11''')
a__				:      Union[str, Any]    		=     parsed_torch_version_base < version.parse('''1.11''')







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						a__ ,						self   ):



    """simple docstring"""



    from torch import _softmax_backward_data

    if is_torch_less_than_1_11:
        return _softmax_backward_data(a__ ,						a__ ,						parent.dim ,						self   )
    else:
        return _softmax_backward_data(a__ ,						a__ ,						parent.dim ,						self.dtype   )







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ = 0   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = index.to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = layer.weight.index_select(a__ ,						a__   ).clone().detach()
    if layer.bias is not None:
        if dim == 1:
            SCREAMING_SNAKE_CASE				:				int				      = layer.bias.clone().detach()
        else:
            SCREAMING_SNAKE_CASE				:				List[Any]				      = layer.bias[index].clone().detach()
    SCREAMING_SNAKE_CASE				:				List[Any]				      = list(layer.weight.size()   )
    SCREAMING_SNAKE_CASE				:				List[Any]				      = len(a__   )
    SCREAMING_SNAKE_CASE				:				str				      = nn.Linear(new_size[1] ,						new_size[0] ,						bias=layer.bias is not None   ).to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = False
    new_layer.weight.copy_(W.contiguous()   )
    SCREAMING_SNAKE_CASE				:				Tuple				      = True
    if layer.bias is not None:
        SCREAMING_SNAKE_CASE				:				Dict				      = False
        new_layer.bias.copy_(b.contiguous()   )
        SCREAMING_SNAKE_CASE				:				Tuple				      = True
    return new_layer






class     a_       (			nn.Module							):



       """simple docstring"""






       def __init__(		self				,       _lowerCamelCase				,       _lowerCamelCase   )					->int:
           super().__init__()
           SCREAMING_SNAKE_CASE				:				List[str]				      = nf
           SCREAMING_SNAKE_CASE				:				List[Any]				      = nn.Parameter(torch.empty(_lowerCamelCase				,       _lowerCamelCase   )   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = nn.Parameter(torch.zeros(_lowerCamelCase   )   )
           nn.init.normal_(self.weight				,       std=0.0_2   )






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase   )					->Tuple:
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = x.size()[:-1] + (self.nf,)
           SCREAMING_SNAKE_CASE				:				int				      = torch.addmm(self.bias				,       x.view(-1				,       x.size(-1   )   )				,       self.weight   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = x.view(_lowerCamelCase   )
           return x







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ = 1   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				str				      = index.to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Any				      = layer.weight.index_select(a__ ,						a__   ).clone().detach()
    if dim == 0:
        SCREAMING_SNAKE_CASE				:				Tuple				      = layer.bias.clone().detach()
    else:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = layer.bias[index].clone().detach()
    SCREAMING_SNAKE_CASE				:				int				      = list(layer.weight.size()   )
    SCREAMING_SNAKE_CASE				:				Tuple				      = len(a__   )
    SCREAMING_SNAKE_CASE				:				Any				      = ConvaD(new_size[1] ,						new_size[0]   ).to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = False
    new_layer.weight.copy_(W.contiguous()   )
    SCREAMING_SNAKE_CASE				:				int				      = True
    SCREAMING_SNAKE_CASE				:				Any				      = False
    new_layer.bias.copy_(b.contiguous()   )
    SCREAMING_SNAKE_CASE				:				int				      = True
    return new_layer







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ = None   ):



    """simple docstring"""



    if isinstance(a__ ,						nn.Linear   ):
        return prune_linear_layer(a__ ,						a__ ,						dim=0 if dim is None else dim   )
    elif isinstance(a__ ,						a__   ):
        return prune_convad_layer(a__ ,						a__ ,						dim=1 if dim is None else dim   )
    else:
        raise ValueError(F"""Can't prune layer of class {layer.__class__}"""   )







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						*a__   ):



    """simple docstring"""



    assert len(a__   ) > 0, F"""{input_tensors} has to be a tuple/list of tensors"""

    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
    SCREAMING_SNAKE_CASE				:				str				      = len(inspect.signature(a__   ).parameters   )
    if num_args_in_forward_chunk_fn != len(a__   ):
        raise ValueError(
            F"""forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(a__   )} input """
            '''tensors are given'''   )

    if chunk_size > 0:
        SCREAMING_SNAKE_CASE				:				Dict				      = input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    F"""All input tenors have to be of the same shape: {tensor_shape}, """
                    F"""found shape {input_tensor.shape[chunk_dim]}"""   )

        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                F"""The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk """
                F"""size {chunk_size}"""   )

        SCREAMING_SNAKE_CASE				:				Any				      = input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = tuple(input_tensor.chunk(a__ ,						dim=a__   ) for input_tensor in input_tensors   )
        # apply forward fn to every tuple
        SCREAMING_SNAKE_CASE				:				Tuple				      = tuple(forward_fn(*a__   ) for input_tensors_chunk in zip(*a__   )   )
        # concatenate output at same dimension
        return torch.cat(a__ ,						dim=a__   )

    return forward_fn(*a__   )







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						a__   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.ones(a__ ,						a__   )
    SCREAMING_SNAKE_CASE				:				str				      = set(a__   ) - already_pruned_heads  # Convert to set and remove already pruned heads
    for head in heads:
        # Compute how many pruned heads are before the head and move the index accordingly
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = head - sum(1 if h < head else 0 for h in already_pruned_heads   )
        SCREAMING_SNAKE_CASE				:				str				      = 0
    SCREAMING_SNAKE_CASE				:				List[str]				      = mask.view(-1   ).contiguous().eq(1   )
    SCREAMING_SNAKE_CASE				:				torch.LongTensor				      = torch.arange(len(a__   )   )[mask].long()
    return heads, index







def 				UpperCAmelCase_(    *a__ ,						a__ = None   ):



    """simple docstring"""



    return torch.meshgrid(*a__ ,						indexing=a__   )







def 				UpperCAmelCase_(    a__   ):



    """simple docstring"""



    return tensor.device, storage_ptr(a__   ), storage_size(a__   )
