





"""simple docstring"""





import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
    import jax
    import jax.numpy as jnp
    import numpy as np

    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
    import torch





@require_flax
class 			snake_case		:

      def __init__(			self   :Optional[Any]						,							_lowerCamelCase   :Optional[Any]						,							_lowerCamelCase   :Dict=1_4						,							_lowerCamelCase   :Any=7						,							_lowerCamelCase   :List[str]=True						,							_lowerCamelCase   :str=True						,							_lowerCamelCase   :Union[str, Any]=True						,							_lowerCamelCase   :Optional[int]=9_9						,							_lowerCamelCase   :List[str]=3_2						,							_lowerCamelCase   :List[str]=5						,							_lowerCamelCase   :Tuple=4						,							_lowerCamelCase   :Tuple=3_7						,							_lowerCamelCase   :Optional[int]="gelu"						,							_lowerCamelCase   :Dict=0.1						,							_lowerCamelCase   :Optional[int]=0.1						,							_lowerCamelCase   :int=5_1_2						,							_lowerCamelCase   :str=0.0_2						,							_lowerCamelCase   :Dict=None						,							):
       __SCREAMING_SNAKE_CASE :			Optional[Any]												=       parent
       __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       batch_size
       __SCREAMING_SNAKE_CASE :			Dict												=       seq_length
       __SCREAMING_SNAKE_CASE :			Dict												=       is_training
       __SCREAMING_SNAKE_CASE :			Optional[Any]												=       use_input_mask
       __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       use_labels
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       vocab_size
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       d_model
       __SCREAMING_SNAKE_CASE :			Dict												=       num_hidden_layers
       __SCREAMING_SNAKE_CASE :			List[Any]												=       num_attention_heads
       __SCREAMING_SNAKE_CASE :			int												=       ffn_dim
       __SCREAMING_SNAKE_CASE :			Dict												=       activation_function
       __SCREAMING_SNAKE_CASE :			int												=       activation_dropout
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       attention_dropout
       __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       max_position_embeddings
       __SCREAMING_SNAKE_CASE :			List[str]												=       initializer_range
       __SCREAMING_SNAKE_CASE :			str												=       None
       __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       0
       __SCREAMING_SNAKE_CASE :			Any												=       2
       __SCREAMING_SNAKE_CASE :			Tuple												=       1

      def 						SCREAMING_SNAKE_CASE_			(			self   :Tuple						):
       __SCREAMING_SNAKE_CASE :			Any												=       np.clip(ids_tensor([self.batch_size, self.seq_length]						,							self.vocab_size						)						,							3						,							self.vocab_size						)

       __SCREAMING_SNAKE_CASE :			Optional[Any]												=       None
       if self.use_input_mask:
        __SCREAMING_SNAKE_CASE :			Optional[Any]												=       random_attention_mask([self.batch_size, self.seq_length]						)

       __SCREAMING_SNAKE_CASE :			str												=       XGLMConfig(
           vocab_size=self.vocab_size						,							d_model=self.hidden_size						,							num_layers=self.num_hidden_layers						,							attention_heads=self.num_attention_heads						,							ffn_dim=self.ffn_dim						,							activation_function=self.activation_function						,							activation_dropout=self.activation_dropout						,							attention_dropout=self.attention_dropout						,							max_position_embeddings=self.max_position_embeddings						,							initializer_range=self.initializer_range						,							use_cache=_lowerCamelCase						,							bos_token_id=self.bos_token_id						,							eos_token_id=self.eos_token_id						,							pad_token_id=self.pad_token_id						,							)

       return (config, input_ids, input_mask)

      def 						SCREAMING_SNAKE_CASE_			(			self   :Optional[int]						):
       __SCREAMING_SNAKE_CASE :			int												=       self.prepare_config_and_inputs()
       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			int												=       config_and_inputs
       __SCREAMING_SNAKE_CASE :			int												=       {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
       return config, inputs_dict

      def 						SCREAMING_SNAKE_CASE_			(			self   :Any						):
       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			List[str]												=       self.prepare_config_and_inputs()

       __SCREAMING_SNAKE_CASE :			List[str]												=       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
       __SCREAMING_SNAKE_CASE :			int												=       ids_tensor([self.batch_size, self.seq_length]						,							vocab_size=2						)

       return (
           config,
           input_ids,
           attention_mask,
           encoder_hidden_states,
           encoder_attention_mask,
       )

      def 						SCREAMING_SNAKE_CASE_			(			self   :Tuple						,							_lowerCamelCase   :Dict						,							_lowerCamelCase   :Optional[Any]						,							_lowerCamelCase   :str						,							_lowerCamelCase   :List[Any]						):
       __SCREAMING_SNAKE_CASE :			Dict												=       2_0
       __SCREAMING_SNAKE_CASE :			int												=       model_class_name(_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE :			Optional[int]												=       model.init_cache(input_ids.shape[0]						,							_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE :			List[str]												=       jnp.ones((input_ids.shape[0], max_decoder_length)						,							dtype='''i4'''						)

       __SCREAMING_SNAKE_CASE :			List[Any]												=       jnp.broadcast_to(
           jnp.arange(input_ids.shape[-1] - 1						)[None, :]						,							(input_ids.shape[0], input_ids.shape[-1] - 1)						)
       __SCREAMING_SNAKE_CASE :			int												=       model(
           input_ids[:, :-1]						,							attention_mask=_lowerCamelCase						,							past_key_values=_lowerCamelCase						,							position_ids=_lowerCamelCase						,							)

       __SCREAMING_SNAKE_CASE :			Any												=       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,							dtype='''i4'''						)
       __SCREAMING_SNAKE_CASE :			Tuple												=       model(
           input_ids[:, -1:]						,							attention_mask=_lowerCamelCase						,							past_key_values=outputs_cache.past_key_values						,							position_ids=_lowerCamelCase						,							)

       __SCREAMING_SNAKE_CASE :			List[str]												=       model(_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE :			Any												=       np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])						)						)
       self.parent.assertTrue(diff < 1e-3						,							msg=f'''Max diff is {diff}'''						)



      def 						SCREAMING_SNAKE_CASE_			(			self   :str						,							_lowerCamelCase   :Union[str, Any]						,							_lowerCamelCase   :int						,							_lowerCamelCase   :int						,							_lowerCamelCase   :Optional[int]						):
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       2_0
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       model_class_name(_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE :			Any												=       jnp.concatenate(
           [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])						)]						,							axis=-1						,							)

       __SCREAMING_SNAKE_CASE :			List[str]												=       model.init_cache(input_ids.shape[0]						,							_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE :			List[Any]												=       jnp.broadcast_to(
           jnp.arange(input_ids.shape[-1] - 1						)[None, :]						,							(input_ids.shape[0], input_ids.shape[-1] - 1)						)

       __SCREAMING_SNAKE_CASE :			Dict												=       model(
           input_ids[:, :-1]						,							attention_mask=_lowerCamelCase						,							past_key_values=_lowerCamelCase						,							position_ids=_lowerCamelCase						,							)
       __SCREAMING_SNAKE_CASE :			int												=       jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,							dtype='''i4'''						)
       __SCREAMING_SNAKE_CASE :			Tuple												=       model(
           input_ids[:, -1:]						,							past_key_values=outputs_cache.past_key_values						,							attention_mask=_lowerCamelCase						,							position_ids=_lowerCamelCase						,							)

       __SCREAMING_SNAKE_CASE :			Tuple												=       model(_lowerCamelCase						,							attention_mask=_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])						)						)
       self.parent.assertTrue(diff < 1e-3						,							msg=f'''Max diff is {diff}'''						)







@require_sentencepiece
@require_flax
class 			snake_case		(   __UpperCAmelCase ,     __UpperCAmelCase ,     unittest.TestCase   ):
      lowerCamelCase__						= (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
      lowerCamelCase__						= (FlaxXGLMForCausalLM,) if is_flax_available() else ()

      def 						SCREAMING_SNAKE_CASE_			(			self   :int						):
       __SCREAMING_SNAKE_CASE :			str												=       FlaxXGLMModelTester(self						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :Tuple						):
       for model_class_name in self.all_model_classes:
        __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_use_cache_forward(_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						)

      def 						SCREAMING_SNAKE_CASE_			(			self   :str						):
       for model_class_name in self.all_model_classes:
        __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			Optional[int]												=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_use_cache_forward_with_attn_mask(
            _lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase						)

      @slow
      def 						SCREAMING_SNAKE_CASE_			(			self   :Optional[int]						):
       __SCREAMING_SNAKE_CASE :			int												=       XGLMTokenizer.from_pretrained('''XGLM'''						,							padding_side='''left'''						)
       __SCREAMING_SNAKE_CASE :			Optional[int]												=       tokenizer(['''Hello this is a long string''', '''Hey''']						,							return_tensors='''np'''						,							padding=_lowerCamelCase						,							truncation=_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE :			Dict												=       FlaxXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''						)
       __SCREAMING_SNAKE_CASE :			Any												=       1
       __SCREAMING_SNAKE_CASE :			int												=       False

       __SCREAMING_SNAKE_CASE :			int												=       jax.jit(model.generate						)

       __SCREAMING_SNAKE_CASE :			str												=       jit_generate(inputs['''input_ids''']						,							attention_mask=inputs['''attention_mask''']						).sequences

       __SCREAMING_SNAKE_CASE :			List[str]												=       tokenizer.batch_decode(_lowerCamelCase						,							skip_special_tokens=_lowerCamelCase						)

       __SCREAMING_SNAKE_CASE :			int												=       [
           '''Hello this is a long string of questions, but I\'m not sure if I\'m''',
           '''Hey, I\'m a newbie to the forum and I\'''',
       ]

       self.assertListEqual(_lowerCamelCase						,							_lowerCamelCase						)

      @is_pt_flax_cross_test
      def 						SCREAMING_SNAKE_CASE_			(			self   :List[Any]						):
       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			List[Any]												=       self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__						):
         # prepare inputs
         __SCREAMING_SNAKE_CASE :			str												=       self._prepare_for_class(_lowerCamelCase						,							_lowerCamelCase						)
         __SCREAMING_SNAKE_CASE :			Optional[int]												=       {k: torch.tensor(v.tolist()						) for k, v in prepared_inputs_dict.items()}

         # load corresponding PyTorch class
         __SCREAMING_SNAKE_CASE :			List[str]												=       model_class.__name__[4:]  # Skip the "Flax" at the beginning
         __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       getattr(_lowerCamelCase						,							_lowerCamelCase						)

         __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			List[Any]												=       pt_inputs['''input_ids'''].shape
         __SCREAMING_SNAKE_CASE :			Any												=       np.random.randint(0						,							seq_length - 1						,							size=(batch_size,)						)
         for batch_idx, start_index in enumerate(_lowerCamelCase						):
          __SCREAMING_SNAKE_CASE :			int												=       0
          __SCREAMING_SNAKE_CASE :			str												=       1
          __SCREAMING_SNAKE_CASE :			Tuple												=       0
          __SCREAMING_SNAKE_CASE :			str												=       1
         __SCREAMING_SNAKE_CASE :			List[Any]												=       pt_model_class(_lowerCamelCase						).eval()
         # Flax models don't use the `use_cache` option and cache is not returned as a default.
         # So we disable `use_cache` here for PyTorch model.
         __SCREAMING_SNAKE_CASE :			str												=       False
         __SCREAMING_SNAKE_CASE :			List[str]												=       model_class(_lowerCamelCase						,							dtype=jnp.floataa						)
         __SCREAMING_SNAKE_CASE :			str												=       convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,							_lowerCamelCase						)
         __SCREAMING_SNAKE_CASE :			List[str]												=       fx_state

         with torch.no_grad():
          __SCREAMING_SNAKE_CASE :			Dict												=       pt_model(**_lowerCamelCase						).to_tuple()

         __SCREAMING_SNAKE_CASE :			Optional[Any]												=       fx_model(**_lowerCamelCase						).to_tuple()
         self.assertEqual(len(_lowerCamelCase						)						,							len(_lowerCamelCase						)						,							'''Output lengths differ between Flax and PyTorch'''						)
         for fx_output, pt_output in zip(_lowerCamelCase						,							_lowerCamelCase						):
          self.assert_almost_equals(fx_output[:, -1]						,							pt_output[:, -1].numpy()						,							4e-2						)

         with tempfile.TemporaryDirectory() as tmpdirname:
          pt_model.save_pretrained(_lowerCamelCase						)
          __SCREAMING_SNAKE_CASE :			Any												=       model_class.from_pretrained(_lowerCamelCase						,							from_pt=_lowerCamelCase						)

         __SCREAMING_SNAKE_CASE :			Dict												=       fx_model_loaded(**_lowerCamelCase						).to_tuple()
         self.assertEqual(
             len(_lowerCamelCase						)						,							len(_lowerCamelCase						)						,							'''Output lengths differ between Flax and PyTorch'''						)
         for fx_output_loaded, pt_output in zip(_lowerCamelCase						,							_lowerCamelCase						):
          self.assert_almost_equals(fx_output_loaded[:, -1]						,							pt_output[:, -1].numpy()						,							4e-2						)

      @is_pt_flax_cross_test
      def 						SCREAMING_SNAKE_CASE_			(			self   :Dict						):
       __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			Any												=       self.model_tester.prepare_config_and_inputs_for_common()
       for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__						):
         # prepare inputs
         __SCREAMING_SNAKE_CASE :			List[str]												=       self._prepare_for_class(_lowerCamelCase						,							_lowerCamelCase						)
         __SCREAMING_SNAKE_CASE :			Dict												=       {k: torch.tensor(v.tolist()						) for k, v in prepared_inputs_dict.items()}

         # load corresponding PyTorch class
         __SCREAMING_SNAKE_CASE :			Any												=       model_class.__name__[4:]  # Skip the "Flax" at the beginning
         __SCREAMING_SNAKE_CASE :			int												=       getattr(_lowerCamelCase						,							_lowerCamelCase						)

         __SCREAMING_SNAKE_CASE :			str												=       pt_model_class(_lowerCamelCase						).eval()
         __SCREAMING_SNAKE_CASE :			Union[str, Any]												=       False
         __SCREAMING_SNAKE_CASE :			int												=       model_class(_lowerCamelCase						,							dtype=jnp.floataa						)

         __SCREAMING_SNAKE_CASE :			Dict												=       load_flax_weights_in_pytorch_model(_lowerCamelCase						,							fx_model.params						)
         __SCREAMING_SNAKE_CASE							,      __SCREAMING_SNAKE_CASE :			Optional[int]												=       pt_inputs['''input_ids'''].shape
         __SCREAMING_SNAKE_CASE :			Dict												=       np.random.randint(0						,							seq_length - 1						,							size=(batch_size,)						)
         for batch_idx, start_index in enumerate(_lowerCamelCase						):
          __SCREAMING_SNAKE_CASE :			List[Any]												=       0
          __SCREAMING_SNAKE_CASE :			str												=       1
          __SCREAMING_SNAKE_CASE :			List[Any]												=       0
          __SCREAMING_SNAKE_CASE :			Dict												=       1

         # make sure weights are tied in PyTorch
         pt_model.tie_weights()

         with torch.no_grad():
          __SCREAMING_SNAKE_CASE :			Any												=       pt_model(**_lowerCamelCase						).to_tuple()

         __SCREAMING_SNAKE_CASE :			Any												=       fx_model(**_lowerCamelCase						).to_tuple()
         self.assertEqual(len(_lowerCamelCase						)						,							len(_lowerCamelCase						)						,							'''Output lengths differ between Flax and PyTorch'''						)
         for fx_output, pt_output in zip(_lowerCamelCase						,							_lowerCamelCase						):
          self.assert_almost_equals(fx_output[:, -1]						,							pt_output[:, -1].numpy()						,							4e-2						)

         with tempfile.TemporaryDirectory() as tmpdirname:
          fx_model.save_pretrained(_lowerCamelCase						)
          __SCREAMING_SNAKE_CASE :			List[str]												=       pt_model_class.from_pretrained(_lowerCamelCase						,							from_flax=_lowerCamelCase						)

         with torch.no_grad():
          __SCREAMING_SNAKE_CASE :			List[str]												=       pt_model_loaded(**_lowerCamelCase						).to_tuple()

         self.assertEqual(
             len(_lowerCamelCase						)						,							len(_lowerCamelCase						)						,							'''Output lengths differ between Flax and PyTorch'''						)
         for fx_output, pt_output in zip(_lowerCamelCase						,							_lowerCamelCase						):
          self.assert_almost_equals(fx_output[:, -1]						,							pt_output[:, -1].numpy()						,							4e-2						)



      @slow
      def 						SCREAMING_SNAKE_CASE_			(			self   :int						):
       for model_class_name in self.all_model_classes:
        __SCREAMING_SNAKE_CASE :			int												=       model_class_name.from_pretrained('''facebook/xglm-564M'''						)
        __SCREAMING_SNAKE_CASE :			Dict												=       model(np.ones((1, 1)						)						)
        self.assertIsNotNone(_lowerCamelCase						)








from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCAmelCase_        =			logging.get_logger(__name__)

lowerCAmelCase_        =			{
    """google/mobilebert-uncased""": """https://huggingface.co/google/mobilebert-uncased/resolve/main/config.json"""
}



class 				_lowerCAmelCase			(	UpperCAmelCase_							):


  '''simple docstring'''
  a_							:				List[Any]					    =MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP
  a_							:				Union[str, Any]					    ="""mobilebert"""







  def __init__(					self					:    int					, UpperCamelCase					:    int=3_05_22					, UpperCamelCase					:    Any=5_12					, UpperCamelCase					:    Any=24					, UpperCamelCase					:    List[str]=4					, UpperCamelCase					:    Any=5_12					, UpperCamelCase					:    str="relu"					, UpperCamelCase					:    Dict=0.0					, UpperCamelCase					:    int=0.1					, UpperCamelCase					:    Optional[Any]=5_12					, UpperCamelCase					:    str=2					, UpperCamelCase					:    Tuple=0.02					, UpperCamelCase					:    Union[str, Any]=1e-1_2					, UpperCamelCase					:    Optional[Any]=0					, UpperCamelCase					:    Union[str, Any]=1_28					, UpperCamelCase					:    Tuple=True					, UpperCamelCase					:    int=True					, UpperCamelCase					:    List[Any]=1_28					, UpperCamelCase					:    Dict=False					, UpperCamelCase					:    List[str]=True					, UpperCamelCase					:    Optional[Any]=4					, UpperCamelCase					:    int="no_norm"					, UpperCamelCase					:    Tuple=True					, UpperCamelCase					:    str=None					, **UpperCamelCase					:    str					, ):


    '''simple docstring'''





    super().__init__(pad_token_id=UpperCamelCase					, **UpperCamelCase							)

    _snake_case      : Union[str, Any]														=		vocab_size
    _snake_case      : Any														=		hidden_size
    _snake_case      : Tuple														=		num_hidden_layers
    _snake_case      : Any														=		num_attention_heads
    _snake_case      : Optional[Any]														=		hidden_act
    _snake_case      : List[Any]														=		intermediate_size
    _snake_case      : List[Any]														=		hidden_dropout_prob
    _snake_case      : Optional[int]														=		attention_probs_dropout_prob
    _snake_case      : int														=		max_position_embeddings
    _snake_case      : Optional[int]														=		type_vocab_size
    _snake_case      : Any														=		initializer_range
    _snake_case      : Optional[int]														=		layer_norm_eps
    _snake_case      : Tuple														=		embedding_size
    _snake_case      : int														=		trigram_input
    _snake_case      : Optional[int]														=		use_bottleneck
    _snake_case      : List[Any]														=		intra_bottleneck_size
    _snake_case      : Any														=		use_bottleneck_attention
    _snake_case      : str														=		key_query_shared_bottleneck
    _snake_case      : List[str]														=		num_feedforward_networks
    _snake_case      : int														=		normalization_type
    _snake_case      : str														=		classifier_activation

    if self.use_bottleneck:
      _snake_case      : Any														=		intra_bottleneck_size
    else:
      _snake_case      : Union[str, Any]														=		hidden_size

    _snake_case      : Dict														=		classifier_dropout




class 				_lowerCAmelCase			(	UpperCAmelCase_							):


  '''simple docstring'''







  @property
  def UpperCamelCase_       (					self					:    List[Any]							):


    '''simple docstring'''





    if self.task == "multiple-choice":
      _snake_case      : str														=		{0: 'batch', 1: 'choice', 2: 'sequence'}
    else:
      _snake_case      : Dict														=		{0: 'batch', 1: 'sequence'}
    return OrderedDict(
        [
            ('input_ids', dynamic_axis),
            ('attention_mask', dynamic_axis),
            ('token_type_ids', dynamic_axis),
        ]							)




