


# 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

SCREAMING_SNAKE_CASE__     					=			get_logger()

SCREAMING_SNAKE_CASE__     					=			None



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






		def __init__(						self				:     int   ,  UpperCAmelCase				:     Any=None   ,  UpperCAmelCase				:     List[Any]=None   ,  **UpperCAmelCase				:     Dict):
				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`.")
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						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:
						SCREAMING_SNAKE_CASE_				:str				     =						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])}.")
						SCREAMING_SNAKE_CASE_				:Any				     =						str(jax.devices()[0])
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						jnp_array_kwargs






		@staticmethod
		def _snake_case		(						):
				import jax

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






		def _snake_case		(						self				:     str   ,  UpperCAmelCase				:     List[Any]):
				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 _snake_case		(						self				:     Optional[int]   ,  UpperCAmelCase				:     Tuple):
				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()

				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						{}

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

						if isinstance(UpperCAmelCase   ,  PIL.Image.Image):
								SCREAMING_SNAKE_CASE_				:Any				     =						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:
						SCREAMING_SNAKE_CASE_				:int				     =						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 _snake_case		(						self				:     str   ,  UpperCAmelCase				:     str):
				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):
						SCREAMING_SNAKE_CASE_				:Optional[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 _snake_case		(						self				:     Dict   ,  UpperCAmelCase				:     dict):
				return map_nested(self._recursive_tensorize   ,  UpperCAmelCase   ,  map_list=UpperCAmelCase)






		def _snake_case		(						self				:     Union[str, Any]   ,  UpperCAmelCase				:     pa.Table):
				SCREAMING_SNAKE_CASE_				:Optional[Any]				     =						self.numpy_arrow_extractor().extract_row(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:int				     =						self.python_features_decoder.decode_row(UpperCAmelCase)
				return self.recursive_tensorize(UpperCAmelCase)






		def _snake_case		(						self				:     Any   ,  UpperCAmelCase				:     pa.Table):
				SCREAMING_SNAKE_CASE_				:str				     =						self.numpy_arrow_extractor().extract_column(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						self.python_features_decoder.decode_column(UpperCAmelCase   ,  pa_table.column_names[0])
				SCREAMING_SNAKE_CASE_				:int				     =						self.recursive_tensorize(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						self._consolidate(UpperCAmelCase)
				return column







		def _snake_case		(						self				:     Tuple   ,  UpperCAmelCase				:     pa.Table):
				SCREAMING_SNAKE_CASE_				:List[Any]				     =						self.numpy_arrow_extractor().extract_batch(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Any				     =						self.python_features_decoder.decode_batch(UpperCAmelCase)
				SCREAMING_SNAKE_CASE_				:Any				     =						self.recursive_tensorize(UpperCAmelCase)
				for column_name in batch:
						SCREAMING_SNAKE_CASE_				:List[str]				     =						self._consolidate(batch[column_name])
				return batch










import unittest

from transformers.utils.backbone_utils import (
    BackboneMixin,
    get_aligned_output_features_output_indices,
    verify_out_features_out_indices,
)



class 						_UpperCAmelCase      (     unittest.TestCase							):






		def _snake_case		(						self				:     Union[str, Any]):
				SCREAMING_SNAKE_CASE_				:Optional[int]				     =						["a", "b", "c"]

				# Defaults to last layer if both are None
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Dict				     =						get_aligned_output_features_output_indices(UpperCAmelCase   ,  UpperCAmelCase   ,  UpperCAmelCase)
				self.assertEqual(UpperCAmelCase   ,  ["c"])
				self.assertEqual(UpperCAmelCase   ,  [2])

				# Out indices set to match out features
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Tuple				     =						get_aligned_output_features_output_indices(["a", "c"]   ,  UpperCAmelCase   ,  UpperCAmelCase)
				self.assertEqual(UpperCAmelCase   ,  ["a", "c"])
				self.assertEqual(UpperCAmelCase   ,  [0, 2])

				# Out features set to match out indices
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:int				     =						get_aligned_output_features_output_indices(UpperCAmelCase   ,  [0, 2]   ,  UpperCAmelCase)
				self.assertEqual(UpperCAmelCase   ,  ["a", "c"])
				self.assertEqual(UpperCAmelCase   ,  [0, 2])

				# Out features selected from negative indices
				SCREAMING_SNAKE_CASE_     ,				SCREAMING_SNAKE_CASE_				:Any				     =						get_aligned_output_features_output_indices(UpperCAmelCase   ,  [-3, -1]   ,  UpperCAmelCase)
				self.assertEqual(UpperCAmelCase   ,  ["a", "c"])
				self.assertEqual(UpperCAmelCase   ,  [-3, -1])






		def _snake_case		(						self				:     str):
				# Stage names must be set
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["a", "b"]   ,  (0, 1)   ,  UpperCAmelCase)

				# Out features must be a list
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(("a", "b")   ,  (0, 1)   ,  ["a", "b"])

				# Out features must be a subset of stage names
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["a", "b"]   ,  (0, 1)   ,  ["a"])

				# Out indices must be a list or tuple
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(UpperCAmelCase   ,  0   ,  ["a", "b"])

				# Out indices must be a subset of stage names
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(UpperCAmelCase   ,  (0, 1)   ,  ["a"])

				# Out features and out indices must be the same length
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["a", "b"]   ,  (0,)   ,  ["a", "b", "c"])

				# Out features should match out indices
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["a", "b"]   ,  (0, 2)   ,  ["a", "b", "c"])

				# Out features and out indices should be in order
				with self.assertRaises(UpperCAmelCase):
						verify_out_features_out_indices(["b", "a"]   ,  (0, 1)   ,  ["a", "b"])

				# Check passes with valid inputs
				verify_out_features_out_indices(["a", "b", "d"]   ,  (0, 1, -1)   ,  ["a", "b", "c", "d"])







		def _snake_case		(						self				:     Tuple):
				SCREAMING_SNAKE_CASE_				:int				     =						BackboneMixin()

				SCREAMING_SNAKE_CASE_				:List[str]				     =						["a", "b", "c"]
				SCREAMING_SNAKE_CASE_				:str				     =						["a", "c"]
				SCREAMING_SNAKE_CASE_				:Union[str, Any]				     =						[0, 2]

				# Check that the output features and indices are set correctly
				self.assertEqual(backbone.out_features   ,  ["a", "c"])
				self.assertEqual(backbone.out_indices   ,  [0, 2])

				# Check out features and indices are updated correctly
				SCREAMING_SNAKE_CASE_				:Tuple				     =						["a", "b"]
				self.assertEqual(backbone.out_features   ,  ["a", "b"])
				self.assertEqual(backbone.out_indices   ,  [0, 1])

				SCREAMING_SNAKE_CASE_				:List[str]				     =						[-3, -1]
				self.assertEqual(backbone.out_features   ,  ["a", "c"])
				self.assertEqual(backbone.out_indices   ,  [-3, -1])






