


'''simple docstring'''
from __future__ import annotations

import unittest

import numpy as np

from transformers import OPTConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, slow

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


if is_tf_available():
		import tensorflow as tf

		from transformers import GPTaTokenizer, TFOPTForCausalLM, TFOPTModel


def       a				(      __a			, __a			, __a=None			, __a=None       )  ->     Optional[Any]:



					'''simple docstring'''






					if attention_mask is None:
										UpperCamelCase__       :Dict							 =			tf.cast(tf.math.not_equal(__a			, config.pad_token_id       )			, tf.inta       )
					return {"input_ids": input_ids, "attention_mask": attention_mask}
@require_tf
class        lowercase     :

		"""simple docstring"""

		_a   			=    OPTConfig
		_a   			=    {}
		_a   			=    'gelu'




		def __init__(		self   ,       UpperCamelCase_   ,       UpperCamelCase_=13   ,       UpperCamelCase_=7   ,       UpperCamelCase_=True   ,       UpperCamelCase_=False   ,       UpperCamelCase_=99   ,       UpperCamelCase_=16   ,       UpperCamelCase_=2   ,       UpperCamelCase_=4   ,       UpperCamelCase_=4   ,       UpperCamelCase_="gelu"   ,       UpperCamelCase_=0.1   ,       UpperCamelCase_=0.1   ,       UpperCamelCase_=20   ,       UpperCamelCase_=2   ,       UpperCamelCase_=1   ,       UpperCamelCase_=0   ,       UpperCamelCase_=16   ,       UpperCamelCase_=16   ,       ):

							'''simple docstring'''


							UpperCamelCase__       :List[str]							 =			parent
							UpperCamelCase__       :str							 =			batch_size
							UpperCamelCase__       :Optional[Any]							 =			seq_length
							UpperCamelCase__       :List[str]							 =			is_training
							UpperCamelCase__       :int							 =			use_labels
							UpperCamelCase__       :List[Any]							 =			vocab_size
							UpperCamelCase__       :int							 =			hidden_size
							UpperCamelCase__       :Optional[int]							 =			num_hidden_layers
							UpperCamelCase__       :Union[str, Any]							 =			num_attention_heads
							UpperCamelCase__       :List[str]							 =			intermediate_size
							UpperCamelCase__       :Tuple							 =			hidden_act
							UpperCamelCase__       :Any							 =			hidden_dropout_prob
							UpperCamelCase__       :Tuple							 =			attention_probs_dropout_prob
							UpperCamelCase__       :int							 =			max_position_embeddings
							UpperCamelCase__       :Union[str, Any]							 =			eos_token_id
							UpperCamelCase__       :Optional[int]							 =			pad_token_id
							UpperCamelCase__       :Optional[Any]							 =			bos_token_id
							UpperCamelCase__       :Tuple							 =			embed_dim
							UpperCamelCase__       :List[Any]							 =			word_embed_proj_dim
							UpperCamelCase__       :List[Any]							 =			False




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Optional[Any]							 =			ids_tensor([self.batch_size, self.seq_length - 1]   ,       self.vocab_size							)
							UpperCamelCase__       :Optional[int]							 =			tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size							)   ,       1							)
							UpperCamelCase__       :int							 =			tf.concat([input_ids, eos_tensor]   ,       axis=1							)

							UpperCamelCase__       :Dict							 =			self.config_cls(
							    vocab_size=self.vocab_size   ,       hidden_size=self.hidden_size   ,       num_hidden_layers=self.num_hidden_layers   ,       num_attention_heads=self.num_attention_heads   ,       ffn_dim=self.intermediate_size   ,       dropout=self.hidden_dropout_prob   ,       attention_dropout=self.attention_probs_dropout_prob   ,       max_position_embeddings=self.max_position_embeddings   ,       eos_token_id=self.eos_token_id   ,       bos_token_id=self.bos_token_id   ,       pad_token_id=self.pad_token_id   ,       embed_dim=self.embed_dim   ,       word_embed_proj_dim=self.word_embed_proj_dim   ,       is_encoder_decoder=UpperCamelCase_   ,       **self.config_updates   ,       )
							UpperCamelCase__       :int							 =			prepare_opt_inputs_dict(UpperCamelCase_   ,       UpperCamelCase_							)
							return config, inputs_dict




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :str							 =			TFOPTModel(config=UpperCamelCase_							)
							UpperCamelCase__       :Dict							 =			inputs_dict['''input_ids''']

							UpperCamelCase__       :Optional[Any]							 =			input_ids[:1, :]
							UpperCamelCase__       :Union[str, Any]							 =			inputs_dict['''attention_mask'''][:1, :]
							UpperCamelCase__       :Optional[Any]							 =			1

							# first forward pass
							UpperCamelCase__       :str							 =			model(UpperCamelCase_   ,       attention_mask=UpperCamelCase_   ,       use_cache=UpperCamelCase_							)

							UpperCamelCase__ ,    UpperCamelCase__       :Any							 =			outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							UpperCamelCase__       :List[Any]							 =			ids_tensor((self.batch_size, 3)   ,       config.vocab_size							)
							UpperCamelCase__       :List[str]							 =			tf.cast(ids_tensor((self.batch_size, 3)   ,       2							)   ,       tf.inta							)

							# append to next input_ids and
							UpperCamelCase__       :str							 =			tf.concat([input_ids, next_tokens]   ,       axis=-1							)
							UpperCamelCase__       :str							 =			tf.concat([attention_mask, next_attn_mask]   ,       axis=-1							)

							UpperCamelCase__       :Tuple							 =			model(UpperCamelCase_   ,       attention_mask=UpperCamelCase_							)[0]
							UpperCamelCase__       :Dict							 =			model(UpperCamelCase_   ,       attention_mask=UpperCamelCase_   ,       past_key_values=UpperCamelCase_							)[0]

							self.parent.assertEqual(next_tokens.shape[1]   ,       output_from_past.shape[1]							)

							# select random slice
							UpperCamelCase__       :Tuple							 =			int(ids_tensor((1,)   ,       output_from_past.shape[-1]							)							)
							UpperCamelCase__       :Any							 =			output_from_no_past[:, -3:, random_slice_idx]
							UpperCamelCase__       :Optional[Any]							 =			output_from_past[:, :, random_slice_idx]

							# test that outputs are equal for slice
							tf.debugging.assert_near(UpperCamelCase_   ,       UpperCamelCase_   ,       rtol=1e-3							)



@require_tf
class        lowercase     (							A__     ,		A__     ,		unittest.TestCase ):

		"""simple docstring"""

		_a   			=    (TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
		_a   			=    (TFOPTForCausalLM,) if is_tf_available() else ()
		_a   			=    (
		    {'feature-extraction': TFOPTModel, 'text-generation': TFOPTForCausalLM} if is_tf_available() else {}
		)
		_a   			=    False
		_a   			=    False
		_a   			=    False
		_a   			=    10




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Union[str, Any]							 =			TFOPTModelTester(self							)
							UpperCamelCase__       :Any							 =			ConfigTester(self   ,       config_class=UpperCamelCase_							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							self.config_tester.run_common_tests()




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Optional[Any]							 =			self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.check_decoder_model_past_large_inputs(*UpperCamelCase_							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__ ,    UpperCamelCase__       :str							 =			self.model_tester.prepare_config_and_inputs_for_common()

							def _get_word_embedding_weight(UpperCamelCase_   ,       UpperCamelCase_							):
												if hasattr(UpperCamelCase_   ,       '''weight'''							):
																	return embedding_layer.weight
												else:
																	# Here we build the word embeddings weights if not exists.
																	# And then we retry to get the attribute once built.
																	model.build()
																	if hasattr(UpperCamelCase_   ,       '''weight'''							):
																						return embedding_layer.weight
																	else:
																						return None

							for model_class in self.all_model_classes:
												for size in [config.vocab_size - 10, config.vocab_size + 10]:
																	# build the embeddings
																	UpperCamelCase__       :Optional[Any]							 =			model_class(config=UpperCamelCase_							)
																	UpperCamelCase__       :Any							 =			_get_word_embedding_weight(UpperCamelCase_   ,       model.get_input_embeddings()							)
																	UpperCamelCase__       :Optional[int]							 =			_get_word_embedding_weight(UpperCamelCase_   ,       model.get_output_embeddings()							)

																	# reshape the embeddings
																	model.resize_token_embeddings(UpperCamelCase_							)
																	UpperCamelCase__       :int							 =			_get_word_embedding_weight(UpperCamelCase_   ,       model.get_input_embeddings()							)
																	UpperCamelCase__       :Union[str, Any]							 =			_get_word_embedding_weight(UpperCamelCase_   ,       model.get_output_embeddings()							)

																	# check that the resized embeddings size matches the desired size.
																	UpperCamelCase__       :List[Any]							 =			size if size is not None else config.vocab_size

																	self.assertEqual(new_input_embeddings.shape[0]   ,       UpperCamelCase_							)

																	# check that weights remain the same after resizing
																	UpperCamelCase__       :Any							 =			True
																	for pa, pa in zip(old_input_embeddings.value()   ,       new_input_embeddings.value()							):
																						if tf.math.reduce_sum(tf.math.abs(pa - pa							)							) > 0:
																											UpperCamelCase__       :List[str]							 =			False
																	self.assertTrue(UpperCamelCase_							)

																	if old_output_embeddings is not None and new_output_embeddings is not None:
																						self.assertEqual(new_output_embeddings.shape[0]   ,       UpperCamelCase_							)

																						UpperCamelCase__       :List[str]							 =			True
																						for pa, pa in zip(old_output_embeddings.value()   ,       new_output_embeddings.value()							):
																											if tf.math.reduce_sum(tf.math.abs(pa - pa							)							) > 0:
																																UpperCamelCase__       :str							 =			False
																						self.assertTrue(UpperCamelCase_							)



def       a				(      __a       )  ->     int:



					'''simple docstring'''






					return tf.constant(__a			, dtype=tf.intaa       )
@require_tf
class        lowercase     (							unittest.TestCase ):

		"""simple docstring"""

		_a   			=    99




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Optional[int]							 =			tf.ones((4, 1)   ,       dtype=tf.intaa							) * 2
							UpperCamelCase__       :Dict							 =			tf.concat([ids_tensor((4, 6)   ,       self.vocab_size - 3							) + 3, eos_column_vector]   ,       axis=1							)
							UpperCamelCase__       :Union[str, Any]							 =			input_ids.shape[0]
							UpperCamelCase__       :Union[str, Any]							 =			OPTConfig(
							    vocab_size=self.vocab_size   ,       hidden_size=24   ,       num_hidden_layers=2   ,       num_attention_heads=2   ,       ffn_dim=32   ,       max_position_embeddings=48   ,       eos_token_id=2   ,       pad_token_id=1   ,       bos_token_id=0   ,       )
							return config, input_ids, batch_size



@require_sentencepiece
@require_tf
class        lowercase     (							unittest.TestCase ):

		"""simple docstring"""

		@slow
		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :List[str]							 =			TFOPTModel.from_pretrained('''facebook/opt-350m'''							)
							UpperCamelCase__       :Any							 =			_long_tensor([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]							)
							UpperCamelCase__       :int							 =			tf.not_equal(UpperCamelCase_   ,       model.config.pad_token_id							)
							with tf.GradientTape():
												UpperCamelCase__       :List[Any]							 =			model(input_ids=UpperCamelCase_   ,       attention_mask=UpperCamelCase_							).last_hidden_state
							UpperCamelCase__       :List[str]							 =			(1, 11, 512)
							self.assertEqual(output.shape   ,       UpperCamelCase_							)
							UpperCamelCase__       :List[str]							 =			tf.constant(
							    [[-0.2873, -1.9218, -0.3033], [-1.2710, -0.1338, -0.1902], [0.4095, 0.1214, -1.3121]]							)
							self.assertTrue(np.allclose(output[:, :3, :3]   ,       UpperCamelCase_   ,       atol=4e-3							)							)

							UpperCamelCase__       :Dict							 =			tf.function(UpperCamelCase_   ,       jit_compile=UpperCamelCase_							)
							UpperCamelCase__       :List[str]							 =			xla_generate(UpperCamelCase_   ,       UpperCamelCase_							)[0]
							self.assertTrue(np.allclose(output[:, :3, :3]   ,       UpperCamelCase_   ,       atol=4e-2							)							)



@require_tf
@slow
class        lowercase     (							unittest.TestCase ):

		"""simple docstring"""

		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							super().setUp()
							UpperCamelCase__       :Union[str, Any]							 =			'''facebook/opt-350m'''




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Dict							 =			TFOPTForCausalLM.from_pretrained(self.path_model							)
							UpperCamelCase__       :Any							 =			GPTaTokenizer.from_pretrained(self.path_model							)

							UpperCamelCase__       :Optional[Any]							 =			[
							    '''Today is a beautiful day and I want to''',
							    '''In the city of''',
							    '''Paris is the capital of France and''',
							    '''Computers and mobile phones have taken''',
							]
							# verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
							UpperCamelCase__       :List[str]							 =			tokenizer(UpperCamelCase_   ,       return_tensors='''tf'''   ,       padding=UpperCamelCase_   ,       add_special_tokens=UpperCamelCase_							)
							UpperCamelCase__       :Tuple							 =			tf.math.reduce_mean(model(inputs.input_ids   ,       attention_mask=inputs.attention_mask							)[0]   ,       axis=-1							)
							UpperCamelCase__       :List[str]							 =			tf.constant(
							    [
							        [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
							        [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
							        [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
							        [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
							    ]							)
							self.assertTrue(np.allclose(UpperCamelCase_   ,       UpperCamelCase_   ,       atol=1e-4							)							)

							UpperCamelCase__       :Any							 =			tf.function(UpperCamelCase_   ,       jit_compile=UpperCamelCase_							)
							UpperCamelCase__       :Optional[Any]							 =			tf.math.reduce_mean(xla_generate(inputs.input_ids   ,       attention_mask=inputs.attention_mask							)[0]   ,       axis=-1							)
							self.assertTrue(np.allclose(UpperCamelCase_   ,       UpperCamelCase_   ,       atol=1e-4							)							)



@require_tf
@slow
class        lowercase     (							unittest.TestCase ):

		"""simple docstring"""

		@property
		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							return [
							    "Today is a beautiful day and I want",
							    "In the city of",
							    "Paris is the capital of France and",
							    "Computers and mobile phones have taken",
							]




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :int							 =			'''facebook/opt-125m'''

							UpperCamelCase__       :Optional[Any]							 =			[
							    '''Today is a beautiful day and I want to''',
							    '''In the city of New York, the city''',
							    '''Paris is the capital of France and the capital''',
							    '''Computers and mobile phones have taken over the''',
							]

							UpperCamelCase__       :Any							 =			[]
							UpperCamelCase__       :Optional[Any]							 =			GPTaTokenizer.from_pretrained(UpperCamelCase_							)
							UpperCamelCase__       :int							 =			TFOPTForCausalLM.from_pretrained(UpperCamelCase_							)

							for prompt in self.prompts:
												UpperCamelCase__       :Optional[int]							 =			tokenizer(UpperCamelCase_   ,       return_tensors='''tf'''							).input_ids

												UpperCamelCase__       :Any							 =			model.generate(UpperCamelCase_   ,       max_length=10							)

												UpperCamelCase__       :Tuple							 =			tokenizer.batch_decode(UpperCamelCase_   ,       skip_special_tokens=UpperCamelCase_							)
												predicted_outputs += generated_string

							self.assertListEqual(UpperCamelCase_   ,       UpperCamelCase_							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :int							 =			'''facebook/opt-350m'''

							UpperCamelCase__       :Tuple							 =			GPTaTokenizer.from_pretrained(UpperCamelCase_							)
							UpperCamelCase__       :Union[str, Any]							 =			TFOPTForCausalLM.from_pretrained(UpperCamelCase_							)

							UpperCamelCase__       :Optional[Any]							 =			'''left'''

							# use different length sentences to test batching
							UpperCamelCase__       :int							 =			[
							    '''Hello, my dog is a little''',
							    '''Today, I''',
							]

							UpperCamelCase__       :Dict							 =			tokenizer(UpperCamelCase_   ,       return_tensors='''tf'''   ,       padding=UpperCamelCase_							)
							UpperCamelCase__       :Optional[int]							 =			inputs['''input_ids''']

							UpperCamelCase__       :Any							 =			model.generate(input_ids=UpperCamelCase_   ,       attention_mask=inputs['''attention_mask''']							)

							UpperCamelCase__       :Optional[Any]							 =			tokenizer(sentences[0]   ,       return_tensors='''tf'''							).input_ids
							UpperCamelCase__       :Optional[int]							 =			model.generate(input_ids=UpperCamelCase_							)

							UpperCamelCase__       :Optional[int]							 =			inputs_non_padded.shape[-1] - tf.math.reduce_sum(
							    tf.cast(inputs['''attention_mask'''][-1]   ,       tf.intaa							)							)
							UpperCamelCase__       :int							 =			tokenizer(sentences[1]   ,       return_tensors='''tf'''							).input_ids
							UpperCamelCase__       :int							 =			model.generate(input_ids=UpperCamelCase_   ,       max_length=model.config.max_length - num_paddings							)

							UpperCamelCase__       :List[Any]							 =			tokenizer.batch_decode(UpperCamelCase_   ,       skip_special_tokens=UpperCamelCase_							)
							UpperCamelCase__       :Dict							 =			tokenizer.decode(output_non_padded[0]   ,       skip_special_tokens=UpperCamelCase_							)
							UpperCamelCase__       :Union[str, Any]							 =			tokenizer.decode(output_padded[0]   ,       skip_special_tokens=UpperCamelCase_							)

							UpperCamelCase__       :Union[str, Any]							 =			[
							    '''Hello, my dog is a little bit of a dork.\nI\'m a little bit''',
							    '''Today, I was in the middle of a conversation with a friend about the''',
							]
							self.assertListEqual(UpperCamelCase_   ,       UpperCamelCase_							)
							self.assertListEqual(UpperCamelCase_   ,       [non_padded_sentence, padded_sentence]							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Dict							 =			'''facebook/opt-350m'''

							UpperCamelCase__       :Optional[Any]							 =			[
							    '''Today is a beautiful day and I want to''',
							    '''In the city of San Francisco, the city''',
							    '''Paris is the capital of France and the capital''',
							    '''Computers and mobile phones have taken over the''',
							]

							UpperCamelCase__       :List[Any]							 =			[]
							UpperCamelCase__       :List[str]							 =			GPTaTokenizer.from_pretrained(UpperCamelCase_							)
							UpperCamelCase__       :int							 =			TFOPTForCausalLM.from_pretrained(UpperCamelCase_							)

							for prompt in self.prompts:
												UpperCamelCase__       :Union[str, Any]							 =			tokenizer(UpperCamelCase_   ,       return_tensors='''tf'''							).input_ids

												UpperCamelCase__       :List[str]							 =			model.generate(UpperCamelCase_   ,       max_length=10							)

												UpperCamelCase__       :Dict							 =			tokenizer.batch_decode(UpperCamelCase_   ,       skip_special_tokens=UpperCamelCase_							)
												predicted_outputs += generated_string

							self.assertListEqual(UpperCamelCase_   ,       UpperCamelCase_							)



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