


'''simple docstring'''
import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.generation import (
		    FlaxForcedBOSTokenLogitsProcessor,
		    FlaxForcedEOSTokenLogitsProcessor,
		    FlaxLogitsProcessorList,
		    FlaxMinLengthLogitsProcessor,
		    FlaxTemperatureLogitsWarper,
		    FlaxTopKLogitsWarper,
		    FlaxTopPLogitsWarper,
		)
@require_flax
class        lowercase     (							unittest.TestCase ):

		"""simple docstring"""

		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_							):

								'''simple docstring'''


								UpperCamelCase__       :Tuple							 =			jnp.ones((batch_size, length)							) / length
								return scores




		def    lowerCAmelCase__    (		self							):

								'''simple docstring'''


								UpperCamelCase__       :Optional[int]							 =			None
								UpperCamelCase__       :Any							 =			20

								UpperCamelCase__       :str							 =			self._get_uniform_logits(batch_size=2   ,       length=lowerCAmelCase__							)

								# tweak scores to not be uniform anymore
								UpperCamelCase__       :str							 =			scores.at[1, 5].set((1 / length) + 0.1							)  # peak, 1st batch
								UpperCamelCase__       :Optional[Any]							 =			scores.at[1, 10].set((1 / length) - 0.4							)  # valley, 1st batch

								# compute softmax
								UpperCamelCase__       :Union[str, Any]							 =			jax.nn.softmax(lowerCAmelCase__   ,       axis=-1							)

								UpperCamelCase__       :List[Any]							 =			FlaxTemperatureLogitsWarper(temperature=0.5							)
								UpperCamelCase__       :Any							 =			FlaxTemperatureLogitsWarper(temperature=1.3							)

								UpperCamelCase__       :Union[str, Any]							 =			jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase__   ,       scores.copy()   ,       cur_len=lowerCAmelCase__							)   ,       axis=-1							)
								UpperCamelCase__       :Dict							 =			jax.nn.softmax(temp_dist_warper_smoother(lowerCAmelCase__   ,       scores.copy()   ,       cur_len=lowerCAmelCase__							)   ,       axis=-1							)

								# uniform distribution stays uniform
								self.assertTrue(jnp.allclose(probs[0, :]   ,       warped_prob_sharp[0, :]   ,       atol=1e-3							)							)
								self.assertTrue(jnp.allclose(probs[0, :]   ,       warped_prob_smooth[0, :]   ,       atol=1e-3							)							)

								# sharp peaks get higher, valleys get lower
								self.assertLess(probs[1, :].max()   ,       warped_prob_sharp[1, :].max()							)
								self.assertGreater(probs[1, :].min()   ,       warped_prob_sharp[1, :].min()							)

								# smooth peaks get lower, valleys get higher
								self.assertGreater(probs[1, :].max()   ,       warped_prob_smooth[1, :].max()							)
								self.assertLess(probs[1, :].min()   ,       warped_prob_smooth[1, :].min()							)




		def    lowerCAmelCase__    (		self							):

								'''simple docstring'''


								UpperCamelCase__       :List[str]							 =			None
								UpperCamelCase__       :Optional[int]							 =			10
								UpperCamelCase__       :int							 =			2

								# create ramp distribution
								UpperCamelCase__       :str							 =			np.broadcast_to(np.arange(lowerCAmelCase__							)[None, :]   ,       (batch_size, vocab_size)							).copy()
								UpperCamelCase__       :List[Any]							 =			ramp_logits[1:, : vocab_size // 2] + vocab_size

								UpperCamelCase__       :Optional[int]							 =			FlaxTopKLogitsWarper(3							)

								UpperCamelCase__       :Union[str, Any]							 =			top_k_warp(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)

								# check that correct tokens are filtered
								self.assertListEqual(jnp.isinf(scores[0]							).tolist()   ,       7 * [True] + 3 * [False]							)
								self.assertListEqual(jnp.isinf(scores[1]							).tolist()   ,       2 * [True] + 3 * [False] + 5 * [True]							)

								# check special case
								UpperCamelCase__       :Optional[int]							 =			5
								UpperCamelCase__       :Union[str, Any]							 =			FlaxTopKLogitsWarper(top_k=1   ,       filter_value=0.0   ,       min_tokens_to_keep=3							)

								UpperCamelCase__       :Optional[Any]							 =			np.broadcast_to(np.arange(lowerCAmelCase__							)[None, :]   ,       (batch_size, length)							).copy()
								UpperCamelCase__       :List[str]							 =			top_k_warp_safety_check(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)

								# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
								self.assertListEqual((scores == 0.0).sum(axis=-1							).tolist()   ,       [2, 2]							)




		def    lowerCAmelCase__    (		self							):

								'''simple docstring'''


								UpperCamelCase__       :Optional[int]							 =			None
								UpperCamelCase__       :Any							 =			10
								UpperCamelCase__       :Union[str, Any]							 =			2

								# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
								UpperCamelCase__       :List[Any]							 =			np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]							)							)

								UpperCamelCase__       :Optional[int]							 =			FlaxTopPLogitsWarper(0.8							)
								UpperCamelCase__       :Optional[Any]							 =			np.exp(top_p_warp(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)							)

								# dist should be filtered to keep min num values so that sum is >= top_p
								# exp (-inf) => 0
								UpperCamelCase__       :Optional[Any]							 =			np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]							)
								self.assertTrue(np.allclose(lowerCAmelCase__   ,       lowerCAmelCase__   ,       atol=1e-3							)							)

								# check edge cases with negative and extreme logits
								UpperCamelCase__       :List[str]							 =			np.broadcast_to(np.arange(lowerCAmelCase__							)[None, :]   ,       (batch_size, vocab_size)							).copy() - (
								    vocab_size // 2
								)

								# make ramp_logits more extreme
								UpperCamelCase__       :int							 =			ramp_logits[1] * 100.0

								# make sure at least 2 tokens are kept
								UpperCamelCase__       :Union[str, Any]							 =			FlaxTopPLogitsWarper(0.9   ,       min_tokens_to_keep=2   ,       filter_value=0.0							)
								UpperCamelCase__       :List[str]							 =			top_p_warp(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)

								# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
								self.assertListEqual((filtered_dist != 0.0).sum(axis=-1							).tolist()   ,       [3, 2]							)




		def    lowerCAmelCase__    (		self							):

								'''simple docstring'''


								UpperCamelCase__       :int							 =			20
								UpperCamelCase__       :List[str]							 =			4
								UpperCamelCase__       :Dict							 =			0

								UpperCamelCase__       :Any							 =			FlaxMinLengthLogitsProcessor(min_length=10   ,       eos_token_id=lowerCAmelCase__							)

								# check that min length is applied at length 5
								UpperCamelCase__       :Optional[Any]							 =			ids_tensor((batch_size, 20)   ,       vocab_size=20							)
								UpperCamelCase__       :List[str]							 =			5
								UpperCamelCase__       :Tuple							 =			self._get_uniform_logits(lowerCAmelCase__   ,       lowerCAmelCase__							)
								UpperCamelCase__       :List[Any]							 =			min_dist_processor(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()   ,       4 * [-float('''inf'''							)]							)

								# check that min length is not applied anymore at length 15
								UpperCamelCase__       :int							 =			self._get_uniform_logits(lowerCAmelCase__   ,       lowerCAmelCase__							)
								UpperCamelCase__       :List[Any]							 =			15
								UpperCamelCase__       :int							 =			min_dist_processor(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								self.assertFalse(jnp.isinf(lowerCAmelCase__							).any()							)




		def    lowerCAmelCase__    (		self							):

								'''simple docstring'''


								UpperCamelCase__       :Optional[Any]							 =			20
								UpperCamelCase__       :int							 =			4
								UpperCamelCase__       :Tuple							 =			0

								UpperCamelCase__       :Optional[int]							 =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__							)

								# check that all scores are -inf except the bos_token_id score
								UpperCamelCase__       :Any							 =			ids_tensor((batch_size, 1)   ,       vocab_size=20							)
								UpperCamelCase__       :Union[str, Any]							 =			1
								UpperCamelCase__       :Any							 =			self._get_uniform_logits(lowerCAmelCase__   ,       lowerCAmelCase__							)
								UpperCamelCase__       :Optional[Any]							 =			logits_processor(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]							).all()							)
								self.assertListEqual(scores[:, bos_token_id].tolist()   ,       4 * [0]							)  # score for bos_token_id shold be zero

								# check that bos_token_id is not forced if current length is greater than 1
								UpperCamelCase__       :Any							 =			3
								UpperCamelCase__       :Optional[int]							 =			self._get_uniform_logits(lowerCAmelCase__   ,       lowerCAmelCase__							)
								UpperCamelCase__       :str							 =			logits_processor(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								self.assertFalse(jnp.isinf(lowerCAmelCase__							).any()							)




		def    lowerCAmelCase__    (		self							):

								'''simple docstring'''


								UpperCamelCase__       :str							 =			20
								UpperCamelCase__       :List[Any]							 =			4
								UpperCamelCase__       :int							 =			0
								UpperCamelCase__       :Optional[int]							 =			5

								UpperCamelCase__       :int							 =			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__   ,       eos_token_id=lowerCAmelCase__							)

								# check that all scores are -inf except the eos_token_id when max_length is reached
								UpperCamelCase__       :Optional[int]							 =			ids_tensor((batch_size, 4)   ,       vocab_size=20							)
								UpperCamelCase__       :Optional[int]							 =			4
								UpperCamelCase__       :str							 =			self._get_uniform_logits(lowerCAmelCase__   ,       lowerCAmelCase__							)
								UpperCamelCase__       :Optional[Any]							 =			logits_processor(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]							).all()							)
								self.assertListEqual(scores[:, eos_token_id].tolist()   ,       4 * [0]							)  # score for eos_token_id should be zero

								# check that eos_token_id is not forced if max_length is not reached
								UpperCamelCase__       :Dict							 =			3
								UpperCamelCase__       :Dict							 =			self._get_uniform_logits(lowerCAmelCase__   ,       lowerCAmelCase__							)
								UpperCamelCase__       :Any							 =			logits_processor(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								self.assertFalse(jnp.isinf(lowerCAmelCase__							).any()							)




		def    lowerCAmelCase__    (		self							):

								'''simple docstring'''


								UpperCamelCase__       :int							 =			4
								UpperCamelCase__       :Optional[int]							 =			10
								UpperCamelCase__       :List[Any]							 =			15
								UpperCamelCase__       :Optional[Any]							 =			2
								UpperCamelCase__       :Union[str, Any]							 =			1
								UpperCamelCase__       :Any							 =			15

								# dummy input_ids and scores
								UpperCamelCase__       :Optional[int]							 =			ids_tensor((batch_size, sequence_length)   ,       lowerCAmelCase__							)
								UpperCamelCase__       :Optional[int]							 =			input_ids.copy()

								UpperCamelCase__       :str							 =			self._get_uniform_logits(lowerCAmelCase__   ,       lowerCAmelCase__							)
								UpperCamelCase__       :Union[str, Any]							 =			scores.copy()

								# instantiate all dist processors
								UpperCamelCase__       :str							 =			FlaxTemperatureLogitsWarper(temperature=0.5							)
								UpperCamelCase__       :str							 =			FlaxTopKLogitsWarper(3							)
								UpperCamelCase__       :Optional[int]							 =			FlaxTopPLogitsWarper(0.8							)

								# instantiate all logits processors
								UpperCamelCase__       :Optional[Any]							 =			FlaxMinLengthLogitsProcessor(min_length=10   ,       eos_token_id=lowerCAmelCase__							)
								UpperCamelCase__       :Dict							 =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__							)
								UpperCamelCase__       :List[Any]							 =			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__   ,       eos_token_id=lowerCAmelCase__							)

								UpperCamelCase__       :str							 =			10

								# no processor list
								UpperCamelCase__       :Optional[int]							 =			temp_dist_warp(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								UpperCamelCase__       :List[str]							 =			top_k_warp(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								UpperCamelCase__       :str							 =			top_p_warp(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								UpperCamelCase__       :Tuple							 =			min_dist_proc(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								UpperCamelCase__       :Optional[Any]							 =			bos_dist_proc(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
								UpperCamelCase__       :Union[str, Any]							 =			eos_dist_proc(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)

								# with processor list
								UpperCamelCase__       :List[str]							 =			FlaxLogitsProcessorList(
								    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
								UpperCamelCase__       :Tuple							 =			processor(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)

								# scores should be equal
								self.assertTrue(jnp.allclose(lowerCAmelCase__   ,       lowerCAmelCase__   ,       atol=1e-3							)							)

								# input_ids should never be changed
								self.assertListEqual(input_ids.tolist()   ,       input_ids_comp.tolist()							)




		def    lowerCAmelCase__    (		self							):

								'''simple docstring'''


								UpperCamelCase__       :Dict							 =			4
								UpperCamelCase__       :int							 =			10
								UpperCamelCase__       :Any							 =			15
								UpperCamelCase__       :List[Any]							 =			2
								UpperCamelCase__       :Any							 =			1
								UpperCamelCase__       :Optional[int]							 =			15

								# dummy input_ids and scores
								UpperCamelCase__       :List[Any]							 =			ids_tensor((batch_size, sequence_length)   ,       lowerCAmelCase__							)
								UpperCamelCase__       :List[str]							 =			input_ids.copy()

								UpperCamelCase__       :Tuple							 =			self._get_uniform_logits(lowerCAmelCase__   ,       lowerCAmelCase__							)
								UpperCamelCase__       :List[Any]							 =			scores.copy()

								# instantiate all dist processors
								UpperCamelCase__       :Any							 =			FlaxTemperatureLogitsWarper(temperature=0.5							)
								UpperCamelCase__       :Dict							 =			FlaxTopKLogitsWarper(3							)
								UpperCamelCase__       :List[str]							 =			FlaxTopPLogitsWarper(0.8							)

								# instantiate all logits processors
								UpperCamelCase__       :str							 =			FlaxMinLengthLogitsProcessor(min_length=10   ,       eos_token_id=lowerCAmelCase__							)
								UpperCamelCase__       :int							 =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__							)
								UpperCamelCase__       :Union[str, Any]							 =			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__   ,       eos_token_id=lowerCAmelCase__							)

								UpperCamelCase__       :Dict							 =			10

								# no processor list
								def run_no_processor_list(UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_							):
														UpperCamelCase__       :List[str]							 =			temp_dist_warp(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
														UpperCamelCase__       :Dict							 =			top_k_warp(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
														UpperCamelCase__       :List[Any]							 =			top_p_warp(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
														UpperCamelCase__       :Union[str, Any]							 =			min_dist_proc(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
														UpperCamelCase__       :List[str]							 =			bos_dist_proc(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
														UpperCamelCase__       :int							 =			eos_dist_proc(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
														return scores

								# with processor list
								def run_processor_list(UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_							):
														UpperCamelCase__       :Union[str, Any]							 =			FlaxLogitsProcessorList(
														    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
														UpperCamelCase__       :Dict							 =			processor(lowerCAmelCase__   ,       lowerCAmelCase__   ,       cur_len=lowerCAmelCase__							)
														return scores

								UpperCamelCase__       :Any							 =			jax.jit(lowerCAmelCase__							)
								UpperCamelCase__       :List[str]							 =			jax.jit(lowerCAmelCase__							)

								UpperCamelCase__       :Tuple							 =			jitted_run_no_processor_list(lowerCAmelCase__   ,       lowerCAmelCase__   ,       lowerCAmelCase__							)
								UpperCamelCase__       :Dict							 =			jitted_run_processor_list(lowerCAmelCase__   ,       lowerCAmelCase__   ,       lowerCAmelCase__							)

								# scores should be equal
								self.assertTrue(jnp.allclose(lowerCAmelCase__   ,       lowerCAmelCase__   ,       atol=1e-3							)							)

								# input_ids should never be changed
								self.assertListEqual(input_ids.tolist()   ,       input_ids_comp.tolist()							)



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

__snake_case											=					get_logger()

__snake_case											=					None
class        lowercase     (							TensorFormatter[Mapping, 'jax.Array', Mapping] ):

		"""simple docstring"""

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

							'''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`.'''							)
							UpperCamelCase__       :Tuple							 =			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:
												UpperCamelCase__       :Optional[Any]							 =			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__       :Optional[int]							 =			str(jax.devices()[0]							)
							UpperCamelCase__       :Tuple							 =			jnp_array_kwargs




		@staticmethod
		def    lowerCAmelCase__    (		):

							'''simple docstring'''


							import jax

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




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''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   ,       UpperCamelCase_							):

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

							UpperCamelCase__       :Optional[int]							 =			{}

							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:
																	UpperCamelCase__       :List[str]							 =			{'''dtype''': jnp.intaa}
												else:
																	UpperCamelCase__       :Union[str, Any]							 =			{'''dtype''': jnp.intaa}
							elif isinstance(UpperCamelCase_   ,       (np.number, np.ndarray)							) and np.issubdtype(value.dtype   ,       np.floating							):
												UpperCamelCase__       :Optional[Any]							 =			{'''dtype''': jnp.floataa}
							elif config.PIL_AVAILABLE and "PIL" in sys.modules:
												import PIL.Image

												if isinstance(UpperCamelCase_   ,       PIL.Image.Image							):
																	UpperCamelCase__       :str							 =			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:
												UpperCamelCase__       :Dict							 =			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   ,       UpperCamelCase_							):

							'''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							):
												UpperCamelCase__       :int							 =			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   ,       UpperCamelCase_							):

							'''simple docstring'''


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




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			self.numpy_arrow_extractor().extract_row(UpperCamelCase_							)
							UpperCamelCase__       :List[Any]							 =			self.python_features_decoder.decode_row(UpperCamelCase_							)
							return self.recursive_tensorize(UpperCamelCase_							)




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			self.numpy_arrow_extractor().extract_column(UpperCamelCase_							)
							UpperCamelCase__       :Tuple							 =			self.python_features_decoder.decode_column(UpperCamelCase_   ,       pa_table.column_names[0]							)
							UpperCamelCase__       :Dict							 =			self.recursive_tensorize(UpperCamelCase_							)
							UpperCamelCase__       :str							 =			self._consolidate(UpperCamelCase_							)
							return column




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :int							 =			self.numpy_arrow_extractor().extract_batch(UpperCamelCase_							)
							UpperCamelCase__       :Union[str, Any]							 =			self.python_features_decoder.decode_batch(UpperCamelCase_							)
							UpperCamelCase__       :List[str]							 =			self.recursive_tensorize(UpperCamelCase_							)
							for column_name in batch:
												UpperCamelCase__       :Optional[int]							 =			self._consolidate(batch[column_name]							)
							return batch