


import collections
import inspect
import unittest
from typing import Dict, List, Tuple

from transformers import MaskFormerSwinConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, torch_device
from transformers.utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch
   from torch import nn

   from transformers import MaskFormerSwinBackbone
   from transformers.models.maskformer import MaskFormerSwinModel







class lowerCAmelCase_		:

 '''simple docstring'''

 def __init__(			self	,   _UpperCAmelCase	,   _UpperCAmelCase=13	,   _UpperCAmelCase=32	,   _UpperCAmelCase=2	,   _UpperCAmelCase=3	,   _UpperCAmelCase=16	,   _UpperCAmelCase=[1, 2, 1]	,   _UpperCAmelCase=[2, 2, 4]	,   _UpperCAmelCase=2	,   _UpperCAmelCase=2.0	,   _UpperCAmelCase=True	,   _UpperCAmelCase=0.0	,   _UpperCAmelCase=0.0	,   _UpperCAmelCase=0.1	,   _UpperCAmelCase="gelu"	,   _UpperCAmelCase=False	,   _UpperCAmelCase=True	,   _UpperCAmelCase=0.02	,   _UpperCAmelCase=1E-5	,   _UpperCAmelCase=True	,   _UpperCAmelCase=None	,   _UpperCAmelCase=True	,   _UpperCAmelCase=10	,   _UpperCAmelCase=8	,   _UpperCAmelCase=["stage1", "stage2", "stage3"]	,   _UpperCAmelCase=[1, 2, 3]	,   ):
       snake_case_ =   parent
       snake_case_ =   batch_size
       snake_case_ =   image_size
       snake_case_ =   patch_size
       snake_case_ =   num_channels
       snake_case_ =   embed_dim
       snake_case_ =   depths
       snake_case_ =   num_heads
       snake_case_ =   window_size
       snake_case_ =   mlp_ratio
       snake_case_ =   qkv_bias
       snake_case_ =   hidden_dropout_prob
       snake_case_ =   attention_probs_dropout_prob
       snake_case_ =   drop_path_rate
       snake_case_ =   hidden_act
       snake_case_ =   use_absolute_embeddings
       snake_case_ =   patch_norm
       snake_case_ =   layer_norm_eps
       snake_case_ =   initializer_range
       snake_case_ =   is_training
       snake_case_ =   scope
       snake_case_ =   use_labels
       snake_case_ =   type_sequence_label_size
       snake_case_ =   encoder_stride
       snake_case_ =   out_features
       snake_case_ =   out_indices



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

       snake_case_ =   None
       if self.use_labels:
             snake_case_ =   ids_tensor([self.batch_size]	,   self.type_sequence_label_size      )

       snake_case_ =   self.get_config()

       return config, pixel_values, labels



 def 		UpperCamelCase__  (			self      ):
       return MaskFormerSwinConfig(
           image_size=self.image_size	,   patch_size=self.patch_size	,   num_channels=self.num_channels	,   embed_dim=self.embed_dim	,   depths=self.depths	,   num_heads=self.num_heads	,   window_size=self.window_size	,   mlp_ratio=self.mlp_ratio	,   qkv_bias=self.qkv_bias	,   hidden_dropout_prob=self.hidden_dropout_prob	,   attention_probs_dropout_prob=self.attention_probs_dropout_prob	,   drop_path_rate=self.drop_path_rate	,   hidden_act=self.hidden_act	,   use_absolute_embeddings=self.use_absolute_embeddings	,   path_norm=self.patch_norm	,   layer_norm_eps=self.layer_norm_eps	,   initializer_range=self.initializer_range	,   encoder_stride=self.encoder_stride	,   out_features=self.out_features	,   out_indices=self.out_indices	,   )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   MaskFormerSwinModel(config=_UpperCAmelCase      )
       model.to(_UpperCAmelCase      )
       model.eval()
       snake_case_ =   model(_UpperCAmelCase      )

       snake_case_ =   ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths      ) - 1))
       snake_case_ =   int(config.embed_dim * 2 ** (len(config.depths      ) - 1)      )

       self.parent.assertEqual(result.last_hidden_state.shape	,   (self.batch_size, expected_seq_len, expected_dim)      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   MaskFormerSwinBackbone(config=_UpperCAmelCase      )
       model.to(_UpperCAmelCase      )
       model.eval()
       snake_case_ =   model(_UpperCAmelCase      )

       # verify feature maps
       self.parent.assertEqual(len(result.feature_maps      )	,   len(config.out_features      )      )
       self.parent.assertListEqual(list(result.feature_maps[0].shape      )	,   [13, 16, 16, 16]      )

       # verify channels
       self.parent.assertEqual(len(model.channels      )	,   len(config.out_features      )      )
       self.parent.assertListEqual(model.channels	,   [16, 32, 64]      )

       # verify ValueError
       with self.parent.assertRaises(_UpperCAmelCase      ):
             snake_case_ =   ['''stem''']
             snake_case_ =   MaskFormerSwinBackbone(config=_UpperCAmelCase      )





 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.prepare_config_and_inputs()
       snake_case_ ,					snake_case_ ,					snake_case_ =   config_and_inputs
       snake_case_ =   {'''pixel_values''': pixel_values}
       return config, inputs_dict







@require_torch
class lowerCAmelCase_		(				lowerCamelCase__      ,		lowerCamelCase__      ,		unittest.TestCase ):

 '''simple docstring'''

 __snake_case				=						(
     (
         MaskFormerSwinModel,
         MaskFormerSwinBackbone,
     )
     if is_torch_available()
     else ()
 )
 __snake_case				=						{"feature-extraction": MaskFormerSwinModel} if is_torch_available() else {}
 __snake_case				=						False
 __snake_case				=						False
 __snake_case				=						False
 __snake_case				=						False
 __snake_case				=						False



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   MaskFormerSwinModelTester(self      )
       snake_case_ =   ConfigTester(self	,   config_class=_UpperCAmelCase	,   embed_dim=37      )



 @require_torch_multi_gpu
 @unittest.skip(
     reason=(
         '''`MaskFormerSwinModel` outputs `hidden_states_spatial_dimensions` which doesn\'t work well with'''
         ''' `nn.DataParallel`'''
     )      )
 def 		UpperCamelCase__  (			self      ):
       pass



 def 		UpperCamelCase__  (			self      ):
       self.create_and_test_config_common_properties()
       self.config_tester.create_and_test_config_to_json_string()
       self.config_tester.create_and_test_config_to_json_file()
       self.config_tester.create_and_test_config_from_and_save_pretrained()
       self.config_tester.create_and_test_config_with_num_labels()
       self.config_tester.check_config_can_be_init_without_params()
       self.config_tester.check_config_arguments_init()



 def 		UpperCamelCase__  (			self      ):
       return



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_backbone(*_UpperCAmelCase      )



 @unittest.skip('''Swin does not use inputs_embeds'''      )
 def 		UpperCamelCase__  (			self      ):
       pass



 @unittest.skip('''Swin does not support feedforward chunking'''      )
 def 		UpperCamelCase__  (			self      ):
       pass



 def 		UpperCamelCase__  (			self      ):
       snake_case_ ,					snake_case_ =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             snake_case_ =   model_class(_UpperCAmelCase      )
             self.assertIsInstance(model.get_input_embeddings()	,   (nn.Module)      )
             snake_case_ =   model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(_UpperCAmelCase	,   nn.Linear      )      )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ ,					snake_case_ =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             snake_case_ =   model_class(_UpperCAmelCase      )
             snake_case_ =   inspect.signature(model.forward      )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             snake_case_ =   [*signature.parameters.keys()]

             snake_case_ =   ['''pixel_values''']
             self.assertListEqual(arg_names[:1]	,   _UpperCAmelCase      )



 @unittest.skip(reason='''MaskFormerSwin is only used as backbone and doesn\'t support output_attentions'''      )
 def 		UpperCamelCase__  (			self      ):
       pass



 @unittest.skip(reason='''MaskFormerSwin is only used as an internal backbone'''      )
 def 		UpperCamelCase__  (			self      ):
       pass



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   model_class(_UpperCAmelCase      )
       model.to(_UpperCAmelCase      )
       model.eval()

       with torch.no_grad():
             snake_case_ =   model(**self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase      )      )

       snake_case_ =   outputs.hidden_states

       snake_case_ =   getattr(
           self.model_tester	,   '''expected_num_hidden_layers'''	,   len(self.model_tester.depths      ) + 1      )
       self.assertEqual(len(_UpperCAmelCase      )	,   _UpperCAmelCase      )

       # Swin has a different seq_length
       snake_case_ =   (
           config.patch_size
           if isinstance(config.patch_size	,   collections.abc.Iterable      )
           else (config.patch_size, config.patch_size)
       )

       snake_case_ =   (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

       self.assertListEqual(
           list(hidden_states[0].shape[-2:]      )	,   [num_patches, self.model_tester.embed_dim]	,   )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ ,					snake_case_ =   self.model_tester.prepare_config_and_inputs_for_common()

       snake_case_ =   (
           self.model_tester.image_size
           if isinstance(self.model_tester.image_size	,   collections.abc.Iterable      )
           else (self.model_tester.image_size, self.model_tester.image_size)
       )

       for model_class in self.all_model_classes:
             snake_case_ =   True
             self.check_hidden_states_output(_UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      )

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             snake_case_ =   True

             self.check_hidden_states_output(_UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      )



 def 		UpperCamelCase__  (			self      ):
       snake_case_ ,					snake_case_ =   self.model_tester.prepare_config_and_inputs_for_common()
       snake_case_ =   3

       snake_case_ =   (
           self.model_tester.image_size
           if isinstance(self.model_tester.image_size	,   collections.abc.Iterable      )
           else (self.model_tester.image_size, self.model_tester.image_size)
       )
       snake_case_ =   (
           config.patch_size
           if isinstance(config.patch_size	,   collections.abc.Iterable      )
           else (config.patch_size, config.patch_size)
       )

       snake_case_ =   image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
       snake_case_ =   image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

       for model_class in self.all_model_classes:
             snake_case_ =   True
             self.check_hidden_states_output(_UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   (padded_height, padded_width)      )

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             snake_case_ =   True
             self.check_hidden_states_output(_UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   (padded_height, padded_width)      )



 @unittest.skip(reason='''MaskFormerSwin doesn\'t have pretrained checkpoints'''      )
 def 		UpperCamelCase__  (			self      ):
       pass



 @unittest.skip(reason='''This will be fixed once MaskFormerSwin is replaced by native Swin'''      )
 def 		UpperCamelCase__  (			self      ):
       pass



 @unittest.skip(reason='''This will be fixed once MaskFormerSwin is replaced by native Swin'''      )
 def 		UpperCamelCase__  (			self      ):
       pass





 def 		UpperCamelCase__  (			self      ):
       snake_case_ ,					snake_case_ =   self.model_tester.prepare_config_and_inputs_for_common()

       def set_nan_tensor_to_zero(_UpperCAmelCase      ):
             snake_case_ =   0
             return t

       def check_equivalence(_UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase={}      ):
             with torch.no_grad():
                   snake_case_ =   model(**_UpperCAmelCase	,   return_dict=_UpperCAmelCase	,   **_UpperCAmelCase      )
                   snake_case_ =   model(**_UpperCAmelCase	,   return_dict=_UpperCAmelCase	,   **_UpperCAmelCase      ).to_tuple()

                   def recursive_check(_UpperCAmelCase	,   _UpperCAmelCase      ):
                         if isinstance(_UpperCAmelCase	,   (List, Tuple)      ):
                               for tuple_iterable_value, dict_iterable_value in zip(_UpperCAmelCase	,   _UpperCAmelCase      ):
                                     recursive_check(_UpperCAmelCase	,   _UpperCAmelCase      )
                         elif isinstance(_UpperCAmelCase	,   _UpperCAmelCase      ):
                               for tuple_iterable_value, dict_iterable_value in zip(
                                   tuple_object.values()	,   dict_object.values()      ):
                                     recursive_check(_UpperCAmelCase	,   _UpperCAmelCase      )
                         elif tuple_object is None:
                               return
                         else:
                               self.assertTrue(
                                   torch.allclose(
                                       set_nan_tensor_to_zero(_UpperCAmelCase      )	,   set_nan_tensor_to_zero(_UpperCAmelCase      )	,   atol=1E-5      )	,   msg=(
                                       '''Tuple and dict output are not equal. Difference:'''
                                       F''' {torch.max(torch.abs(tuple_object - dict_object      )      )}. Tuple has `nan`:'''
                                       F''' {torch.isnan(_UpperCAmelCase      ).any()} and `inf`: {torch.isinf(_UpperCAmelCase      )}. Dict has'''
                                       F''' `nan`: {torch.isnan(_UpperCAmelCase      ).any()} and `inf`: {torch.isinf(_UpperCAmelCase      )}.'''
                                   )	,   )

                   recursive_check(_UpperCAmelCase	,   _UpperCAmelCase      )

       for model_class in self.all_model_classes:
             snake_case_ =   model_class(_UpperCAmelCase      )
             model.to(_UpperCAmelCase      )
             model.eval()

             snake_case_ =   self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase      )
             snake_case_ =   self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase      )
             check_equivalence(_UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      )

             snake_case_ =   self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase	,   return_labels=_UpperCAmelCase      )
             snake_case_ =   self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase	,   return_labels=_UpperCAmelCase      )
             check_equivalence(_UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      )

             snake_case_ =   self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase      )
             snake_case_ =   self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase      )
             check_equivalence(_UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   {'''output_hidden_states''': True}      )

             snake_case_ =   self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase	,   return_labels=_UpperCAmelCase      )
             snake_case_ =   self._prepare_for_class(_UpperCAmelCase	,   _UpperCAmelCase	,   return_labels=_UpperCAmelCase      )
             check_equivalence(_UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   {'''output_hidden_states''': True}      )







@require_torch
class lowerCAmelCase_		(				unittest.TestCase      ,		lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						(MaskFormerSwinBackbone,) if is_torch_available() else ()
 __snake_case				=						MaskFormerSwinConfig



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   MaskFormerSwinModelTester(self      )





 def 		UpperCamelCase__  (			self      ):
       snake_case_ ,					snake_case_ =   self.model_tester.prepare_config_and_inputs_for_common()
       snake_case_ =   inputs_dict['''pixel_values'''].shape[0]

       for backbone_class in self.all_model_classes:
             snake_case_ =   backbone_class(_UpperCAmelCase      )
             backbone.to(_UpperCAmelCase      )
             backbone.eval()

             snake_case_ =   backbone(**_UpperCAmelCase      )

             # Test default outputs and verify feature maps
             self.assertIsInstance(outputs.feature_maps	,   _UpperCAmelCase      )
             self.assertTrue(len(outputs.feature_maps      ) == len(backbone.channels      )      )
             for feature_map, n_channels in zip(outputs.feature_maps	,   backbone.channels      ):
                   self.assertTrue(feature_map.shape[:2]	,   (batch_size, n_channels)      )
             self.assertIsNone(outputs.hidden_states      )
             self.assertIsNone(outputs.attentions      )

             # Test output_hidden_states=True
             snake_case_ =   backbone(**_UpperCAmelCase	,   output_hidden_states=_UpperCAmelCase      )
             self.assertIsNotNone(outputs.hidden_states      )
             self.assertTrue(len(outputs.hidden_states      )	,   len(backbone.stage_names      )      )
             # We skip the stem layer
             for hidden_states, n_channels in zip(outputs.hidden_states[1:]	,   backbone.channels      ):
                   for hidden_state in hidden_states:
                         # Hidden states are in the format (batch_size, (height * width), n_channels)
                         snake_case_ ,					snake_case_ ,					snake_case_ =   hidden_state.shape
                         self.assertTrue((h_batch_size, h_n_channels)	,   (batch_size, n_channels)      )

            # Test output_attentions=True
             if self.has_attentions:
                   snake_case_ =   backbone(**_UpperCAmelCase	,   output_attentions=_UpperCAmelCase      )
                   self.assertIsNotNone(outputs.attentions      )



# 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

UpperCAmelCase    =					get_logger()

UpperCAmelCase    =					None







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

 '''simple docstring'''

 def __init__(			self	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   **_UpperCAmelCase      ):
       super().__init__(features=_UpperCAmelCase      )
       import jax
       from jaxlib.xla_client import Device

       if isinstance(_UpperCAmelCase	,   _UpperCAmelCase      ):
             raise ValueError(
                 F'''Expected {device} to be a `str` not {type(_UpperCAmelCase      )}, 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`.'''      )
       snake_case_ =   device if isinstance(_UpperCAmelCase	,   _UpperCAmelCase      ) 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:
             snake_case_ =   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]      )}.'''      )
             snake_case_ =   str(jax.devices()[0]      )
       snake_case_ =   jnp_array_kwargs



 @staticmethod
 def 		UpperCamelCase__  (			):
       import jax

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



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       import jax
       import jax.numpy as jnp

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



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       import jax
       import jax.numpy as jnp

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

       snake_case_ =   {}

       if isinstance(_UpperCAmelCase	,   (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:
                   snake_case_ =   {'''dtype''': jnp.intaa}
             else:
                   snake_case_ =   {'''dtype''': jnp.intaa}
       elif isinstance(_UpperCAmelCase	,   (np.number, np.ndarray)      ) and np.issubdtype(value.dtype	,   np.floating      ):
             snake_case_ =   {'''dtype''': jnp.floataa}
       elif config.PIL_AVAILABLE and "PIL" in sys.modules:
             import PIL.Image

             if isinstance(_UpperCAmelCase	,   PIL.Image.Image      ):
                   snake_case_ =   np.asarray(_UpperCAmelCase      )

        # 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:
             snake_case_ =   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(_UpperCAmelCase	,   **{**default_dtype, **self.jnp_array_kwargs}      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       import jax

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

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



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       return map_nested(self._recursive_tensorize	,   _UpperCAmelCase	,   map_list=_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       snake_case_ =   self.numpy_arrow_extractor().extract_row(_UpperCAmelCase      )
       snake_case_ =   self.python_features_decoder.decode_row(_UpperCAmelCase      )
       return self.recursive_tensorize(_UpperCAmelCase      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       snake_case_ =   self.numpy_arrow_extractor().extract_column(_UpperCAmelCase      )
       snake_case_ =   self.python_features_decoder.decode_column(_UpperCAmelCase	,   pa_table.column_names[0]      )
       snake_case_ =   self.recursive_tensorize(_UpperCAmelCase      )
       snake_case_ =   self._consolidate(_UpperCAmelCase      )
       return column





 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase      ):
       snake_case_ =   self.numpy_arrow_extractor().extract_batch(_UpperCAmelCase      )
       snake_case_ =   self.python_features_decoder.decode_batch(_UpperCAmelCase      )
       snake_case_ =   self.recursive_tensorize(_UpperCAmelCase      )
       for column_name in batch:
             snake_case_ =   self._consolidate(batch[column_name]      )
       return batch