


'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


_a      :      Optional[int]       				=      {"configuration_gptj": ["GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTJConfig", "GPTJOnnxConfig"]}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _a      :      Optional[Any]       				=      [
       "GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST",
       "GPTJForCausalLM",
       "GPTJForQuestionAnswering",
       "GPTJForSequenceClassification",
       "GPTJModel",
       "GPTJPreTrainedModel",
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _a      :      List[Any]       				=      [
       "TFGPTJForCausalLM",
       "TFGPTJForQuestionAnswering",
       "TFGPTJForSequenceClassification",
       "TFGPTJModel",
       "TFGPTJPreTrainedModel",
   ]

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _a      :      Dict       				=      [
       "FlaxGPTJForCausalLM",
       "FlaxGPTJModel",
       "FlaxGPTJPreTrainedModel",
   ]


if TYPE_CHECKING:
   from .configuration_gptj import GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTJConfig, GPTJOnnxConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_gptj import (
          GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST,
          GPTJForCausalLM,
          GPTJForQuestionAnswering,
          GPTJForSequenceClassification,
          GPTJModel,
          GPTJPreTrainedModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_gptj import (
          TFGPTJForCausalLM,
          TFGPTJForQuestionAnswering,
          TFGPTJForSequenceClassification,
          TFGPTJModel,
          TFGPTJPreTrainedModel,
      )

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel, FlaxGPTJPreTrainedModel

else:
   import sys

   _a      :      Tuple       				=      _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)








'''simple docstring'''




from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types





def      _a		(lowercase__ :      torch.Tensor       ,      lowercase__ :      torch.Tensor       ,      lowercase__ :      bool	) -> torch.Tensor:






     """simple docstring"""
     # The first operation in a checkpoint can't be in-place, but it's
     # nice to have in-place addition during inference. Thus...
     if not inplace:
          __snake_case    			=	ma + ma
     else:
          ma += ma

     return ma





def      _a		(lowercase__ :      torch.Tensor       ,      lowercase__ :      List[int]	) -> torch.Tensor:






     """simple docstring"""
     __snake_case    			=	-1 * len(lowercase__	)
     __snake_case    			=	list(range(len(tensor.shape[:zero_index]	)	)	)
     return tensor.permute(first_inds + [zero_index + i for i in inds]	)





def      _a		(lowercase__ :      torch.Tensor       ,      lowercase__ :      int	) -> torch.Tensor:






     """simple docstring"""
     return t.reshape(t.shape[:-no_dims] + (-1,)	)





def      _a		(lowercase__ :      torch.Tensor       ,      lowercase__ :      torch.Tensor       ,      lowercase__ :      int       ,      lowercase__ :      float = 1e-4	) -> torch.Tensor:






     """simple docstring"""
     __snake_case    			=	mask.expand(*value.shape	)
     return torch.sum(mask * value       ,      dim=lowercase__	) / (eps + torch.sum(lowercase__       ,      dim=lowercase__	))





def      _a		(lowercase__ :      torch.Tensor       ,      lowercase__ :      torch.types.Number = 2.31_25       ,      lowercase__ :      torch.types.Number = 21.68_75       ,      lowercase__ :      int = 6_4	) -> torch.Tensor:






     """simple docstring"""
     __snake_case    			=	torch.linspace(lowercase__       ,      lowercase__       ,      no_bins - 1       ,      device=pts.device	)
     __snake_case    			=	torch.sqrt(torch.sum((pts.unsqueeze(-2	) - pts.unsqueeze(-3	)) ** 2       ,      dim=-1	)	)
     return torch.bucketize(lowercase__       ,      lowercase__	)





def      _a		(lowercase__ :      Callable[[list], Any]       ,      lowercase__ :      List[dict]	) -> dict:






     """simple docstring"""
     __snake_case    			=	dicts[0]
     __snake_case    			=	{}
     for k, v in first.items():
          __snake_case    			=	[d[k] for d in dicts]
          if isinstance(lowercase__       ,      lowercase__	):
               __snake_case    			=	dict_multimap(lowercase__       ,      lowercase__	)
          else:
               __snake_case    			=	fn(lowercase__	)

     return new_dict





def      _a		(lowercase__ :      torch.Tensor       ,      lowercase__ :      torch.Tensor	) -> torch.Tensor:






     """simple docstring"""
     __snake_case    			=	v_bins.view(((1,) * len(x.shape	)) + (len(lowercase__	),)	)
     __snake_case    			=	x[..., None] - reshaped_bins
     __snake_case    			=	torch.argmin(torch.abs(lowercase__	)       ,      dim=-1	)
     return nn.functional.one_hot(lowercase__       ,      num_classes=len(lowercase__	)	).float()





def      _a		(lowercase__ :      torch.Tensor       ,      lowercase__ :      torch.Tensor       ,      lowercase__ :      int = 0       ,      lowercase__ :      int = 0	) -> torch.Tensor:






     """simple docstring"""
     __snake_case    			=	[]
     for i, s in enumerate(data.shape[:no_batch_dims]	):
          __snake_case    			=	torch.arange(lowercase__	)
          __snake_case    			=	r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape	) - i - 1)))	)
          ranges.append(lowercase__	)

     __snake_case    			=	[slice(lowercase__	) for _ in range(len(data.shape	) - no_batch_dims	)]
     __snake_case    			=	inds
     ranges.extend(lowercase__	)
     # Matt note: Editing this to get around the behaviour of using a list as an array index changing
     # in recent Numpy versions
     return data[tuple(lowercase__	)]


_a      :      str       				=      TypeVar("T")





def      _a		(lowercase__ :      Callable[[T], Any]       ,      lowercase__ :      Dict[Any, Union[dict, list, tuple, T]]       ,      lowercase__ :      Type[T]	) -> Dict[Any, Union[dict, list, tuple, Any]]:






     """simple docstring"""
     __snake_case    			=	{}
     for k, v in dic.items():
          if isinstance(lowercase__       ,      lowercase__	):
               __snake_case    			=	dict_map(lowercase__       ,      lowercase__       ,      lowercase__	)
          else:
               __snake_case    			=	tree_map(lowercase__       ,      lowercase__       ,      lowercase__	)

     return new_dict





@overload
def      _a		(lowercase__ :      Callable[[T], Any]       ,      lowercase__ :      T       ,      lowercase__ :      Type[T]	) -> Any:






     """simple docstring"""
     ...





@overload
def      _a		(lowercase__ :      Callable[[T], Any]       ,      lowercase__ :      dict       ,      lowercase__ :      Type[T]	) -> dict:






     """simple docstring"""
     ...





@overload
def      _a		(lowercase__ :      Callable[[T], Any]       ,      lowercase__ :      list       ,      lowercase__ :      Type[T]	) -> list:






     """simple docstring"""
     ...





@overload
def      _a		(lowercase__ :      Callable[[T], Any]       ,      lowercase__ :      tuple       ,      lowercase__ :      Type[T]	) -> tuple:






     """simple docstring"""
     ...





def      _a		(lowercase__ :      str       ,      lowercase__ :      Any       ,      lowercase__ :      Dict	) -> Union[str, Any]:






     """simple docstring"""
     if isinstance(lowercase__       ,      lowercase__	):
          return dict_map(lowercase__       ,      lowercase__       ,      lowercase__	)
     elif isinstance(lowercase__       ,      lowercase__	):
          return [tree_map(lowercase__       ,      lowercase__       ,      lowercase__	) for x in tree]
     elif isinstance(lowercase__       ,      lowercase__	):
          return tuple(tree_map(lowercase__       ,      lowercase__       ,      lowercase__	) for x in tree	)
     elif isinstance(lowercase__       ,      lowercase__	):
          return fn(lowercase__	)
     else:
          print(type(lowercase__	)	)
          raise ValueError('Not supported'	)


_a      :      Union[str, Any]       				=      partial(tree_map, leaf_type=torch.Tensor)




