


import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       MODEL_FOR_PRETRAINING_MAPPING,
       MegatronBertForCausalLM,
       MegatronBertForMaskedLM,
       MegatronBertForMultipleChoice,
       MegatronBertForNextSentencePrediction,
       MegatronBertForPreTraining,
       MegatronBertForQuestionAnswering,
       MegatronBertForSequenceClassification,
       MegatronBertForTokenClassification,
       MegatronBertModel,
   )






class    UpperCAmelCase__	:



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_=13 ,       A_=7 ,       A_=True ,       A_=True ,       A_=True ,       A_=True ,       A_=99 ,       A_=64 ,       A_=32 ,       A_=5 ,       A_=4 ,       A_=37 ,       A_="gelu" ,       A_=0.1 ,       A_=0.1 ,       A_=512 ,       A_=16 ,       A_=2 ,       A_=0.02 ,       A_=3 ,       A_=4 ,       A_=None ,       )  ->	Optional[Any]:
         __UpperCamelCase   							=parent
         __UpperCamelCase   							=batch_size
         __UpperCamelCase   							=seq_length
         __UpperCamelCase   							=is_training
         __UpperCamelCase   							=use_input_mask
         __UpperCamelCase   							=use_token_type_ids
         __UpperCamelCase   							=use_labels
         __UpperCamelCase   							=vocab_size
         __UpperCamelCase   							=hidden_size
         __UpperCamelCase   							=embedding_size
         __UpperCamelCase   							=num_hidden_layers
         __UpperCamelCase   							=num_attention_heads
         __UpperCamelCase   							=intermediate_size
         __UpperCamelCase   							=hidden_act
         __UpperCamelCase   							=hidden_dropout_prob
         __UpperCamelCase   							=attention_probs_dropout_prob
         __UpperCamelCase   							=max_position_embeddings
         __UpperCamelCase   							=type_vocab_size
         __UpperCamelCase   							=type_sequence_label_size
         __UpperCamelCase   							=initializer_range
         __UpperCamelCase   							=num_labels
         __UpperCamelCase   							=num_choices
         __UpperCamelCase   							=scope






    def _a (      self   )  ->	Optional[int]:
         __UpperCamelCase   							=ids_tensor([self.batch_size, self.seq_length] ,       self.vocab_size   )

         __UpperCamelCase   							=None
         if self.use_input_mask:
              __UpperCamelCase   							=random_attention_mask([self.batch_size, self.seq_length]   )

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

         __UpperCamelCase   							=None
         __UpperCamelCase   							=None
         __UpperCamelCase   							=None
         if self.use_labels:
              __UpperCamelCase   							=ids_tensor([self.batch_size] ,       self.type_sequence_label_size   )
              __UpperCamelCase   							=ids_tensor([self.batch_size, self.seq_length] ,       self.num_labels   )
              __UpperCamelCase   							=ids_tensor([self.batch_size] ,       self.num_choices   )

         __UpperCamelCase   							=self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






    def _a (      self   )  ->	Tuple:
         return MegatronBertConfig(
             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 ,       embedding_size=self.embedding_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=UpperCamelCase_ ,       initializer_range=self.initializer_range ,       )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	List[Any]:
         __UpperCamelCase   							=MegatronBertModel(config=UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.eval()
         __UpperCamelCase   							=model(UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       token_type_ids=UpperCamelCase_   )
         __UpperCamelCase   							=model(UpperCamelCase_ ,       token_type_ids=UpperCamelCase_   )
         __UpperCamelCase   							=model(UpperCamelCase_   )

         self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, self.seq_length, self.hidden_size)   )
         self.parent.assertEqual(result.pooler_output.shape ,       (self.batch_size, self.hidden_size)   )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	Union[str, Any]:
         __UpperCamelCase   							=MegatronBertForMaskedLM(config=UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.eval()
         __UpperCamelCase   							=model(UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       token_type_ids=UpperCamelCase_ ,       labels=UpperCamelCase_   )
         self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.seq_length, self.vocab_size)   )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	Union[str, Any]:
         __UpperCamelCase   							=MegatronBertForCausalLM(config=UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.eval()
         __UpperCamelCase   							=model(UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       token_type_ids=UpperCamelCase_ ,       labels=UpperCamelCase_   )
         self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.seq_length, self.vocab_size)   )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	List[Any]:
         __UpperCamelCase   							=MegatronBertForNextSentencePrediction(config=UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.eval()
         __UpperCamelCase   							=model(
             UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       token_type_ids=UpperCamelCase_ ,       labels=UpperCamelCase_ ,       )
         self.parent.assertEqual(result.logits.shape ,       (self.batch_size, 2)   )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	Union[str, Any]:
         __UpperCamelCase   							=MegatronBertForPreTraining(config=UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.eval()
         __UpperCamelCase   							=model(
             UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       token_type_ids=UpperCamelCase_ ,       labels=UpperCamelCase_ ,       next_sentence_label=UpperCamelCase_ ,       )
         self.parent.assertEqual(result.prediction_logits.shape ,       (self.batch_size, self.seq_length, self.vocab_size)   )
         self.parent.assertEqual(result.seq_relationship_logits.shape ,       (self.batch_size, 2)   )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	List[Any]:
         __UpperCamelCase   							=MegatronBertForQuestionAnswering(config=UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.eval()
         __UpperCamelCase   							=model(
             UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       token_type_ids=UpperCamelCase_ ,       start_positions=UpperCamelCase_ ,       end_positions=UpperCamelCase_ ,       )
         self.parent.assertEqual(result.start_logits.shape ,       (self.batch_size, self.seq_length)   )
         self.parent.assertEqual(result.end_logits.shape ,       (self.batch_size, self.seq_length)   )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	Tuple:
         __UpperCamelCase   							=self.num_labels
         __UpperCamelCase   							=MegatronBertForSequenceClassification(UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.eval()
         __UpperCamelCase   							=model(UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       token_type_ids=UpperCamelCase_ ,       labels=UpperCamelCase_   )
         self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_labels)   )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	Union[str, Any]:
         __UpperCamelCase   							=self.num_labels
         __UpperCamelCase   							=MegatronBertForTokenClassification(config=UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.eval()
         __UpperCamelCase   							=model(UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       token_type_ids=UpperCamelCase_ ,       labels=UpperCamelCase_   )
         self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.seq_length, self.num_labels)   )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_ ,       A_   )  ->	List[str]:
         __UpperCamelCase   							=self.num_choices
         __UpperCamelCase   							=MegatronBertForMultipleChoice(config=UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.eval()
         __UpperCamelCase   							=input_ids.unsqueeze(1   ).expand(-1 ,       self.num_choices ,       -1   ).contiguous()
         __UpperCamelCase   							=token_type_ids.unsqueeze(1   ).expand(-1 ,       self.num_choices ,       -1   ).contiguous()
         __UpperCamelCase   							=input_mask.unsqueeze(1   ).expand(-1 ,       self.num_choices ,       -1   ).contiguous()
         __UpperCamelCase   							=model(
             UpperCamelCase_ ,       attention_mask=UpperCamelCase_ ,       token_type_ids=UpperCamelCase_ ,       labels=UpperCamelCase_ ,       )
         self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_choices)   )






    def _a (      self   )  ->	int:
         __UpperCamelCase   							=self.prepare_config_and_inputs()
         (
             (
             __UpperCamelCase       
         )							,				(
             __UpperCamelCase       
         )							,				(
             __UpperCamelCase       
         )							,				(
             __UpperCamelCase       
         )							,				(
             __UpperCamelCase       
         )							,				(
             __UpperCamelCase       
         )							,				(
             __UpperCamelCase       
         )							,				       
         )   							=config_and_inputs
         __UpperCamelCase   							={'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
         return config, inputs_dict






@require_torch
class    UpperCAmelCase__	( __lowerCamelCase       ,  __lowerCamelCase       ,  unittest.TestCase		):



    """simple docstring"""






    UpperCAmelCase__						:		Tuple              =  (
        (
            MegatronBertModel,
            MegatronBertForMaskedLM,
            MegatronBertForCausalLM,
            MegatronBertForMultipleChoice,
            MegatronBertForNextSentencePrediction,
            MegatronBertForPreTraining,
            MegatronBertForQuestionAnswering,
            MegatronBertForSequenceClassification,
            MegatronBertForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    UpperCAmelCase__						:		Optional[int]              =  (
        {
            'feature-extraction': MegatronBertModel,
            'fill-mask': MegatronBertForMaskedLM,
            'question-answering': MegatronBertForQuestionAnswering,
            'text-classification': MegatronBertForSequenceClassification,
            'text-generation': MegatronBertForCausalLM,
            'token-classification': MegatronBertForTokenClassification,
            'zero-shot': MegatronBertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    UpperCAmelCase__						:		Optional[int]              =  True
    # test_resize_embeddings = False
    UpperCAmelCase__						:		List[str]              =  False






    def _a (      self ,       A_ ,       A_ ,       A_=False   )  ->	Union[str, Any]:
         __UpperCamelCase   							=super()._prepare_for_class(UpperCamelCase_ ,       UpperCamelCase_ ,       return_labels=UpperCamelCase_   )

         if return_labels:
              if model_class in get_values(UpperCamelCase_   ):
                   __UpperCamelCase   							=torch.zeros(
                       (self.model_tester.batch_size, self.model_tester.seq_length) ,       dtype=torch.long ,       device=UpperCamelCase_   )
                   __UpperCamelCase   							=torch.zeros(
                       self.model_tester.batch_size ,       dtype=torch.long ,       device=UpperCamelCase_   )
         return inputs_dict






    def _a (      self   )  ->	Dict:
         __UpperCamelCase   							=MegatronBertModelTester(self   )
         __UpperCamelCase   							=ConfigTester(self ,       config_class=UpperCamelCase_ ,       hidden_size=37   )






    def _a (      self   )  ->	Optional[Any]:
         self.config_tester.run_common_tests()






    def _a (      self   )  ->	Optional[Any]:
         __UpperCamelCase   							=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase_   )






    def _a (      self   )  ->	List[Any]:
         __UpperCamelCase   							=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase_   )






    def _a (      self   )  ->	str:
         __UpperCamelCase   							=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCamelCase_   )






    def _a (      self   )  ->	Optional[Any]:
         __UpperCamelCase   							=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase_   )






    def _a (      self   )  ->	Union[str, Any]:
         __UpperCamelCase   							=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase_   )






    def _a (      self   )  ->	List[Any]:
         __UpperCamelCase   							=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase_   )






    def _a (      self   )  ->	Optional[Any]:
         __UpperCamelCase   							=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase_   )






    def _a (      self   )  ->	Tuple:
         __UpperCamelCase   							=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCamelCase_   )

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Optional[int]    ):
     return torch.tensor(
         A__      ,					dtype=torch.long      ,					device=A__      ,					)


_A			=						1e-4






@require_torch
@require_sentencepiece
@require_tokenizers
class    UpperCAmelCase__	( unittest.TestCase		):



    """simple docstring"""






    @slow
    @unittest.skip('Model is not available.'   )
    def _a (      self   )  ->	str:
         __UpperCamelCase   							='nvidia/megatron-bert-uncased-345m'
         if "MYDIR" in os.environ:
              __UpperCamelCase   							=os.path.join(os.environ['MYDIR'] ,       UpperCamelCase_   )
         __UpperCamelCase   							=MegatronBertModel.from_pretrained(UpperCamelCase_   )
         model.to(UpperCamelCase_   )
         model.half()
         __UpperCamelCase   							=_long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]   )
         with torch.no_grad():
              __UpperCamelCase   							=model(UpperCamelCase_   )[0]
         __UpperCamelCase   							=torch.Size((1, 9, 1024)   )
         self.assertEqual(output.shape ,       UpperCamelCase_   )

         __UpperCamelCase   							=[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
         for ii in range(3   ):
              for jj in range(3   ):
                   __UpperCamelCase   							=output[0, ii, jj]
                   __UpperCamelCase   							=expected[3 * ii + jj]
                   __UpperCamelCase   							='ii={} jj={} a={} b={}'.format(UpperCamelCase_ ,       UpperCamelCase_ ,       UpperCamelCase_ ,       UpperCamelCase_   )
                   self.assertTrue(math.isclose(UpperCamelCase_ ,       UpperCamelCase_ ,       rel_tol=UpperCamelCase_ ,       abs_tol=UpperCamelCase_   ) ,       msg=UpperCamelCase_   )









from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


_A			=						logging.get_logger(__name__)


if is_vision_available():
   import PIL






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    UpperCAmelCase__						:		Optional[int]              =  ["pixel_values"]






    def __init__(      self ,       A_ = True ,       A_ = None ,       A_ = PILImageResampling.BICUBIC ,       A_ = True ,       A_ = None ,       A_ = True ,       A_ = 1 / 255 ,       A_ = True ,       A_ = None ,       A_ = None ,       A_ = True ,       **A_ ,       )  ->	None:
         super().__init__(**A_   )
         __UpperCamelCase   							=size if size is not None else {'shortest_edge': 224}
         __UpperCamelCase   							=get_size_dict(A_ ,       default_to_square=A_   )
         __UpperCamelCase   							=crop_size if crop_size is not None else {'height': 224, 'width': 224}
         __UpperCamelCase   							=get_size_dict(A_ ,       default_to_square=A_ ,       param_name='crop_size'   )

         __UpperCamelCase   							=do_resize
         __UpperCamelCase   							=size
         __UpperCamelCase   							=resample
         __UpperCamelCase   							=do_center_crop
         __UpperCamelCase   							=crop_size
         __UpperCamelCase   							=do_rescale
         __UpperCamelCase   							=rescale_factor
         __UpperCamelCase   							=do_normalize
         __UpperCamelCase   							=image_mean if image_mean is not None else OPENAI_CLIP_MEAN
         __UpperCamelCase   							=image_std if image_std is not None else OPENAI_CLIP_STD
         __UpperCamelCase   							=do_convert_rgb






    def _a (      self ,       A_ ,       A_ ,       A_ = PILImageResampling.BICUBIC ,       A_ = None ,       **A_ ,       )  ->	np.ndarray:
         __UpperCamelCase   							=get_size_dict(A_ ,       default_to_square=A_   )
         if "shortest_edge" not in size:
              raise ValueError(f'The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'   )
         __UpperCamelCase   							=get_resize_output_image_size(A_ ,       size=size['shortest_edge'] ,       default_to_square=A_   )
         return resize(A_ ,       size=A_ ,       resample=A_ ,       data_format=A_ ,       **A_   )






    def _a (      self ,       A_ ,       A_ ,       A_ = None ,       **A_ ,       )  ->	np.ndarray:
         __UpperCamelCase   							=get_size_dict(A_   )
         if "height" not in size or "width" not in size:
              raise ValueError(f'The `size` parameter must contain the keys (height, width). Got {size.keys()}'   )
         return center_crop(A_ ,       size=(size['height'], size['width']) ,       data_format=A_ ,       **A_   )






    def _a (      self ,       A_ ,       A_ ,       A_ = None ,       **A_ ,       )  ->	Union[str, Any]:
         return rescale(A_ ,       scale=A_ ,       data_format=A_ ,       **A_   )






    def _a (      self ,       A_ ,       A_ ,       A_ ,       A_ = None ,       **A_ ,       )  ->	np.ndarray:
         return normalize(A_ ,       mean=A_ ,       std=A_ ,       data_format=A_ ,       **A_   )






    def _a (      self ,       A_ ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = None ,       A_ = ChannelDimension.FIRST ,       **A_ ,       )  ->	PIL.Image.Image:
         __UpperCamelCase   							=do_resize if do_resize is not None else self.do_resize
         __UpperCamelCase   							=size if size is not None else self.size
         __UpperCamelCase   							=get_size_dict(A_ ,       param_name='size' ,       default_to_square=A_   )
         __UpperCamelCase   							=resample if resample is not None else self.resample
         __UpperCamelCase   							=do_center_crop if do_center_crop is not None else self.do_center_crop
         __UpperCamelCase   							=crop_size if crop_size is not None else self.crop_size
         __UpperCamelCase   							=get_size_dict(A_ ,       param_name='crop_size' ,       default_to_square=A_   )
         __UpperCamelCase   							=do_rescale if do_rescale is not None else self.do_rescale
         __UpperCamelCase   							=rescale_factor if rescale_factor is not None else self.rescale_factor
         __UpperCamelCase   							=do_normalize if do_normalize is not None else self.do_normalize
         __UpperCamelCase   							=image_mean if image_mean is not None else self.image_mean
         __UpperCamelCase   							=image_std if image_std is not None else self.image_std
         __UpperCamelCase   							=do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

         __UpperCamelCase   							=make_list_of_images(A_   )

         if not valid_images(A_   ):
              raise ValueError(
                  'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '
                  'torch.Tensor, tf.Tensor or jax.ndarray.'   )

         if do_resize and size is None:
              raise ValueError('Size must be specified if do_resize is True.'   )

         if do_center_crop and crop_size is None:
              raise ValueError('Crop size must be specified if do_center_crop is True.'   )

         if do_rescale and rescale_factor is None:
              raise ValueError('Rescale factor must be specified if do_rescale is True.'   )

         if do_normalize and (image_mean is None or image_std is None):
              raise ValueError('Image mean and std must be specified if do_normalize is True.'   )

         # PIL RGBA images are converted to RGB
         if do_convert_rgb:
              __UpperCamelCase   							=[convert_to_rgb(A_   ) for image in images]

         # All transformations expect numpy arrays.
         __UpperCamelCase   							=[to_numpy_array(A_   ) for image in images]

         if do_resize:
              __UpperCamelCase   							=[self.resize(image=A_ ,       size=A_ ,       resample=A_   ) for image in images]

         if do_center_crop:
              __UpperCamelCase   							=[self.center_crop(image=A_ ,       size=A_   ) for image in images]

         if do_rescale:
              __UpperCamelCase   							=[self.rescale(image=A_ ,       scale=A_   ) for image in images]

         if do_normalize:
              __UpperCamelCase   							=[self.normalize(image=A_ ,       mean=A_ ,       std=A_   ) for image in images]

         __UpperCamelCase   							=[to_channel_dimension_format(A_ ,       A_   ) for image in images]

         __UpperCamelCase   							={'pixel_values': images}
         return BatchFeature(data=A_ ,       tensor_type=A_   )





