


import json
import os
from functools import lru_cache
from typing import Dict, List, Optional, Tuple, Union

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...tokenization_utils_base import BatchEncoding, EncodedInput
from ...utils import PaddingStrategy, logging


__lowerCAmelCase      :					Tuple    							=logging.get_logger(__name__)


__lowerCAmelCase      :					Optional[int]    							={'vocab_file': 'vocab.json', 'merges_file': 'merges.txt'}

# See all LED models at https://huggingface.co/models?filter=LED
__lowerCAmelCase      :					Optional[int]    							={
    'vocab_file': {
        'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/vocab.json',
    },
    'merges_file': {
        'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/merges.txt',
    },
    'tokenizer_file': {
        'allenai/led-base-16384': 'https://huggingface.co/allenai/led-base-16384/resolve/main/tokenizer.json',
    },
}

__lowerCAmelCase      :					Optional[Any]    							={
    'allenai/led-base-16384': 1_6_3_8_4,
}


@lru_cache()
# Copied from transformers.models.bart.tokenization_bart.bytes_to_unicode
def 			_UpperCamelCase				( ):
						__SCREAMING_SNAKE_CASE		:	Any						     =						(
						    list(range(ord('''!''' )       ,			ord('''~''' ) + 1 ) ) + list(range(ord('''¡''' )       ,			ord('''¬''' ) + 1 ) ) + list(range(ord('''®''' )       ,			ord('''ÿ''' ) + 1 ) )
						)
						__SCREAMING_SNAKE_CASE		:	int						     =						bs[:]
						__SCREAMING_SNAKE_CASE		:	List[str]						     =						0
						for b in range(2**8 ):
												if b not in bs:
																		bs.append(lowercase__ )
																		cs.append(2**8 + n )
																		n += 1
						__SCREAMING_SNAKE_CASE		:	Dict						     =						[chr(lowercase__ ) for n in cs]
						return dict(zip(lowercase__       ,			lowercase__ ) )


def 			_UpperCamelCase				( lowercase__ ):
						__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						set()
						__SCREAMING_SNAKE_CASE		:	List[Any]						     =						word[0]
						for char in word[1:]:
												pairs.add((prev_char, char) )
												__SCREAMING_SNAKE_CASE		:	List[Any]						     =						char
						return pairs
class 			_lowercase (     __a			):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__							:     Optional[int]         =						VOCAB_FILES_NAMES
				SCREAMING_SNAKE_CASE__							:     Any         =						PRETRAINED_VOCAB_FILES_MAP
				SCREAMING_SNAKE_CASE__							:     str         =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				SCREAMING_SNAKE_CASE__							:     Any         =						['''input_ids''', '''attention_mask''']


				def __init__(			self     :List[Any]     ,     lowerCAmelCase__     :int     ,     lowerCAmelCase__     :List[Any]     ,     lowerCAmelCase__     :Union[str, Any]="replace"     ,     lowerCAmelCase__     :int="<s>"     ,     lowerCAmelCase__     :List[Any]="</s>"     ,     lowerCAmelCase__     :List[Any]="</s>"     ,     lowerCAmelCase__     :Any="<s>"     ,     lowerCAmelCase__     :int="<unk>"     ,     lowerCAmelCase__     :List[str]="<pad>"     ,     lowerCAmelCase__     :str="<mask>"     ,     lowerCAmelCase__     :Optional[Any]=False     ,     **lowerCAmelCase__     :Optional[int]     ,     )						->     int:
										__SCREAMING_SNAKE_CASE		:	int						     =						AddedToken(_lowercase     ,     lstrip=_lowercase     ,     rstrip=_lowercase      ) if isinstance(_lowercase     ,     _lowercase      ) else bos_token
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						AddedToken(_lowercase     ,     lstrip=_lowercase     ,     rstrip=_lowercase      ) if isinstance(_lowercase     ,     _lowercase      ) else eos_token
										__SCREAMING_SNAKE_CASE		:	Tuple						     =						AddedToken(_lowercase     ,     lstrip=_lowercase     ,     rstrip=_lowercase      ) if isinstance(_lowercase     ,     _lowercase      ) else sep_token
										__SCREAMING_SNAKE_CASE		:	Any						     =						AddedToken(_lowercase     ,     lstrip=_lowercase     ,     rstrip=_lowercase      ) if isinstance(_lowercase     ,     _lowercase      ) else cls_token
										__SCREAMING_SNAKE_CASE		:	Tuple						     =						AddedToken(_lowercase     ,     lstrip=_lowercase     ,     rstrip=_lowercase      ) if isinstance(_lowercase     ,     _lowercase      ) else unk_token
										__SCREAMING_SNAKE_CASE		:	List[str]						     =						AddedToken(_lowercase     ,     lstrip=_lowercase     ,     rstrip=_lowercase      ) if isinstance(_lowercase     ,     _lowercase      ) else pad_token

										# Mask token behave like a normal word, i.e. include the space before it
										__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						AddedToken(_lowercase     ,     lstrip=_lowercase     ,     rstrip=_lowercase      ) if isinstance(_lowercase     ,     _lowercase      ) else mask_token

										super().__init__(
										    errors=_lowercase     ,     bos_token=_lowercase     ,     eos_token=_lowercase     ,     unk_token=_lowercase     ,     sep_token=_lowercase     ,     cls_token=_lowercase     ,     pad_token=_lowercase     ,     mask_token=_lowercase     ,     add_prefix_space=_lowercase     ,     **_lowercase     ,     )

										with open(_lowercase     ,     encoding='''utf-8'''      ) as vocab_handle:
																__SCREAMING_SNAKE_CASE		:	Any						     =						json.load(_lowercase      )
										__SCREAMING_SNAKE_CASE		:	Dict						     =						{v: k for k, v in self.encoder.items()}
										__SCREAMING_SNAKE_CASE		:	Dict						     =						errors  # how to handle errors in decoding
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						bytes_to_unicode()
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						{v: k for k, v in self.byte_encoder.items()}
										with open(_lowercase     ,     encoding='''utf-8'''      ) as merges_handle:
																__SCREAMING_SNAKE_CASE		:	List[Any]						     =						merges_handle.read().split('''\n'''      )[1:-1]
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						[tuple(merge.split()      ) for merge in bpe_merges]
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						dict(zip(_lowercase     ,     range(len(_lowercase      )      )      )      )
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						{}
										__SCREAMING_SNAKE_CASE		:	List[str]						     =						add_prefix_space

										# Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
										__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						re.compile(r'''\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+'''      )


				@property
				# Copied from transformers.models.bart.tokenization_bart.BartTokenizer.vocab_size
				def 						__magic_name__(			self     :Optional[Any]      )						->     int:
										return len(self.encoder      )


				def 						__magic_name__(			self     :int      )						->     Tuple:
										return dict(self.encoder     ,     **self.added_tokens_encoder      )


				def 						__magic_name__(			self     :str     ,     lowerCAmelCase__     :Tuple      )						->     Optional[int]:



										if token in self.cache:
																return self.cache[token]
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						tuple(_lowercase      )
										__SCREAMING_SNAKE_CASE		:	List[str]						     =						get_pairs(_lowercase      )

										if not pairs:
																return token

										while True:
																__SCREAMING_SNAKE_CASE		:	str						     =						min(_lowercase     ,     key=lambda lowerCAmelCase__      : self.bpe_ranks.get(_lowercase     ,     float('''inf'''      )      )      )
																if bigram not in self.bpe_ranks:
																						break
																__SCREAMING_SNAKE_CASE		:	Any						     =						bigram
																__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						[]
																__SCREAMING_SNAKE_CASE		:	Dict						     =						0
																while i < len(_lowercase      ):
																						try:
																												__SCREAMING_SNAKE_CASE		:	List[str]						     =						word.index(_lowercase     ,     _lowercase      )
																						except ValueError:
																												new_word.extend(word[i:]      )
																												break
																						else:
																												new_word.extend(word[i:j]      )
																												__SCREAMING_SNAKE_CASE		:	str						     =						j

																						if word[i] == first and i < len(_lowercase      ) - 1 and word[i + 1] == second:
																												new_word.append(first + second      )
																												i += 2
																						else:
																												new_word.append(word[i]      )
																												i += 1
																__SCREAMING_SNAKE_CASE		:	List[str]						     =						tuple(_lowercase      )
																__SCREAMING_SNAKE_CASE		:	Tuple						     =						new_word
																if len(_lowercase      ) == 1:
																						break
																else:
																						__SCREAMING_SNAKE_CASE		:	Any						     =						get_pairs(_lowercase      )
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						''' '''.join(_lowercase      )
										__SCREAMING_SNAKE_CASE		:	int						     =						word
										return word


				def 						__magic_name__(			self     :Dict     ,     lowerCAmelCase__     :Optional[int]      )						->     Union[str, Any]:
										__SCREAMING_SNAKE_CASE		:	Any						     =						[]
										for token in re.findall(self.pat     ,     _lowercase      ):
																__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						''''''.join(
																    self.byte_encoder[b] for b in token.encode('''utf-8'''      )      )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
																bpe_tokens.extend(bpe_token for bpe_token in self.bpe(_lowercase      ).split(''' '''      )      )
										return bpe_tokens


				def 						__magic_name__(			self     :Union[str, Any]     ,     lowerCAmelCase__     :Dict      )						->     Tuple:
										return self.encoder.get(_lowercase     ,     self.encoder.get(self.unk_token      )      )


				def 						__magic_name__(			self     :Union[str, Any]     ,     lowerCAmelCase__     :Any      )						->     int:
										return self.decoder.get(_lowercase      )


				def 						__magic_name__(			self     :str     ,     lowerCAmelCase__     :Union[str, Any]      )						->     Any:
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						''''''.join(_lowercase      )
										__SCREAMING_SNAKE_CASE		:	List[str]						     =						bytearray([self.byte_decoder[c] for c in text]      ).decode('''utf-8'''     ,     errors=self.errors      )
										return text


				def 						__magic_name__(			self     :Dict     ,     lowerCAmelCase__     :str     ,     lowerCAmelCase__     :Optional[str] = None      )						->     Any:



										if not os.path.isdir(_lowercase      ):
																logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''      )
																return
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						os.path.join(
										    _lowercase     ,     (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']      )
										__SCREAMING_SNAKE_CASE		:	Dict						     =						os.path.join(
										    _lowercase     ,     (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']      )

										with open(_lowercase     ,     '''w'''     ,     encoding='''utf-8'''      ) as f:
																f.write(json.dumps(self.encoder     ,     indent=2     ,     sort_keys=_lowercase     ,     ensure_ascii=_lowercase      ) + '''\n'''      )

										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						0
										with open(_lowercase     ,     '''w'''     ,     encoding='''utf-8'''      ) as writer:
																writer.write('''#version: 0.2\n'''      )
																for bpe_tokens, token_index in sorted(self.bpe_ranks.items()     ,     key=lambda lowerCAmelCase__      : kv[1]      ):
																						if index != token_index:
																												logger.warning(
																												    f'''Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'''
																												    ''' Please check that the tokenizer is not corrupted!'''      )
																												__SCREAMING_SNAKE_CASE		:	int						     =						token_index
																						writer.write(''' '''.join(_lowercase      ) + '''\n'''      )
																						index += 1

										return vocab_file, merge_file


				def 						__magic_name__(			self     :Dict     ,     lowerCAmelCase__     :List[int]     ,     lowerCAmelCase__     :Optional[List[int]] = None      )						->     List[str]:



										if token_ids_a is None:
																return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
										__SCREAMING_SNAKE_CASE		:	int						     =						[self.cls_token_id]
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						[self.sep_token_id]
										return cls + token_ids_a + sep + sep + token_ids_a + sep


				def 						__magic_name__(			self     :Any     ,     lowerCAmelCase__     :List[int]     ,     lowerCAmelCase__     :Optional[List[int]] = None     ,     lowerCAmelCase__     :bool = False      )						->     Tuple:



										if already_has_special_tokens:
																return super().get_special_tokens_mask(
																    token_ids_a=_lowercase     ,     token_ids_a=_lowercase     ,     already_has_special_tokens=_lowercase      )

										if token_ids_a is None:
																return [1] + ([0] * len(_lowercase      )) + [1]
										return [1] + ([0] * len(_lowercase      )) + [1, 1] + ([0] * len(_lowercase      )) + [1]


				def 						__magic_name__(			self     :Dict     ,     lowerCAmelCase__     :List[int]     ,     lowerCAmelCase__     :Optional[List[int]] = None      )						->     Any:
										__SCREAMING_SNAKE_CASE		:	int						     =						[self.sep_token_id]
										__SCREAMING_SNAKE_CASE		:	Dict						     =						[self.cls_token_id]

										if token_ids_a is None:
																return len(cls + token_ids_a + sep      ) * [0]
										return len(cls + token_ids_a + sep + sep + token_ids_a + sep      ) * [0]


				def 						__magic_name__(			self     :str     ,     lowerCAmelCase__     :Any     ,     lowerCAmelCase__     :Dict=False     ,     **lowerCAmelCase__     :Union[str, Any]      )						->     Optional[Any]:
										__SCREAMING_SNAKE_CASE		:	int						     =						kwargs.pop('''add_prefix_space'''     ,     self.add_prefix_space      )
										if (is_split_into_words or add_prefix_space) and (len(_lowercase      ) > 0 and not text[0].isspace()):
																__SCREAMING_SNAKE_CASE		:	Any						     =						''' ''' + text
										return (text, kwargs)



				def 						__magic_name__(			self     :List[Any]     ,     lowerCAmelCase__     :Union[Dict[str, EncodedInput], BatchEncoding]     ,     lowerCAmelCase__     :Optional[int] = None     ,     lowerCAmelCase__     :PaddingStrategy = PaddingStrategy.DO_NOT_PAD     ,     lowerCAmelCase__     :Optional[int] = None     ,     lowerCAmelCase__     :Optional[bool] = None     ,     )						->     Any:
										__SCREAMING_SNAKE_CASE		:	List[str]						     =						super()._pad(
										    encoded_inputs=_lowercase     ,     max_length=_lowercase     ,     padding_strategy=_lowercase     ,     pad_to_multiple_of=_lowercase     ,     return_attention_mask=_lowercase     ,     )

										# Load from model defaults
										if return_attention_mask is None:
																__SCREAMING_SNAKE_CASE		:	int						     =						'''attention_mask''' in self.model_input_names

										if return_attention_mask and "global_attention_mask" in encoded_inputs:
																__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						encoded_inputs[self.model_input_names[0]]
																# `global_attention_mask` need to have the same length as other (sequential) inputs.
																__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						len(encoded_inputs['''global_attention_mask''']      ) != len(_lowercase      )

																if needs_to_be_padded:
																						__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						len(_lowercase      ) - len(encoded_inputs['''global_attention_mask''']      )

																						if self.padding_side == "right":
																												# Use `-1` since `0` in `global_attention_mask` means `local attention` instead of `not to attend`
																												__SCREAMING_SNAKE_CASE		:	List[str]						     =						(
																												    encoded_inputs['''global_attention_mask'''] + [-1] * difference
																												)
																						elif self.padding_side == "left":
																												__SCREAMING_SNAKE_CASE		:	Dict						     =						[-1] * difference + encoded_inputs[
																												    '''global_attention_mask'''
																												]
																						else:
																												raise ValueError('''Invalid padding strategy:''' + str(self.padding_side      )      )

										return encoded_inputs

import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax.numpy as jnp

  from transformers.models.albert.modeling_flax_albert import (
      FlaxAlbertForMaskedLM,
      FlaxAlbertForMultipleChoice,
      FlaxAlbertForPreTraining,
      FlaxAlbertForQuestionAnswering,
      FlaxAlbertForSequenceClassification,
      FlaxAlbertForTokenClassification,
      FlaxAlbertModel,
  )




class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




   def __init__(						self       :	Any		,	_lowercase       :	Optional[int]		,	_lowercase       :	Tuple=13		,	_lowercase       :	str=7		,	_lowercase       :	List[Any]=True		,	_lowercase       :	Optional[int]=True		,	_lowercase       :	str=True		,	_lowercase       :	Optional[int]=True		,	_lowercase       :	Dict=99		,	_lowercase       :	List[Any]=32		,	_lowercase       :	List[str]=5		,	_lowercase       :	str=4		,	_lowercase       :	int=37		,	_lowercase       :	List[str]="gelu"		,	_lowercase       :	Any=0.1		,	_lowercase       :	Optional[int]=0.1		,	_lowercase       :	Dict=512		,	_lowercase       :	int=16		,	_lowercase       :	Optional[Any]=2		,	_lowercase       :	Dict=0.02		,	_lowercase       :	List[Any]=4		,	):



    """simple docstring"""



    _UpperCamelCase:  List[str]    						=			parent
    _UpperCamelCase:  int    						=			batch_size
    _UpperCamelCase:  List[str]    						=			seq_length
    _UpperCamelCase:  Optional[int]    						=			is_training
    _UpperCamelCase:  Optional[Any]    						=			use_attention_mask
    _UpperCamelCase:  Any    						=			use_token_type_ids
    _UpperCamelCase:  List[str]    						=			use_labels
    _UpperCamelCase:  Optional[int]    						=			vocab_size
    _UpperCamelCase:  List[str]    						=			hidden_size
    _UpperCamelCase:  Union[str, Any]    						=			num_hidden_layers
    _UpperCamelCase:  Any    						=			num_attention_heads
    _UpperCamelCase:  List[str]    						=			intermediate_size
    _UpperCamelCase:  Union[str, Any]    						=			hidden_act
    _UpperCamelCase:  Dict    						=			hidden_dropout_prob
    _UpperCamelCase:  List[str]    						=			attention_probs_dropout_prob
    _UpperCamelCase:  str    						=			max_position_embeddings
    _UpperCamelCase:  Dict    						=			type_vocab_size
    _UpperCamelCase:  Tuple    						=			type_sequence_label_size
    _UpperCamelCase:  List[Any]    						=			initializer_range
    _UpperCamelCase:  str    						=			num_choices



   def 	lowerCAmelCase    (						self       :	str				):



    """simple docstring"""



    _UpperCamelCase:  str    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size				)

    _UpperCamelCase:  Any    						=			None
    if self.use_attention_mask:
     _UpperCamelCase:  List[str]    						=			random_attention_mask([self.batch_size, self.seq_length]				)

    _UpperCamelCase:  Optional[Any]    						=			None
    if self.use_token_type_ids:
     _UpperCamelCase:  Optional[Any]    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.type_vocab_size				)

    _UpperCamelCase:  Tuple    						=			AlbertConfig(
        vocab_size=self.vocab_size		,	hidden_size=self.hidden_size		,	num_hidden_layers=self.num_hidden_layers		,	num_attention_heads=self.num_attention_heads		,	intermediate_size=self.intermediate_size		,	hidden_act=self.hidden_act		,	hidden_dropout_prob=self.hidden_dropout_prob		,	attention_probs_dropout_prob=self.attention_probs_dropout_prob		,	max_position_embeddings=self.max_position_embeddings		,	type_vocab_size=self.type_vocab_size		,	is_decoder=_lowercase		,	initializer_range=self.initializer_range		,	)

    return config, input_ids, token_type_ids, attention_mask



   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    _UpperCamelCase:  List[Any]    						=			self.prepare_config_and_inputs()
    _UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase:  int    						=			config_and_inputs
    _UpperCamelCase:  int    						=			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
    return config, inputs_dict




@require_flax
class 			__magic_name__	(		__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     Any	      = (
       (
           FlaxAlbertModel,
           FlaxAlbertForPreTraining,
           FlaxAlbertForMaskedLM,
           FlaxAlbertForMultipleChoice,
           FlaxAlbertForQuestionAnswering,
           FlaxAlbertForSequenceClassification,
           FlaxAlbertForTokenClassification,
           FlaxAlbertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )



   def 	lowerCAmelCase    (						self       :	Union[str, Any]				):



    """simple docstring"""



    _UpperCamelCase:  Optional[int]    						=			FlaxAlbertModelTester(self				)



   @slow
   def 	lowerCAmelCase    (						self       :	Union[str, Any]				):



    """simple docstring"""



    for model_class_name in self.all_model_classes:
     _UpperCamelCase:  int    						=			model_class_name.from_pretrained('''albert-base-v2'''				)
     _UpperCamelCase:  List[Any]    						=			model(np.ones((1, 1)				)				)
     self.assertIsNotNone(_lowercase				)




@require_flax
class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




   @slow
   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    _UpperCamelCase:  Optional[int]    						=			FlaxAlbertModel.from_pretrained('''albert-base-v2'''				)
    _UpperCamelCase:  List[Any]    						=			np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]				)
    _UpperCamelCase:  Optional[int]    						=			np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]				)
    _UpperCamelCase:  Optional[int]    						=			model(_lowercase		,	attention_mask=_lowercase				)[0]
    _UpperCamelCase:  Tuple    						=			(1, 11, 768)
    self.assertEqual(output.shape		,	_lowercase				)
    _UpperCamelCase:  Tuple    						=			np.array(
        [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]				)

    self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]		,	_lowercase		,	atol=1E-4				)				)