





# 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	:					Optional[Any]			  =					get_logger()

__lowerCamelCase	:					Optional[dict]			  =					None



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





						def __init__(			self	,							A_=None	,							A_=None	,							**A_       ):
							'''simple docstring'''




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

							if isinstance(A_	,							A_       ):
								raise ValueError(
								    F"""Expected {device} to be a `str` not {type(A_       )}, 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     : int	     =       device if isinstance(A_	,							A_       ) 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     : Union[str, 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     : Any	     =       str(jax.devices()[0]       )
							UpperCamelCase     : Optional[Any]	     =       jnp_array_kwargs





						@staticmethod
						def 	__UpperCamelCase(			):
							'''simple docstring'''




							import jax

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





						def 	__UpperCamelCase(			self	,							A_       ):
							'''simple docstring'''




							import jax
							import jax.numpy as jnp

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





						def 	__UpperCamelCase(			self	,							A_       ):
							'''simple docstring'''




							import jax
							import jax.numpy as jnp

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

							UpperCamelCase     : Any	     =       {}

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

								if isinstance(A_	,							PIL.Image.Image       ):
									UpperCamelCase     : str	     =       np.asarray(A_       )

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





						def 	__UpperCamelCase(			self	,							A_       ):
							'''simple docstring'''




							import jax

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

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





						def 	__UpperCamelCase(			self	,							A_       ):
							'''simple docstring'''




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





						def 	__UpperCamelCase(			self	,							A_       ):
							'''simple docstring'''




							UpperCamelCase     : Optional[Any]	     =       self.numpy_arrow_extractor().extract_row(A_       )
							UpperCamelCase     : Tuple	     =       self.python_features_decoder.decode_row(A_       )
							return self.recursive_tensorize(A_       )





						def 	__UpperCamelCase(			self	,							A_       ):
							'''simple docstring'''




							UpperCamelCase     : List[Any]	     =       self.numpy_arrow_extractor().extract_column(A_       )
							UpperCamelCase     : List[str]	     =       self.python_features_decoder.decode_column(A_	,							pa_table.column_names[0]       )
							UpperCamelCase     : Any	     =       self.recursive_tensorize(A_       )
							UpperCamelCase     : int	     =       self._consolidate(A_       )
							return column





						def 	__UpperCamelCase(			self	,							A_       ):
							'''simple docstring'''




							UpperCamelCase     : Optional[Any]	     =       self.numpy_arrow_extractor().extract_batch(A_       )
							UpperCamelCase     : Optional[Any]	     =       self.python_features_decoder.decode_batch(A_       )
							UpperCamelCase     : Dict	     =       self.recursive_tensorize(A_       )
							for column_name in batch:
								UpperCamelCase     : Optional[Any]	     =       self._consolidate(batch[column_name]       )
							return batch







from __future__ import annotations



def 		A_    (					_lowerCAmelCase							,       _lowerCAmelCase					)						->			list[str]:
 if partitions <= 0:
  raise ValueError("partitions must be a positive number!"					)
 if partitions > number_of_bytes:
  raise ValueError("partitions can not > number_of_bytes!"					)
 UpperCamelCase     : str	     =       number_of_bytes // partitions
 UpperCamelCase     : List[Any]	     =       []
 for i in range(_lowerCAmelCase					):
  UpperCamelCase     : Optional[Any]	     =       i * bytes_per_partition + 1
  UpperCamelCase     : Any	     =       (
      number_of_bytes if i == partitions - 1 else (i + 1) * bytes_per_partition
  )
  allocation_list.append(F"""{start_bytes}-{end_bytes}"""					)
 return allocation_list


if __name__ == "__main__":
    import doctest

    doctest.testmod()
