






'''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

lowerCamelCase_												=					get_logger()

lowerCamelCase_												=					None

class    _UpperCAmelCase			(     TensorFormatter[Mapping, '''jax.Array''', Mapping]						):




   """simple docstring"""

   def __init__(		self :	int      ,      __UpperCAmelCase :	int=None      ,      __UpperCAmelCase :	List[str]=None      ,      **__UpperCAmelCase :	Optional[int]		):





         '''simple docstring'''





         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`."		)
         _A			     =	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:
               _A			     =	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]		)}.'''		)
               _A			     =	str(jax.devices()[0]		)
         _A			     =	jnp_array_kwargs

   @staticmethod
   def 						lowerCAmelCase (		):





         '''simple docstring'''





         import jax

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

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





         '''simple docstring'''





         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 						lowerCAmelCase (		self :	int      ,      __UpperCAmelCase :	List[str]		):





         '''simple docstring'''





         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()

         _A			     =	{}

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

               if isinstance(__UpperCAmelCase      ,      PIL.Image.Image		):
                     _A			     =	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:
               _A			     =	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 						lowerCAmelCase (		self :	Optional[Any]      ,      __UpperCAmelCase :	int		):





         '''simple docstring'''





         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		):
               _A			     =	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 						lowerCAmelCase (		self :	List[Any]      ,      __UpperCAmelCase :	dict		):





         '''simple docstring'''





         return map_nested(self._recursive_tensorize      ,      __UpperCAmelCase      ,      map_list=__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	str      ,      __UpperCAmelCase :	pa.Table		):





         '''simple docstring'''





         _A			     =	self.numpy_arrow_extractor().extract_row(__UpperCAmelCase		)
         _A			     =	self.python_features_decoder.decode_row(__UpperCAmelCase		)
         return self.recursive_tensorize(__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	List[str]      ,      __UpperCAmelCase :	pa.Table		):





         '''simple docstring'''





         _A			     =	self.numpy_arrow_extractor().extract_column(__UpperCAmelCase		)
         _A			     =	self.python_features_decoder.decode_column(__UpperCAmelCase      ,      pa_table.column_names[0]		)
         _A			     =	self.recursive_tensorize(__UpperCAmelCase		)
         _A			     =	self._consolidate(__UpperCAmelCase		)
         return column

   def 						lowerCAmelCase (		self :	Union[str, Any]      ,      __UpperCAmelCase :	pa.Table		):





         '''simple docstring'''





         _A			     =	self.numpy_arrow_extractor().extract_batch(__UpperCAmelCase		)
         _A			     =	self.python_features_decoder.decode_batch(__UpperCAmelCase		)
         _A			     =	self.recursive_tensorize(__UpperCAmelCase		)
         for column_name in batch:
               _A			     =	self._consolidate(batch[column_name]		)
         return batch













from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


lowercase  							=					{
    "configuration_tapas": ["TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP", "TapasConfig"],
    "tokenization_tapas": ["TapasTokenizer"],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	lowercase  							=					[
	    "TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST",
	    "TapasForMaskedLM",
	    "TapasForQuestionAnswering",
	    "TapasForSequenceClassification",
	    "TapasModel",
	    "TapasPreTrainedModel",
	    "load_tf_weights_in_tapas",
	]
try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	lowercase  							=					[
	    "TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST",
	    "TFTapasForMaskedLM",
	    "TFTapasForQuestionAnswering",
	    "TFTapasForSequenceClassification",
	    "TFTapasModel",
	    "TFTapasPreTrainedModel",
	]


if TYPE_CHECKING:
	from .configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
	from .tokenization_tapas import TapasTokenizer

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tapas import (
		    TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TapasForMaskedLM,
		    TapasForQuestionAnswering,
		    TapasForSequenceClassification,
		    TapasModel,
		    TapasPreTrainedModel,
		    load_tf_weights_in_tapas,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_tapas import (
		    TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFTapasForMaskedLM,
		    TFTapasForQuestionAnswering,
		    TFTapasForSequenceClassification,
		    TFTapasModel,
		    TFTapasPreTrainedModel,
		)


else:
	import sys

	lowercase  							=					_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





