"""simple docstring"""
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__   =					get_logger()

SCREAMING_SNAKE_CASE__   =					None



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



				"""simple docstring"""

				def __init__(			self						,      lowerCAmelCase=None						,      lowerCAmelCase=None						,      **lowerCAmelCase      ):






										"""simple docstring"""

										super().__init__(features=_lowerCamelCase      )
										import jax
										from jaxlib.xla_client import Device

										if isinstance(_lowerCamelCase						,      _lowerCamelCase      ):
																raise ValueError(
																    F"""Expected {device} to be a `str` not {type(_lowerCamelCase      )}, 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(_lowerCamelCase						,      _lowerCamelCase      ) 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 						snake_case	(			):






										"""simple docstring"""

										import jax

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

				def 						snake_case	(			self						,      lowerCAmelCase      ):






										"""simple docstring"""

										import jax
										import jax.numpy as jnp

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

				def 						snake_case	(			self						,      lowerCAmelCase      ):






										"""simple docstring"""

										import jax
										import jax.numpy as jnp

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

										snake_case        =     {}

										if isinstance(_lowerCamelCase						,      (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(_lowerCamelCase						,      (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(_lowerCamelCase						,      PIL.Image.Image      ):
																						snake_case        =     np.asarray(_lowerCamelCase      )

        # 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(_lowerCamelCase						,      **{**default_dtype, **self.jnp_array_kwargs}      )

				def 						snake_case	(			self						,      lowerCAmelCase      ):






										"""simple docstring"""

										import jax

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

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

				def 						snake_case	(			self						,      lowerCAmelCase      ):






										"""simple docstring"""

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

				def 						snake_case	(			self						,      lowerCAmelCase      ):






										"""simple docstring"""

										snake_case        =     self.numpy_arrow_extractor().extract_row(_lowerCamelCase      )
										snake_case        =     self.python_features_decoder.decode_row(_lowerCamelCase      )
										return self.recursive_tensorize(_lowerCamelCase      )

				def 						snake_case	(			self						,      lowerCAmelCase      ):






										"""simple docstring"""

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

				def 						snake_case	(			self						,      lowerCAmelCase      ):






										"""simple docstring"""

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




"""simple docstring"""
from typing import TYPE_CHECKING

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


SCREAMING_SNAKE_CASE__   =					{
    "configuration_xlm": ["XLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "XLMConfig", "XLMOnnxConfig"],
    "tokenization_xlm": ["XLMTokenizer"],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					SCREAMING_SNAKE_CASE__   =					[
					    "XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "XLMForMultipleChoice",
					    "XLMForQuestionAnswering",
					    "XLMForQuestionAnsweringSimple",
					    "XLMForSequenceClassification",
					    "XLMForTokenClassification",
					    "XLMModel",
					    "XLMPreTrainedModel",
					    "XLMWithLMHeadModel",
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					SCREAMING_SNAKE_CASE__   =					[
					    "TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "TFXLMForMultipleChoice",
					    "TFXLMForQuestionAnsweringSimple",
					    "TFXLMForSequenceClassification",
					    "TFXLMForTokenClassification",
					    "TFXLMMainLayer",
					    "TFXLMModel",
					    "TFXLMPreTrainedModel",
					    "TFXLMWithLMHeadModel",
					]


if TYPE_CHECKING:
					from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
					from .tokenization_xlm import XLMTokenizer

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_xlm import (
										    XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
										    XLMForMultipleChoice,
										    XLMForQuestionAnswering,
										    XLMForQuestionAnsweringSimple,
										    XLMForSequenceClassification,
										    XLMForTokenClassification,
										    XLMModel,
										    XLMPreTrainedModel,
										    XLMWithLMHeadModel,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_xlm import (
										    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFXLMForMultipleChoice,
										    TFXLMForQuestionAnsweringSimple,
										    TFXLMForSequenceClassification,
										    TFXLMForTokenClassification,
										    TFXLMMainLayer,
										    TFXLMModel,
										    TFXLMPreTrainedModel,
										    TFXLMWithLMHeadModel,
										)

else:
					import sys

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


