'''simple docstring'''





from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TFRoFormerForCausalLM,
        TFRoFormerForMaskedLM,
        TFRoFormerForMultipleChoice,
        TFRoFormerForQuestionAnswering,
        TFRoFormerForSequenceClassification,
        TFRoFormerForTokenClassification,
        TFRoFormerModel,
    )
    from transformers.models.roformer.modeling_tf_roformer import (
        TFRoFormerSelfAttention,
        TFRoFormerSinusoidalPositionalEmbedding,
    )




class 	__A				:



    def __init__(self   :   Union[str, Any]       ,      __a   :   Dict       ,      __a   :   Dict=13       ,      __a   :   List[str]=7       ,      __a   :   Optional[Any]=True       ,      __a   :   Any=True       ,      __a   :   List[Any]=True       ,      __a   :   Any=True       ,      __a   :   Any=99       ,      __a   :   Any=32       ,      __a   :   Union[str, Any]=2       ,      __a   :   Optional[int]=4       ,      __a   :   Optional[int]=37       ,      __a   :   Union[str, Any]="gelu"       ,      __a   :   Dict=0.1       ,      __a   :   str=0.1       ,      __a   :   Any=512       ,      __a   :   Tuple=16       ,      __a   :   Optional[int]=2       ,      __a   :   List[str]=0.02       ,      __a   :   Any=3       ,      __a   :   str=4       ,      __a   :   int=None       ,      ):
         UpperCAmelCase_	   =       parent
         UpperCAmelCase_	   =       13
         UpperCAmelCase_	   =       7
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       99
         UpperCAmelCase_	   =       32
         UpperCAmelCase_	   =       2
         UpperCAmelCase_	   =       4
         UpperCAmelCase_	   =       37
         UpperCAmelCase_	   =       "gelu"
         UpperCAmelCase_	   =       0.1
         UpperCAmelCase_	   =       0.1
         UpperCAmelCase_	   =       512
         UpperCAmelCase_	   =       16
         UpperCAmelCase_	   =       2
         UpperCAmelCase_	   =       0.02
         UpperCAmelCase_	   =       3
         UpperCAmelCase_	   =       4
         UpperCAmelCase_	   =       None



    def  _lowercase						(self   :   Any			):
         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_	   =       RoFormerConfig(
             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       ,      initializer_range=self.initializer_range       ,      return_dict=__a       ,      )

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



    def  _lowercase						(self   :   int       ,      __a   :   Tuple       ,      __a   :   Dict       ,      __a   :   Union[str, Any]       ,      __a   :   List[str]       ,      __a   :   Union[str, Any]       ,      __a   :   List[str]       ,      __a   :   Dict			):
         UpperCAmelCase_	   =       TFRoFormerModel(config=__a			)
         UpperCAmelCase_	   =       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

         UpperCAmelCase_	   =       [input_ids, input_mask]
         UpperCAmelCase_	   =       model(__a			)

         UpperCAmelCase_	   =       model(__a			)

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



    def  _lowercase						(self   :   List[Any]       ,      __a   :   Optional[Any]       ,      __a   :   Optional[Any]       ,      __a   :   Tuple       ,      __a   :   Optional[Any]       ,      __a   :   Tuple       ,      __a   :   Any       ,      __a   :   str			):
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       TFRoFormerForCausalLM(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         UpperCAmelCase_	   =       model(__a			)["logits"]
         self.parent.assertListEqual(
             list(prediction_scores.numpy().shape			)       ,      [self.batch_size, self.seq_length, self.vocab_size]			)



    def  _lowercase						(self   :   Dict       ,      __a   :   int       ,      __a   :   Optional[int]       ,      __a   :   Union[str, Any]       ,      __a   :   Any       ,      __a   :   Optional[int]       ,      __a   :   Union[str, Any]       ,      __a   :   List[Any]			):
         UpperCAmelCase_	   =       TFRoFormerForMaskedLM(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         UpperCAmelCase_	   =       model(__a			)
         self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.seq_length, self.vocab_size)			)



    def  _lowercase						(self   :   int       ,      __a   :   List[Any]       ,      __a   :   Any       ,      __a   :   Dict       ,      __a   :   List[Any]       ,      __a   :   Any       ,      __a   :   List[str]       ,      __a   :   int			):
         UpperCAmelCase_	   =       self.num_labels
         UpperCAmelCase_	   =       TFRoFormerForSequenceClassification(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         UpperCAmelCase_	   =       model(__a			)
         self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.num_labels)			)



    def  _lowercase						(self   :   Optional[Any]       ,      __a   :   str       ,      __a   :   Any       ,      __a   :   Optional[Any]       ,      __a   :   Tuple       ,      __a   :   Tuple       ,      __a   :   Any       ,      __a   :   Any			):
         UpperCAmelCase_	   =       self.num_choices
         UpperCAmelCase_	   =       TFRoFormerForMultipleChoice(config=__a			)
         UpperCAmelCase_	   =       tf.tile(tf.expand_dims(__a       ,      1			)       ,      (1, self.num_choices, 1)			)
         UpperCAmelCase_	   =       tf.tile(tf.expand_dims(__a       ,      1			)       ,      (1, self.num_choices, 1)			)
         UpperCAmelCase_	   =       tf.tile(tf.expand_dims(__a       ,      1			)       ,      (1, self.num_choices, 1)			)
         UpperCAmelCase_	   =       {
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
             "token_type_ids": multiple_choice_token_type_ids,
         }
         UpperCAmelCase_	   =       model(__a			)
         self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.num_choices)			)



    def  _lowercase						(self   :   Union[str, Any]       ,      __a   :   Any       ,      __a   :   int       ,      __a   :   Any       ,      __a   :   Optional[int]       ,      __a   :   Tuple       ,      __a   :   Any       ,      __a   :   Any			):
         UpperCAmelCase_	   =       self.num_labels
         UpperCAmelCase_	   =       TFRoFormerForTokenClassification(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         UpperCAmelCase_	   =       model(__a			)
         self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.seq_length, self.num_labels)			)



    def  _lowercase						(self   :   Optional[Any]       ,      __a   :   Optional[int]       ,      __a   :   Any       ,      __a   :   Dict       ,      __a   :   Optional[int]       ,      __a   :   Optional[Any]       ,      __a   :   Any       ,      __a   :   Dict			):
         UpperCAmelCase_	   =       TFRoFormerForQuestionAnswering(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         UpperCAmelCase_	   =       model(__a			)
         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  _lowercase						(self   :   Union[str, Any]			):
         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_tf
class 	__A				(			UpperCamelCase__		,     UpperCamelCase__		,     unittest.TestCase     ):
    a__	:						Union[str, Any]	       =  (
        (
            TFRoFormerModel,
            TFRoFormerForCausalLM,
            TFRoFormerForMaskedLM,
            TFRoFormerForQuestionAnswering,
            TFRoFormerForSequenceClassification,
            TFRoFormerForTokenClassification,
            TFRoFormerForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    a__	:						List[Any]	       =  (
        {
            """feature-extraction""": TFRoFormerModel,
            """fill-mask""": TFRoFormerForMaskedLM,
            """question-answering""": TFRoFormerForQuestionAnswering,
            """text-classification""": TFRoFormerForSequenceClassification,
            """text-generation""": TFRoFormerForCausalLM,
            """token-classification""": TFRoFormerForTokenClassification,
            """zero-shot""": TFRoFormerForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )

    a__	:						List[str]	       =  False
    a__	:						Optional[Any]	       =  False



    def  _lowercase						(self   :   int       ,      __a   :   Any       ,      __a   :   str       ,      __a   :   Union[str, Any]       ,      __a   :   Optional[Any]       ,      __a   :   Optional[int]			):
         if pipeline_test_casse_name == "TextGenerationPipelineTests":
              return True

         return False



    def  _lowercase						(self   :   List[str]			):
         UpperCAmelCase_	   =       TFRoFormerModelTester(self			)
         UpperCAmelCase_	   =       ConfigTester(self       ,      config_class=__a       ,      hidden_size=37			)



    def  _lowercase						(self   :   Any			):
         self.config_tester.run_common_tests()



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__a			)



    def  _lowercase						(self   :   Dict			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__a			)



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_lm_head(*__a			)



    def  _lowercase						(self   :   int			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*__a			)



    def  _lowercase						(self   :   str			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__a			)



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__a			)



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__a			)






    @slow
    def  _lowercase						(self   :   int			):
         UpperCAmelCase_	   =       TFRoFormerModel.from_pretrained("junnyu/roformer_chinese_base"			)
         self.assertIsNotNone(__a			)




@require_tf
class 	__A				(			unittest.TestCase     ):



    @slow
    def  _lowercase						(self   :   Optional[Any]			):
         UpperCAmelCase_	   =       TFRoFormerForMaskedLM.from_pretrained("junnyu/roformer_chinese_base"			)
         UpperCAmelCase_	   =       tf.constant([[0, 1, 2, 3, 4, 5]]			)
         UpperCAmelCase_	   =       model(__a			)[0]

         # TODO Replace vocab size
         UpperCAmelCase_	   =       50000

         UpperCAmelCase_	   =       [1, 6, vocab_size]
         self.assertEqual(output.shape       ,      __a			)

         print(output[:, :3, :3]			)

         # TODO Replace values below with what was printed above.
         UpperCAmelCase_	   =       tf.constant(
             [
                 [
                     [-0.12_05_33_41, -1.0_26_49_01, 0.29_22_19_46],
                     [-1.5_13_37_83, 0.19_74_33, 0.15_19_06_07],
                     [-5.0_13_54_03, -3.90_02_56, -0.84_03_87_64],
                 ]
             ]			)
         tf.debugging.assert_near(output[:, :3, :3]       ,      __a       ,      atol=1E-4			)




@require_tf
class 	__A				(			unittest.TestCase     ):
    a__	:						Union[str, Any]	       =  1e-4



    def  _lowercase						(self   :   str			):
         UpperCAmelCase_	   =       tf.constant([[4, 10]]			)
         UpperCAmelCase_	   =       TFRoFormerSinusoidalPositionalEmbedding(num_positions=6       ,      embedding_dim=6			)

         UpperCAmelCase_	   =       emba(input_ids.shape			)
         UpperCAmelCase_	   =       tf.constant(
             [[0.00_00, 0.00_00, 0.00_00, 1.00_00, 1.00_00, 1.00_00], [0.84_15, 0.04_64, 0.00_22, 0.54_03, 0.99_89, 1.00_00]]			)

         tf.debugging.assert_near(__a       ,      __a       ,      atol=self.tolerance			)






    def  _lowercase						(self   :   Optional[Any]			):
         UpperCAmelCase_	   =       tf.constant(
             [
                 [0.00_00, 0.00_00, 0.00_00, 0.00_00, 0.00_00],
                 [0.84_15, 0.82_19, 0.80_20, 0.78_19, 0.76_17],
                 [0.90_93, 0.93_64, 0.95_81, 0.97_49, 0.98_70],
             ]			)
         UpperCAmelCase_	   =       TFRoFormerSinusoidalPositionalEmbedding(num_positions=512       ,      embedding_dim=512			)
         emba([2, 16, 512]			)
         UpperCAmelCase_	   =       emba.weight[:3, :5]

         tf.debugging.assert_near(__a       ,      __a       ,      atol=self.tolerance			)




@require_tf
class 	__A				(			unittest.TestCase     ):
    a__	:						Optional[int]	       =  1e-4



    def  _lowercase						(self   :   Optional[int]			):
         # 2,12,16,64
         UpperCAmelCase_	   =       tf.reshape(tf.range(2 * 12 * 16 * 64       ,      dtype=tf.floataa			)       ,      shape=(2, 12, 16, 64)			) / 100

         UpperCAmelCase_	   =       -tf.reshape(tf.range(2 * 12 * 16 * 64       ,      dtype=tf.floataa			)       ,      shape=(2, 12, 16, 64)			) / 100

         UpperCAmelCase_	   =       TFRoFormerSinusoidalPositionalEmbedding(num_positions=32       ,      embedding_dim=64			)
         UpperCAmelCase_	   =       embed_positions([2, 16, 768]			)[None, None, :, :]

         UpperCAmelCase_   ,			UpperCAmelCase_	   =       TFRoFormerSelfAttention.apply_rotary_position_embeddings(
             __a       ,      __a       ,      __a			)

         UpperCAmelCase_	   =       tf.constant(
             [
                 [0.00_00, 0.01_00, 0.02_00, 0.03_00, 0.04_00, 0.05_00, 0.06_00, 0.07_00],
                 [-0.20_12, 0.88_97, 0.02_63, 0.94_01, 0.20_74, 0.94_63, 0.34_81, 0.93_43],
                 [-1.70_57, 0.62_71, -1.21_45, 1.38_97, -0.63_03, 1.76_47, -0.11_73, 1.89_85],
                 [-2.17_31, -1.63_97, -2.73_58, 0.28_54, -2.18_40, 1.71_83, -1.30_18, 2.48_71],
                 [0.27_17, -3.61_73, -2.92_06, -2.19_88, -3.66_38, 0.38_58, -2.91_55, 2.29_80],
                 [3.98_59, -2.15_80, -0.79_84, -4.49_04, -4.11_81, -2.02_52, -4.47_82, 1.12_53],
             ]			)
         UpperCAmelCase_	   =       tf.constant(
             [
                 [0.00_00, -0.01_00, -0.02_00, -0.03_00, -0.04_00, -0.05_00, -0.06_00, -0.07_00],
                 [0.20_12, -0.88_97, -0.02_63, -0.94_01, -0.20_74, -0.94_63, -0.34_81, -0.93_43],
                 [1.70_57, -0.62_71, 1.21_45, -1.38_97, 0.63_03, -1.76_47, 0.11_73, -1.89_85],
                 [2.17_31, 1.63_97, 2.73_58, -0.28_54, 2.18_40, -1.71_83, 1.30_18, -2.48_71],
                 [-0.27_17, 3.61_73, 2.92_06, 2.19_88, 3.66_38, -0.38_58, 2.91_55, -2.29_80],
                 [-3.98_59, 2.15_80, 0.79_84, 4.49_04, 4.11_81, 2.02_52, 4.47_82, -1.12_53],
             ]			)

         tf.debugging.assert_near(query_layer[0, 0, :6, :8]       ,      __a       ,      atol=self.tolerance			)
         tf.debugging.assert_near(key_layer[0, 0, :6, :8]       ,      __a       ,      atol=self.tolerance			)





'''simple docstring'''





# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
    import jax
    import jaxlib

SCREAMING_SNAKE_CASE_:					List[str]  =get_logger()

SCREAMING_SNAKE_CASE_:					Optional[dict]  =None




class 	__A				(			TensorFormatter[Mapping, """jax.Array""", Mapping]     ):



    def __init__(self   :   List[Any]       ,      __a   :   Optional[int]=None       ,      __a   :   Any=None       ,      **__a   :   Dict			):
         super().__init__(features=__a			)
         import jax
         from jaxlib.xla_client import Device

         if isinstance(__a       ,      __a			):
              raise ValueError(
                  f"""Expected {device} to be a `str` not {type(__a			)}, as `jaxlib.xla_extension.Device` """
                  "is not serializable neither with `pickle` nor with `dill`. Instead you can surround "
                  "the device with `str()` to get its string identifier that will be internally mapped "
                  "to the actual `jaxlib.xla_extension.Device`."			)
         UpperCAmelCase_	   =       device if isinstance(__a       ,      __a			) else str(jax.devices()[0]			)
         # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
         # with `pickle` nor with `dill`, so we need to use a global variable instead
         global DEVICE_MAPPING
         if DEVICE_MAPPING is None:
              UpperCAmelCase_	   =       self._map_devices_to_str()
         if self.device not in list(DEVICE_MAPPING.keys()			):
              logger.warning(
                  f"""Device with string identifier {self.device} not listed among the available """
                  f"""devices: {list(DEVICE_MAPPING.keys()			)}, so falling back to the default """
                  f"""device: {str(jax.devices()[0]			)}."""			)
              UpperCAmelCase_	   =       str(jax.devices()[0]			)
         UpperCAmelCase_	   =       jnp_array_kwargs



    @staticmethod
    def  _lowercase						():
         import jax

         return {str(__a			): device for device in jax.devices()}



    def  _lowercase						(self   :   str       ,      __a   :   Tuple			):
         import jax
         import jax.numpy as jnp

         if isinstance(__a       ,      __a			) and column:
              if all(
                  isinstance(__a       ,      jax.Array			) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column			):
                   return jnp.stack(__a       ,      axis=0			)
         return column



    def  _lowercase						(self   :   Any       ,      __a   :   Optional[int]			):
         import jax
         import jax.numpy as jnp

         if isinstance(__a       ,      (str, bytes, type(__a			))			):
              return value
         elif isinstance(__a       ,      (np.character, np.ndarray)			) and np.issubdtype(value.dtype       ,      np.character			):
              return value.tolist()

         UpperCAmelCase_	   =       {}

         if isinstance(__a       ,      (np.number, np.ndarray)			) and np.issubdtype(value.dtype       ,      np.integer			):
              # the default int precision depends on the jax config
              # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
              if jax.config.jax_enable_xaa:
                   UpperCAmelCase_	   =       {"dtype": jnp.intaa}
              else:
                   UpperCAmelCase_	   =       {"dtype": jnp.intaa}
         elif isinstance(__a       ,      (np.number, np.ndarray)			) and np.issubdtype(value.dtype       ,      np.floating			):
              UpperCAmelCase_	   =       {"dtype": jnp.floataa}
         elif config.PIL_AVAILABLE and "PIL" in sys.modules:
              import PIL.Image

              if isinstance(__a       ,      PIL.Image.Image			):
                   UpperCAmelCase_	   =       np.asarray(__a			)

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
         global DEVICE_MAPPING
         if DEVICE_MAPPING is None:
              UpperCAmelCase_	   =       self._map_devices_to_str()

         with jax.default_device(DEVICE_MAPPING[self.device]			):
              # calling jnp.array on a np.ndarray does copy the data
              # see https://github.com/google/jax/issues/4486
              return jnp.array(__a       ,      **{**default_dtype, **self.jnp_array_kwargs}			)



    def  _lowercase						(self   :   int       ,      __a   :   Any			):
         import jax

         # support for torch, tf, jax etc.
         if config.TORCH_AVAILABLE and "torch" in sys.modules:
              import torch

              if isinstance(__a       ,      torch.Tensor			):
                   return self._tensorize(data_struct.detach().cpu().numpy()[()]			)
         if hasattr(__a       ,      "__array__"			) and not isinstance(__a       ,      jax.Array			):
              UpperCAmelCase_	   =       data_struct.__array__()
         # support for nested types like struct of list of struct
         if isinstance(__a       ,      np.ndarray			):
              if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
                   return self._consolidate([self.recursive_tensorize(__a			) for substruct in data_struct]			)
         elif isinstance(__a       ,      (list, tuple)			):
              return self._consolidate([self.recursive_tensorize(__a			) for substruct in data_struct]			)
         return self._tensorize(__a			)



    def  _lowercase						(self   :   Union[str, Any]       ,      __a   :   dict			):
         return map_nested(self._recursive_tensorize       ,      __a       ,      map_list=__a			)



    def  _lowercase						(self   :   str       ,      __a   :   pa.Table			):
         UpperCAmelCase_	   =       self.numpy_arrow_extractor().extract_row(__a			)
         UpperCAmelCase_	   =       self.python_features_decoder.decode_row(__a			)
         return self.recursive_tensorize(__a			)



    def  _lowercase						(self   :   Tuple       ,      __a   :   pa.Table			):
         UpperCAmelCase_	   =       self.numpy_arrow_extractor().extract_column(__a			)
         UpperCAmelCase_	   =       self.python_features_decoder.decode_column(__a       ,      pa_table.column_names[0]			)
         UpperCAmelCase_	   =       self.recursive_tensorize(__a			)
         UpperCAmelCase_	   =       self._consolidate(__a			)
         return column






    def  _lowercase						(self   :   str       ,      __a   :   pa.Table			):
         UpperCAmelCase_	   =       self.numpy_arrow_extractor().extract_batch(__a			)
         UpperCAmelCase_	   =       self.python_features_decoder.decode_batch(__a			)
         UpperCAmelCase_	   =       self.recursive_tensorize(__a			)
         for column_name in batch:
              UpperCAmelCase_	   =       self._consolidate(batch[column_name]			)
         return batch




