'''simple docstring'''




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


__snake_case  : Tuple		       =    [nn.LayerNorm]

__snake_case  : List[Any]		       =    logging.get_logger(__name__)

__snake_case  : Dict		       =    version.parse(version.parse(torch.__version__).base_version)

__snake_case  : Any		       =    parsed_torch_version_base >= version.parse('2.0')
__snake_case  : int		       =    parsed_torch_version_base >= version.parse('1.12')
__snake_case  : int		       =    parsed_torch_version_base >= version.parse('1.11')
__snake_case  : List[str]		       =    parsed_torch_version_base < version.parse('1.11')





def 	_UpperCAmelCase       (				_UpperCamelCase       :      Optional[Any],  _UpperCamelCase       :      Optional[Any],  _UpperCamelCase       :      List[Any],  _UpperCamelCase       :      Dict,  self       :      List[str]  )      -> Optional[int]:
    from torch import _softmax_backward_data

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





def 	_UpperCAmelCase       (				_UpperCamelCase       :      nn.Linear,  _UpperCamelCase       :      torch.LongTensor,  _UpperCamelCase       :      int = 0  )      -> nn.Linear:
    A_      =  index.to(layer.weight.device  )
    A_      =  layer.weight.index_select(_UpperCamelCase,  _UpperCamelCase  ).clone().detach()
    if layer.bias is not None:
        if dim == 1:
            A_      =  layer.bias.clone().detach()
        else:
            A_      =  layer.bias[index].clone().detach()
    A_      =  list(layer.weight.size()  )
    A_      =  len(_UpperCamelCase  )
    A_      =  nn.Linear(new_size[1],  new_size[0],  bias=layer.bias is not None  ).to(layer.weight.device  )
    A_      =  False
    new_layer.weight.copy_(W.contiguous()  )
    A_      =  True
    if layer.bias is not None:
        A_      =  False
        new_layer.bias.copy_(b.contiguous()  )
        A_      =  True
    return new_layer




class       __UpperCAmelCase							(							nn.Module	):





       '''simple docstring'''



       def __init__( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)						->	Optional[Any]:
           super().__init__()
           A_      =  nf
           A_      =  nn.Parameter(torch.empty(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)					)
           A_      =  nn.Parameter(torch.zeros(_SCREAMING_SNAKE_CASE					)					)
           nn.init.normal_(self.weight       ,	std=0.02					)

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE					)						->	Optional[Any]:
           A_      =  x.size()[:-1] + (self.nf,)
           A_      =  torch.addmm(self.bias       ,	x.view(-1       ,	x.size(-1					)					)       ,	self.weight					)
           A_      =  x.view(_SCREAMING_SNAKE_CASE					)
           return x





def 	_UpperCAmelCase       (				_UpperCamelCase       :      ConvaD,  _UpperCamelCase       :      torch.LongTensor,  _UpperCamelCase       :      int = 1  )      -> ConvaD:
    A_      =  index.to(layer.weight.device  )
    A_      =  layer.weight.index_select(_UpperCamelCase,  _UpperCamelCase  ).clone().detach()
    if dim == 0:
        A_      =  layer.bias.clone().detach()
    else:
        A_      =  layer.bias[index].clone().detach()
    A_      =  list(layer.weight.size()  )
    A_      =  len(_UpperCamelCase  )
    A_      =  ConvaD(new_size[1],  new_size[0]  ).to(layer.weight.device  )
    A_      =  False
    new_layer.weight.copy_(W.contiguous()  )
    A_      =  True
    A_      =  False
    new_layer.bias.copy_(b.contiguous()  )
    A_      =  True
    return new_layer





def 	_UpperCAmelCase       (				_UpperCamelCase       :      Union[nn.Linear, ConvaD],  _UpperCamelCase       :      torch.LongTensor,  _UpperCamelCase       :      Optional[int] = None  )      -> Union[nn.Linear, ConvaD]:
    if isinstance(_UpperCamelCase,  nn.Linear  ):
        return prune_linear_layer(_UpperCamelCase,  _UpperCamelCase,  dim=0 if dim is None else dim  )
    elif isinstance(_UpperCamelCase,  _UpperCamelCase  ):
        return prune_convad_layer(_UpperCamelCase,  _UpperCamelCase,  dim=1 if dim is None else dim  )
    else:
        raise ValueError(F'''Can\'t prune layer of class {layer.__class__}'''  )





def 	_UpperCAmelCase       (				_UpperCamelCase       :      Callable[..., torch.Tensor],  _UpperCamelCase       :      int,  _UpperCamelCase       :      int,  *_UpperCamelCase       :      List[str]  )      -> torch.Tensor:
    assert len(_UpperCamelCase  ) > 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
    A_      =  len(inspect.signature(_UpperCamelCase  ).parameters  )
    if num_args_in_forward_chunk_fn != len(_UpperCamelCase  ):
        raise ValueError(
            F'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(_UpperCamelCase  )} input '''
            '''tensors are given'''  )

    if chunk_size > 0:
        A_      =  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}'''  )

        A_      =  input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        A_      =  tuple(input_tensor.chunk(_UpperCamelCase,  dim=_UpperCamelCase  ) for input_tensor in input_tensors  )
        # apply forward fn to every tuple
        A_      =  tuple(forward_fn(*_UpperCamelCase  ) for input_tensors_chunk in zip(*_UpperCamelCase  )  )
        # concatenate output at same dimension
        return torch.cat(_UpperCamelCase,  dim=_UpperCamelCase  )

    return forward_fn(*_UpperCamelCase  )





def 	_UpperCAmelCase       (				_UpperCamelCase       :      List[int],  _UpperCamelCase       :      int,  _UpperCamelCase       :      int,  _UpperCamelCase       :      Set[int]  )      -> Tuple[Set[int], torch.LongTensor]:
    A_      =  torch.ones(_UpperCamelCase,  _UpperCamelCase  )
    A_      =  set(_UpperCamelCase  ) - 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
        A_      =  head - sum(1 if h < head else 0 for h in already_pruned_heads  )
        A_      =  0
    A_      =  mask.view(-1  ).contiguous().eq(1  )
    A_      =  torch.arange(len(_UpperCamelCase  )  )[mask].long()
    return heads, index





def 	_UpperCAmelCase       (				*_UpperCamelCase       :      Union[torch.Tensor, List[torch.Tensor]],  _UpperCamelCase       :      Optional[str] = None  )      -> Tuple[torch.Tensor, ...]:
    return torch.meshgrid(*_UpperCamelCase,  indexing=_UpperCamelCase  )





def 	_UpperCAmelCase       (				_UpperCamelCase       :      torch.Tensor  )      -> Tuple[torch.device, int, int]:
    return tensor.device, storage_ptr(_UpperCamelCase  ), storage_size(_UpperCamelCase  )







'''simple docstring'''




import itertools





def 	_UpperCAmelCase       (				_UpperCamelCase       :      int = 1_00_00_00  )      -> int:
    A_      =  [1]

    for i in itertools.count(len(_UpperCamelCase  )  ):
        A_      =  0
        for j in itertools.count(1  ):
            A_      =  -1 if j % 2 == 0 else +1
            A_      =  (j * j * 3 - j) // 2
            if index > i:
                break
            item += partitions[i - index] * sign
            item %= number
            index += j
            if index > i:
                break
            item += partitions[i - index] * sign
            item %= number

        if item == 0:
            return i
        partitions.append(_UpperCamelCase  )

    return 0


if __name__ == "__main__":
  import doctest

  doctest.testmod()

  print(F"""{solution() = }""")






