

"""simple docstring"""




from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


__UpperCAmelCase			=					[
    'openmmlab/upernet-convnext-tiny',
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
__UpperCAmelCase			=					'UperNetConfig'







class   __lowercase      (     nn.Module       ):

  def __init__(	self		:      int       ,A		:      int       ,A		:      int       ,A		:      Union[int, Tuple[int, int]]       ,A		:      Union[int, Tuple[int, int], str] = 0       ,A		:      bool = False       ,A		:      Union[int, Tuple[int, int]] = 1       ,):






   '''simple docstring'''
   super().__init__()
   UpperCAmelCase__       :	Any        =       nn.Convad(
       in_channels=A       ,out_channels=A       ,kernel_size=A       ,padding=A       ,bias=A       ,dilation=A       ,)
   UpperCAmelCase__       :	List[Any]        =       nn.BatchNormad(A     )
   UpperCAmelCase__       :	Optional[Any]        =       nn.ReLU()




  def  __lowercase							(	self		:      List[Any]       ,A		:      torch.Tensor     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Any        =       self.conv(A     )
   UpperCAmelCase__       :	Optional[Any]        =       self.batch_norm(A     )
   UpperCAmelCase__       :	Dict        =       self.activation(A     )

   return output







class   __lowercase      (     nn.Module       ):

  def __init__(	self		:      int       ,A		:      int       ,A		:      int       ,A		:      int     ):






   '''simple docstring'''
   super().__init__()
   UpperCAmelCase__       :	int        =       [
       nn.AdaptiveAvgPoolad(A     ),
       UperNetConvModule(A       ,A       ,kernel_size=1     ),
   ]
   for i, layer in enumerate(self.layers     ):
    self.add_module(str(A     )       ,A     )




  def  __lowercase							(	self		:      Dict       ,A		:      torch.Tensor     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[int]        =       input
   for layer in self.layers:
    UpperCAmelCase__       :	Optional[int]        =       layer(A     )
   return hidden_state







class   __lowercase      (     nn.Module       ):

  def __init__(	self		:      Union[str, Any]       ,A		:      Tuple[int, ...]       ,A		:      int       ,A		:      int       ,A		:      bool     ):






   '''simple docstring'''
   super().__init__()
   UpperCAmelCase__       :	Dict        =       pool_scales
   UpperCAmelCase__       :	Optional[int]        =       align_corners
   UpperCAmelCase__       :	Optional[int]        =       in_channels
   UpperCAmelCase__       :	List[str]        =       channels
   UpperCAmelCase__       :	Tuple        =       []
   for i, pool_scale in enumerate(A     ):
    UpperCAmelCase__       :	str        =       UperNetPyramidPoolingBlock(pool_scale=A       ,in_channels=A       ,channels=A     )
    self.blocks.append(A     )
    self.add_module(str(A     )       ,A     )




  def  __lowercase							(	self		:      List[Any]       ,A		:      torch.Tensor     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[Any]        =       []
   for ppm in self.blocks:
    UpperCAmelCase__       :	Tuple        =       ppm(A     )
    UpperCAmelCase__       :	List[str]        =       nn.functional.interpolate(
        A       ,size=x.size()[2:]       ,mode="""bilinear"""       ,align_corners=self.align_corners     )
    ppm_outs.append(A     )
   return ppm_outs







class   __lowercase      (     nn.Module       ):

  def __init__(	self		:      Union[str, Any]       ,A		:      str       ,A		:      Optional[Any]     ):






   '''simple docstring'''
   super().__init__()

   UpperCAmelCase__       :	str        =       config
   UpperCAmelCase__       :	Union[str, Any]        =       config.pool_scales  # e.g. (1, 2, 3, 6)
   UpperCAmelCase__       :	Tuple        =       in_channels
   UpperCAmelCase__       :	Union[str, Any]        =       config.hidden_size
   UpperCAmelCase__       :	str        =       False
   UpperCAmelCase__       :	Optional[Any]        =       nn.Convad(self.channels       ,config.num_labels       ,kernel_size=1     )

   # PSP Module
   UpperCAmelCase__       :	int        =       UperNetPyramidPoolingModule(
       self.pool_scales       ,self.in_channels[-1]       ,self.channels       ,align_corners=self.align_corners       ,)
   UpperCAmelCase__       :	Tuple        =       UperNetConvModule(
       self.in_channels[-1] + len(self.pool_scales     ) * self.channels       ,self.channels       ,kernel_size=3       ,padding=1       ,)
   # FPN Module
   UpperCAmelCase__       :	Tuple        =       nn.ModuleList()
   UpperCAmelCase__       :	Optional[int]        =       nn.ModuleList()
   for in_channels in self.in_channels[:-1]:  # skip the top layer
    UpperCAmelCase__       :	Optional[Any]        =       UperNetConvModule(A       ,self.channels       ,kernel_size=1     )
    UpperCAmelCase__       :	Dict        =       UperNetConvModule(self.channels       ,self.channels       ,kernel_size=3       ,padding=1     )
    self.lateral_convs.append(A     )
    self.fpn_convs.append(A     )

   UpperCAmelCase__       :	Tuple        =       UperNetConvModule(
       len(self.in_channels     ) * self.channels       ,self.channels       ,kernel_size=3       ,padding=1       ,)

  def  __lowercase							(	self		:      Dict     ):






   '''simple docstring'''
   self.apply(self._init_weights     )

  def  __lowercase							(	self		:      Dict       ,A		:      List[Any]     ):






   '''simple docstring'''
   if isinstance(A       ,nn.Convad     ):
    module.weight.data.normal_(mean=0.0       ,std=self.config.initializer_range     )
    if module.bias is not None:
     module.bias.data.zero_()

  def  __lowercase							(	self		:      str       ,A		:      Tuple     ):






   '''simple docstring'''
   UpperCAmelCase__       :	Optional[int]        =       inputs[-1]
   UpperCAmelCase__       :	Any        =       [x]
   psp_outs.extend(self.psp_modules(A     )     )
   UpperCAmelCase__       :	Union[str, Any]        =       torch.cat(A       ,dim=1     )
   UpperCAmelCase__       :	Dict        =       self.bottleneck(A     )

   return output




  def  __lowercase							(	self		:      List[Any]       ,A		:      torch.Tensor     ):






   '''simple docstring'''
   # build laterals
   UpperCAmelCase__       :	Optional[int]        =       [lateral_conv(encoder_hidden_states[i]     ) for i, lateral_conv in enumerate(self.lateral_convs     )]

   laterals.append(self.psp_forward(A     )     )

   # build top-down path
   UpperCAmelCase__       :	Optional[int]        =       len(A     )
   for i in range(used_backbone_levels - 1       ,0       ,-1     ):
    UpperCAmelCase__       :	Optional[Any]        =       laterals[i - 1].shape[2:]
    UpperCAmelCase__       :	List[str]        =       laterals[i - 1] + nn.functional.interpolate(
        laterals[i]       ,size=A       ,mode="""bilinear"""       ,align_corners=self.align_corners     )

   # build outputs
   UpperCAmelCase__       :	List[Any]        =       [self.fpn_convs[i](laterals[i]     ) for i in range(used_backbone_levels - 1     )]
   # append psp feature
   fpn_outs.append(laterals[-1]     )

   for i in range(used_backbone_levels - 1       ,0       ,-1     ):
    UpperCAmelCase__       :	List[Any]        =       nn.functional.interpolate(
        fpn_outs[i]       ,size=fpn_outs[0].shape[2:]       ,mode="""bilinear"""       ,align_corners=self.align_corners     )
   UpperCAmelCase__       :	str        =       torch.cat(A       ,dim=1     )
   UpperCAmelCase__       :	Dict        =       self.fpn_bottleneck(A     )
   UpperCAmelCase__       :	Tuple        =       self.classifier(A     )

   return output







class   __lowercase      (     nn.Module       ):

  def __init__(	self		:      Optional[Any]       ,A		:      Tuple       ,A		:      int = 2       ,A		:      int = 3       ,A		:      Union[int, Tuple[int, int]] = 1     ):






   '''simple docstring'''
   super().__init__()

   UpperCAmelCase__       :	Optional[Any]        =       config
   UpperCAmelCase__       :	List[str]        =       config.auxiliary_in_channels
   UpperCAmelCase__       :	int        =       config.auxiliary_channels
   UpperCAmelCase__       :	List[Any]        =       config.auxiliary_num_convs
   UpperCAmelCase__       :	Any        =       config.auxiliary_concat_input
   UpperCAmelCase__       :	Tuple        =       in_index

   UpperCAmelCase__       :	Optional[int]        =       (kernel_size // 2) * dilation
   UpperCAmelCase__       :	Optional[Any]        =       []
   convs.append(
       UperNetConvModule(
           self.in_channels       ,self.channels       ,kernel_size=A       ,padding=A       ,dilation=A     )     )
   for i in range(self.num_convs - 1     ):
    convs.append(
        UperNetConvModule(
            self.channels       ,self.channels       ,kernel_size=A       ,padding=A       ,dilation=A     )     )
   if self.num_convs == 0:
    UpperCAmelCase__       :	Tuple        =       nn.Identity()
   else:
    UpperCAmelCase__       :	List[str]        =       nn.Sequential(*A     )
   if self.concat_input:
    UpperCAmelCase__       :	str        =       UperNetConvModule(
        self.in_channels + self.channels       ,self.channels       ,kernel_size=A       ,padding=kernel_size // 2     )

   UpperCAmelCase__       :	int        =       nn.Convad(self.channels       ,config.num_labels       ,kernel_size=1     )

  def  __lowercase							(	self		:      Optional[int]     ):






   '''simple docstring'''
   self.apply(self._init_weights     )

  def  __lowercase							(	self		:      Any       ,A		:      Any     ):






   '''simple docstring'''
   if isinstance(A       ,nn.Convad     ):
    module.weight.data.normal_(mean=0.0       ,std=self.config.initializer_range     )
    if module.bias is not None:
     module.bias.data.zero_()




  def  __lowercase							(	self		:      Optional[int]       ,A		:      torch.Tensor     ):






   '''simple docstring'''
   # just take the relevant feature maps
   UpperCAmelCase__       :	int        =       encoder_hidden_states[self.in_index]
   UpperCAmelCase__       :	int        =       self.convs(A     )
   if self.concat_input:
    UpperCAmelCase__       :	str        =       self.conv_cat(torch.cat([hidden_states, output]       ,dim=1     )     )
   UpperCAmelCase__       :	Tuple        =       self.classifier(A     )
   return output







class   __lowercase      (     __lowerCamelCase       ):
  snake_case_												=							UperNetConfig
  snake_case_												=							"""pixel_values"""
  snake_case_												=							True

  def  __lowercase							(	self		:      Dict       ,A		:      str     ):






   '''simple docstring'''
   if isinstance(A       ,A     ):
    module.backbone.init_weights()
    module.decode_head.init_weights()
    module.auxiliary_head.init_weights()

  def  __lowercase							(	self		:      str     ):






   '''simple docstring'''
   self.backbone.init_weights()
   self.decode_head.init_weights()
   self.auxiliary_head.init_weights()




  def  __lowercase							(	self		:      Union[str, Any]       ,A		:      List[Any]       ,A		:      int=False     ):






   '''simple docstring'''
   if isinstance(A       ,A     ):
    UpperCAmelCase__       :	str        =       value







__UpperCAmelCase			=					r'\n    Parameters:\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n        config ([`UperNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'

__UpperCAmelCase			=					r'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using\n            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.\n        output_attentions (`bool`, *optional*):\n            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See\n            `attentions` under returned tensors for more detail.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under\n            returned tensors for more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'







@add_start_docstrings(
    """UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes."""						,      __lowerCamelCase						,      )
class   __lowercase      (     __lowerCamelCase       ):

  def __init__(	self		:      Dict       ,A		:      Optional[int]     ):






   '''simple docstring'''
   super().__init__(A     )

   UpperCAmelCase__       :	Any        =       AutoBackbone.from_config(config.backbone_config     )

   # Semantic segmentation head(s)
   UpperCAmelCase__       :	Any        =       UperNetHead(A       ,in_channels=self.backbone.channels     )
   UpperCAmelCase__       :	Dict        =       UperNetFCNHead(A     ) if config.use_auxiliary_head else None

   # Initialize weights and apply final processing
   self.post_init()




  @add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format("""batch_size, sequence_length"""     )     )
  @replace_return_docstrings(output_type=A       ,config_class=_CONFIG_FOR_DOC     )
  def  __lowercase							(	self		:      Optional[int]       ,A		:      Optional[torch.Tensor] = None       ,A		:      Optional[bool] = None       ,A		:      Optional[bool] = None       ,A		:      Optional[torch.Tensor] = None       ,A		:      Optional[bool] = None       ,):






   '''simple docstring'''
   UpperCAmelCase__       :	List[str]        =       return_dict if return_dict is not None else self.config.use_return_dict
   UpperCAmelCase__       :	str        =       (
       output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
   )
   UpperCAmelCase__       :	Tuple        =       output_attentions if output_attentions is not None else self.config.output_attentions

   UpperCAmelCase__       :	int        =       self.backbone.forward_with_filtered_kwargs(
       A       ,output_hidden_states=A       ,output_attentions=A     )
   UpperCAmelCase__       :	List[Any]        =       outputs.feature_maps

   UpperCAmelCase__       :	int        =       self.decode_head(A     )
   UpperCAmelCase__       :	Optional[Any]        =       nn.functional.interpolate(A       ,size=pixel_values.shape[2:]       ,mode="""bilinear"""       ,align_corners=A     )

   UpperCAmelCase__       :	int        =       None
   if self.auxiliary_head is not None:
    UpperCAmelCase__       :	Union[str, Any]        =       self.auxiliary_head(A     )
    UpperCAmelCase__       :	Tuple        =       nn.functional.interpolate(
        A       ,size=pixel_values.shape[2:]       ,mode="""bilinear"""       ,align_corners=A     )

   UpperCAmelCase__       :	Optional[int]        =       None
   if labels is not None:
    if self.config.num_labels == 1:
     raise ValueError("""The number of labels should be greater than one"""     )
    else:
     # compute weighted loss
     UpperCAmelCase__       :	Optional[Any]        =       CrossEntropyLoss(ignore_index=self.config.loss_ignore_index     )
     UpperCAmelCase__       :	Optional[int]        =       loss_fct(A       ,A     )
     UpperCAmelCase__       :	Any        =       loss_fct(A       ,A     )
     UpperCAmelCase__       :	int        =       main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

   if not return_dict:
    if output_hidden_states:
     UpperCAmelCase__       :	Optional[int]        =       (logits,) + outputs[1:]
    else:
     UpperCAmelCase__       :	Dict        =       (logits,) + outputs[2:]
    return ((loss,) + output) if loss is not None else output

   return SemanticSegmenterOutput(
       loss=A       ,logits=A       ,hidden_states=outputs.hidden_states       ,attentions=outputs.attentions       ,)











import unittest

from transformers import JukeboxTokenizer
from transformers.testing_utils import require_torch


class 							UpperCamelCase_       (					unittest.TestCase	):

    '''simple docstring'''






    UpperCAmelCase__				    = JukeboxTokenizer
    UpperCAmelCase__				    = {
        '''artist''': '''Zac Brown Band''',
        '''genres''': '''Country''',
        '''lyrics''': '''I met a traveller from an antique land,
        Who said "Two vast and trunkless legs of stone
        Stand in the desert. . . . Near them, on the sand,
        Half sunk a shattered visage lies, whose frown,
        And wrinkled lip, and sneer of cold command,
        Tell that its sculptor well those passions read
        Which yet survive, stamped on these lifeless things,
        The hand that mocked them, and the heart that fed;
        And on the pedestal, these words appear:
        My name is Ozymandias, King of Kings;
        Look on my Works, ye Mighty, and despair!
        Nothing beside remains. Round the decay
        Of that colossal Wreck, boundless and bare
        The lone and level sands stretch far away
        ''',
    }


    @require_torch
    def 			SCREAMING_SNAKE_CASE				(     self :      Dict)		->int:





           '''simple docstring'''

           import torch

           A__ 					=       JukeboxTokenizer.from_pretrained('''openai/jukebox-1b-lyrics''')
           A__ 					=       tokenizer(**self.metas)['''input_ids''']
           # fmt: off
           A__ 					=       [
               torch.tensor([[
                   0, 0, 0, 7_169, 507, 9, 76, 39, 31, 46, 76, 27,
                   76, 46, 44, 27, 48, 31, 38, 38, 31, 44, 76, 32,
                   44, 41, 39, 76, 27, 40, 76, 27, 40, 46, 35, 43,
                   47, 31, 76, 38, 27, 40, 30, 64, 78, 76, 76, 76,
                   76, 76, 76, 76, 76, 23, 34, 41, 76, 45, 27, 35,
                   30, 76, 71, 20, 49, 41, 76, 48, 27, 45, 46, 76,
                   27, 40, 30, 76, 46, 44, 47, 40, 37, 38, 31, 45,
                   45, 76, 38, 31, 33, 45, 76, 41, 32, 76, 45, 46,
                   41, 40, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
                   19, 46, 27, 40, 30, 76, 35, 40, 76, 46, 34, 31,
                   76, 30, 31, 45, 31, 44, 46, 63, 76, 63, 76, 63,
                   76, 63, 76, 14, 31, 27, 44, 76, 46, 34, 31, 39,
                   64, 76, 41, 40, 76, 46, 34, 31, 76, 45, 27, 40,
                   30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76, 8,
                   27, 38, 32, 76, 45, 47, 40, 37, 76, 27, 76, 45,
                   34, 27, 46, 46, 31, 44, 31, 30, 76, 48, 35, 45,
                   27, 33, 31, 76, 38, 35, 31, 45, 64, 76, 49, 34,
                   41, 45, 31, 76, 32, 44, 41, 49, 40, 64, 78, 76,
                   76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76, 49,
                   44, 35, 40, 37, 38, 31, 30, 76, 38, 35, 42, 64,
                   76, 27, 40, 30, 76, 45, 40, 31, 31, 44, 76, 41,
                   32, 76, 29, 41, 38, 30, 76, 29, 41, 39, 39, 27,
                   40, 30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76,
                   20, 31, 38, 38, 76, 46, 34, 27, 46, 76, 35, 46,
                   45, 76, 45, 29, 47, 38, 42, 46, 41, 44, 76, 49,
                   31, 38, 38, 76, 46, 34, 41, 45, 31, 76, 42, 27,
                   45, 45, 35, 41, 40, 45, 76, 44, 31, 27, 30, 78,
                   76, 76, 76, 76, 76, 76, 76, 76, 23, 34, 35, 29,
                   34, 76, 51, 31, 46, 76, 45, 47, 44, 48, 35, 48,
                   31, 64, 76, 45, 46, 27, 39, 42, 31, 30, 76, 41,
                   40, 76, 46, 34, 31, 45, 31, 76, 38, 35, 32, 31,
                   38, 31, 45, 45, 76, 46, 34, 35, 40, 33, 45, 64,
                   78, 76, 76, 76, 76, 76, 76, 76, 76, 20, 34, 31,
                   76, 34, 27, 40, 30, 76, 46, 34, 27, 46, 76, 39,
                   41, 29, 37, 31, 30, 76, 46, 34, 31, 39, 64, 76,
                   27, 40, 30, 76, 46, 34, 31, 76, 34, 31, 27, 44,
                   46, 76, 46, 34, 27, 46, 76, 32, 31, 30, 66, 78,
                   76, 76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76,
                   41, 40, 76, 46, 34, 31, 76, 42, 31, 30, 31, 45,
                   46, 27, 38, 64, 76, 46, 34, 31, 45, 31, 76, 49,
                   41, 44, 30, 45, 76, 27, 42, 42, 31, 27, 44, 65,
                   78, 76, 76, 76, 76, 76, 76, 76, 76, 13, 51, 76,
                   40, 27, 39, 31, 76, 35, 45, 76, 15, 52, 51, 39,
                   27, 40, 30, 35, 27, 45, 64, 76, 11, 35, 40, 33,
                   76, 41, 32, 76, 11, 35, 40, 33, 45, 66, 78, 76,
                   76, 76, 76, 76, 76, 76, 76, 12, 41, 41, 37, 76,
                   41, 40, 76, 39, 51, 76, 23, 41, 44, 37, 45, 64,
                   76, 51, 31, 76, 13, 35, 33, 34, 46, 51, 64, 76,
                   27, 40, 30, 76, 30, 31, 45, 42, 27, 35, 44, 67,
                   78, 76, 76, 76, 76, 76, 76, 76, 76, 14, 41, 46,
                   34, 35, 40, 33, 76, 28, 31, 45, 35, 30, 31, 76,
                   44, 31, 39, 27, 35, 40, 45, 63, 76, 18, 41, 47,
                   40, 30, 76, 46, 34, 31, 76, 30, 31, 29, 27, 51,
                   78, 76, 76, 76, 76, 76, 76, 76, 76, 15, 32, 76,
                   46, 34, 27, 46, 76, 29, 41, 38, 41, 45, 45, 27,
                   38, 76, 23, 44, 31, 29, 37, 64, 76, 28, 41, 47,
                   40, 30, 38, 31, 45, 45, 76, 27, 40, 30, 76, 28,
                   27, 44, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
                   20, 34, 31, 76, 38, 41, 40, 31, 76, 27, 40, 30,
                   76, 38, 31, 48, 31, 38, 76, 45, 27, 40, 30, 45,
                   76, 45, 46, 44, 31, 46, 29, 34, 76, 32, 27, 44,
                   76, 27, 49, 27, 51, 78, 76, 76, 76, 76, 76, 76,
                   76, 76]]),
               torch.tensor([[0, 0, 0, 1_069, 11]]),
               torch.tensor([[0, 0, 0, 1_069, 11]]),
           ]
           # fmt: on
           self.assertTrue(torch.allclose(tokens[0]   ,					EXPECTED_OUTPUT[0]))
           self.assertTrue(torch.allclose(tokens[1]   ,					EXPECTED_OUTPUT[1]))
           self.assertTrue(torch.allclose(tokens[2]   ,					EXPECTED_OUTPUT[2]))





    @require_torch
    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->Optional[int]:





           '''simple docstring'''

           import torch

           A__ 					=       JukeboxTokenizer.from_pretrained('''openai/jukebox-5b-lyrics''')
           A__ 					=       tokenizer(**self.metas)['''input_ids''']
           # fmt: off
           A__ 					=       [
               torch.tensor([[
                   0, 0, 0, 1_069, 11, -1, -1, -1, -1, 9, 77, 39,
                   31, 46, 77, 27, 77, 46, 44, 27, 48, 31, 38, 38,
                   31, 44, 77, 32, 44, 41, 39, 77, 27, 40, 77, 27,
                   40, 46, 35, 43, 47, 31, 77, 38, 27, 40, 30, 64,
                   79, 77, 77, 77, 77, 77, 77, 77, 77, 23, 34, 41,
                   77, 45, 27, 35, 30, 77, 72, 20, 49, 41, 77, 48,
                   27, 45, 46, 77, 27, 40, 30, 77, 46, 44, 47, 40,
                   37, 38, 31, 45, 45, 77, 38, 31, 33, 45, 77, 41,
                   32, 77, 45, 46, 41, 40, 31, 79, 77, 77, 77, 77,
                   77, 77, 77, 77, 19, 46, 27, 40, 30, 77, 35, 40,
                   77, 46, 34, 31, 77, 30, 31, 45, 31, 44, 46, 63,
                   77, 63, 77, 63, 77, 63, 77, 14, 31, 27, 44, 77,
                   46, 34, 31, 39, 64, 77, 41, 40, 77, 46, 34, 31,
                   77, 45, 27, 40, 30, 64, 79, 77, 77, 77, 77, 77,
                   77, 77, 77, 8, 27, 38, 32, 77, 45, 47, 40, 37,
                   77, 27, 77, 45, 34, 27, 46, 46, 31, 44, 31, 30,
                   77, 48, 35, 45, 27, 33, 31, 77, 38, 35, 31, 45,
                   64, 77, 49, 34, 41, 45, 31, 77, 32, 44, 41, 49,
                   40, 64, 79, 77, 77, 77, 77, 77, 77, 77, 77, 1,
                   40, 30, 77, 49, 44, 35, 40, 37, 38, 31, 30, 77,
                   38, 35, 42, 64, 77, 27, 40, 30, 77, 45, 40, 31,
                   31, 44, 77, 41, 32, 77, 29, 41, 38, 30, 77, 29,
                   41, 39, 39, 27, 40, 30, 64, 79, 77, 77, 77, 77,
                   77, 77, 77, 77, 20, 31, 38, 38, 77, 46, 34, 27,
                   46, 77, 35, 46, 45, 77, 45, 29, 47, 38, 42, 46,
                   41, 44, 77, 49, 31, 38, 38, 77, 46, 34, 41, 45,
                   31, 77, 42, 27, 45, 45, 35, 41, 40, 45, 77, 44,
                   31, 27, 30, 79, 77, 77, 77, 77, 77, 77, 77, 77,
                   23, 34, 35, 29, 34, 77, 51, 31, 46, 77, 45, 47,
                   44, 48, 35, 48, 31, 64, 77, 45, 46, 27, 39, 42,
                   31, 30, 77, 41, 40, 77, 46, 34, 31, 45, 31, 77,
                   38, 35, 32, 31, 38, 31, 45, 45, 77, 46, 34, 35,
                   40, 33, 45, 64, 79, 77, 77, 77, 77, 77, 77, 77,
                   77, 20, 34, 31, 77, 34, 27, 40, 30, 77, 46, 34,
                   27, 46, 77, 39, 41, 29, 37, 31, 30, 77, 46, 34,
                   31, 39, 64, 77, 27, 40, 30, 77, 46, 34, 31, 77,
                   34, 31, 27, 44, 46, 77, 46, 34, 27, 46, 77, 32,
                   31, 30, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77,
                   1, 40, 30, 77, 41, 40, 77, 46, 34, 31, 77, 42,
                   31, 30, 31, 45, 46, 27, 38, 64, 77, 46, 34, 31,
                   45, 31, 77, 49, 41, 44, 30, 45, 77, 27, 42, 42,
                   31, 27, 44, 65, 79, 77, 77, 77, 77, 77, 77, 77,
                   77, 13, 51, 77, 40, 27, 39, 31, 77, 35, 45, 77,
                   15, 52, 51, 39, 27, 40, 30, 35, 27, 45, 64, 77,
                   11, 35, 40, 33, 77, 41, 32, 77, 11, 35, 40, 33,
                   45, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77, 12,
                   41, 41, 37, 77, 41, 40, 77, 39, 51, 77, 23, 41,
                   44, 37, 45, 64, 77, 51, 31, 77, 13, 35, 33, 34,
                   46, 51, 64, 77, 27, 40, 30, 77, 30, 31, 45, 42,
                   27, 35, 44, 67, 79, 77, 77, 77, 77, 77, 77, 77,
                   77, 14, 41, 46, 34, 35, 40, 33, 77, 28, 31, 45,
                   35, 30, 31, 77, 44, 31, 39, 27, 35, 40, 45, 63,
                   77, 18, 41, 47, 40, 30, 77, 46, 34, 31, 77, 30,
                   31, 29, 27, 51, 79, 77, 77, 77, 77, 77, 77, 77,
                   77, 15, 32, 77, 46, 34, 27, 46, 77, 29, 41, 38,
                   41, 45, 45, 27, 38, 77, 23, 44, 31, 29, 37, 64,
                   77, 28, 41, 47, 40, 30, 38, 31, 45, 45, 77, 27,
                   40, 30, 77, 28, 27, 44, 31, 79, 77, 77, 77, 77,
                   77, 77, 77, 77, 20, 34, 31, 77, 38, 41, 40, 31,
                   77, 27, 40, 30, 77, 38, 31, 48, 31, 38, 77, 45,
                   27, 40, 30, 45, 77, 45, 46, 44, 31, 46, 29, 34,
                   77, 32, 27, 44, 77, 27, 49, 27, 51, 79, 77, 77,
                   77, 77, 77, 77, 77, 77]]),
               torch.tensor([[0, 0, 0, 1_069, 11, -1, -1, -1, -1]]),
               torch.tensor([[0, 0, 0, 1_069, 11, -1, -1, -1, -1]]),
           ]
           # fmt: on
           self.assertTrue(torch.allclose(tokens[0]   ,					EXPECTED_OUTPUT[0]))
           self.assertTrue(torch.allclose(tokens[1]   ,					EXPECTED_OUTPUT[1]))
           self.assertTrue(torch.allclose(tokens[2]   ,					EXPECTED_OUTPUT[2]))


