code: 
  
  
  
  
  
  '''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 		UpperCamelCase       (							_lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      bool ):
      # 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:
          A__       =    ma + ma
      else:
          ma += ma
  
      return ma
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      List[int] ):
      A__       =    -1 * len(_lowerCamelCase )
      A__       =    list(range(len(tensor.shape[:zero_index] ) ) )
      return tensor.permute(first_inds + [zero_index + i for i in inds] )
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      int ):
      return t.reshape(t.shape[:-no_dims] + (-1,) )
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      int		, _lowerCamelCase      :      float = 1e-4 ):
      A__       =    mask.expand(*value.shape )
      return torch.sum(mask * value		, dim=_lowerCamelCase ) / (eps + torch.sum(_lowerCamelCase		, dim=_lowerCamelCase ))
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      torch.types.Number = 2.3_1_2_5		, _lowerCamelCase      :      torch.types.Number = 2_1.6_8_7_5		, _lowerCamelCase      :      int = 64 ):
      A__       =    torch.linspace(_lowerCamelCase		, _lowerCamelCase		, no_bins - 1		, device=pts.device )
      A__       =    torch.sqrt(torch.sum((pts.unsqueeze(-2 ) - pts.unsqueeze(-3 )) ** 2		, dim=-1 ) )
      return torch.bucketize(_lowerCamelCase		, _lowerCamelCase )
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      Callable[[list], Any]		, _lowerCamelCase      :      List[dict] ):
      A__       =    dicts[0]
      A__       =    {}
      for k, v in first.items():
          A__       =    [d[k] for d in dicts]
          if isinstance(_lowerCamelCase		, _lowerCamelCase ):
              A__       =    dict_multimap(_lowerCamelCase		, _lowerCamelCase )
          else:
              A__       =    fn(_lowerCamelCase )
  
      return new_dict
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      torch.Tensor ):
      A__       =    v_bins.view(((1,) * len(x.shape )) + (len(_lowerCamelCase ),) )
      A__       =    x[..., None] - reshaped_bins
      A__       =    torch.argmin(torch.abs(_lowerCamelCase )		, dim=-1 )
      return nn.functional.one_hot(_lowerCamelCase		, num_classes=len(_lowerCamelCase ) ).float()
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      torch.Tensor		, _lowerCamelCase      :      int = 0		, _lowerCamelCase      :      int = 0 ):
      A__       =    []
      for i, s in enumerate(data.shape[:no_batch_dims] ):
          A__       =    torch.arange(_lowerCamelCase )
          A__       =    r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape ) - i - 1))) )
          ranges.append(_lowerCamelCase )
  
      A__       =    [slice(_lowerCamelCase ) for _ in range(len(data.shape ) - no_batch_dims )]
      A__       =    inds
      ranges.extend(_lowerCamelCase )
      # 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(_lowerCamelCase )]
  
  
  __lowerCAmelCase							:			Tuple				=TypeVar("T")
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      Callable[[T], Any]		, _lowerCamelCase      :      Dict[Any, Union[dict, list, tuple, T]]		, _lowerCamelCase      :      Type[T] ):
      A__       =    {}
      for k, v in dic.items():
          if isinstance(_lowerCamelCase		, _lowerCamelCase ):
              A__       =    dict_map(_lowerCamelCase		, _lowerCamelCase		, _lowerCamelCase )
          else:
              A__       =    tree_map(_lowerCamelCase		, _lowerCamelCase		, _lowerCamelCase )
  
      return new_dict
  
  
  
  
  
  @overload
  def 		UpperCamelCase       (							_lowerCamelCase      :      Callable[[T], Any]		, _lowerCamelCase      :      T		, _lowerCamelCase      :      Type[T] ):
      ...
  
  
  
  
  
  @overload
  def 		UpperCamelCase       (							_lowerCamelCase      :      Callable[[T], Any]		, _lowerCamelCase      :      dict		, _lowerCamelCase      :      Type[T] ):
      ...
  
  
  
  
  
  @overload
  def 		UpperCamelCase       (							_lowerCamelCase      :      Callable[[T], Any]		, _lowerCamelCase      :      list		, _lowerCamelCase      :      Type[T] ):
      ...
  
  
  
  
  
  @overload
  def 		UpperCamelCase       (							_lowerCamelCase      :      Callable[[T], Any]		, _lowerCamelCase      :      tuple		, _lowerCamelCase      :      Type[T] ):
      ...
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      Any		, _lowerCamelCase      :      Optional[int]		, _lowerCamelCase      :      List[Any] ):
      if isinstance(_lowerCamelCase		, _lowerCamelCase ):
          return dict_map(_lowerCamelCase		, _lowerCamelCase		, _lowerCamelCase )
      elif isinstance(_lowerCamelCase		, _lowerCamelCase ):
          return [tree_map(_lowerCamelCase		, _lowerCamelCase		, _lowerCamelCase ) for x in tree]
      elif isinstance(_lowerCamelCase		, _lowerCamelCase ):
          return tuple(tree_map(_lowerCamelCase		, _lowerCamelCase		, _lowerCamelCase ) for x in tree )
      elif isinstance(_lowerCamelCase		, _lowerCamelCase ):
          return fn(_lowerCamelCase )
      else:
          print(type(_lowerCamelCase ) )
          raise ValueError("Not supported" )
  
  
  
  __lowerCAmelCase							:			Dict				=partial(tree_map, leaf_type=torch.Tensor)
  
  
  
  
style_context: 
  
  
  
  
  
  '''simple docstring'''
  
  import os
  import tempfile
  import unittest
  from typing import List
  
  from transformers.models.esm.tokenization_esm import VOCAB_FILES_NAMES, EsmTokenizer
  from transformers.testing_utils import require_tokenizers
  from transformers.tokenization_utils import PreTrainedTokenizer
  from transformers.tokenization_utils_base import PreTrainedTokenizerBase
  
  
  
  
  
  
  @require_tokenizers
  class      UpperCAmelCase			( unittest.TestCase       ):
   __lowercase     	=       EsmTokenizer
   def   UpperCAmelCase_							(       self   :Any							)->							int:
       super().setUp()
       A__       =    tempfile.mkdtemp()
       # fmt: off
       A__       =    ["<cls>", "<pad>", "<eos>", "<unk>", "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", "O", ".", "-", "<null_1>", "<mask>"]  # noqa: E501
       # fmt: on
       A__       =    os.path.join(self.tmpdirname		, VOCAB_FILES_NAMES["vocab_file"]							)
       with open(self.vocab_file		, "w"		, encoding="utf-8"							) as vocab_writer:
           vocab_writer.write("".join([x + "\n" for x in vocab_tokens]							)							)
   def   UpperCAmelCase_							(       self   :Any		, **lowercase_   :int							)->							List[PreTrainedTokenizerBase]:
       return [self.get_tokenizer(**lowercase_							)]
   def   UpperCAmelCase_							(       self   :Any		, **lowercase_   :Dict							)->							PreTrainedTokenizer:
       return self.tokenizer_class.from_pretrained(self.tmpdirname		, **lowercase_							)
   def   UpperCAmelCase_							(       self   :int							)->							Dict:
       A__       =    self.tokenizer_class(self.vocab_file							)
  
       A__       =    tokenizer.tokenize("LAGVS"							)
       self.assertListEqual(lowercase_		, ["L", "A", "G", "V", "S"]							)
       self.assertListEqual(tokenizer.convert_tokens_to_ids(lowercase_							)		, [4, 5, 6, 7, 8]							)
   def   UpperCAmelCase_							(       self   :Tuple							)->							Tuple:
       A__       =    self.tokenizer_class(self.vocab_file							)
  
       A__       =    "LAGVS"
       self.assertListEqual(tokenizer.encode(lowercase_							)		, [0, 4, 5, 6, 7, 8, 2]							)
   def   UpperCAmelCase_							(       self   :Any							)->							Tuple:
       A__       =    self.tokenizer_class(self.vocab_file							)
  
       A__       =    ["LAGVS", "WCB"]
       A__       =    tokenizer(lowercase_		, padding=lowercase_							)["input_ids"]
  
       self.assertListEqual(lowercase_		, [[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2]]							)
   def   UpperCAmelCase_							(       self   :str							)->							Any:
       A__       =    self.tokenizer_class(self.vocab_file							)
  
       A__       =    ["LAGVS", "WCB"]
       A__       =    tokenizer(lowercase_		, padding=lowercase_							)["input_ids"]
  
       self.assertListEqual(lowercase_		, [[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2, 1, 1]]							)
  
  
  
   def   UpperCAmelCase_							(       self   :Optional[Any]							)->							List[str]:
       A__       =    self.get_tokenizers(fast=lowercase_							)
       for tokenizer in tokenizers:
           with self.subTest(F"{tokenizer.__class__.__name__}"							):
               A__       =    "<unk>"
               A__       =    "<mask>"
  
               A__       =    tokenizer.tokenize(lowercase_							)
               A__       =    tokenizer.tokenize(lowercase_							)
  
               self.assertEqual(len(lowercase_							)		, 1							)
               self.assertEqual(len(lowercase_							)		, 1							)
               self.assertEqual(token_a[0]		, lowercase_							)
               self.assertEqual(token_a[0]		, lowercase_							)
  
  
  
  