

# 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

__UpperCAmelCase    				=     get_logger()

__UpperCAmelCase    				=     None





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



 """simple docstring"""







 def __init__(		self :	str     ,						A_ :	int=None     ,						A_ :	List[Any]=None     ,						**A_ :	Any       ):






      '''simple docstring'''




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

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

 @staticmethod
 def    __magic_name__       (		):






      '''simple docstring'''




      import jax

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

 def    __magic_name__       (		self :	Union[str, Any]     ,						A_ :	Tuple       ):






      '''simple docstring'''




      import jax
      import jax.numpy as jnp

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

 def    __magic_name__       (		self :	str     ,						A_ :	str       ):






      '''simple docstring'''




      import jax
      import jax.numpy as jnp

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

      _lowerCAmelCase   :					List[Any]        =    {}

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

           if isinstance(__SCREAMING_SNAKE_CASE     ,						PIL.Image.Image       ):
                _lowerCAmelCase   :					List[str]        =    np.asarray(__SCREAMING_SNAKE_CASE       )

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

 def    __magic_name__       (		self :	Tuple     ,						A_ :	str       ):






      '''simple docstring'''




      import jax

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

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

 def    __magic_name__       (		self :	Tuple     ,						A_ :	dict       ):






      '''simple docstring'''




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

 def    __magic_name__       (		self :	str     ,						A_ :	pa.Table       ):






      '''simple docstring'''




      _lowerCAmelCase   :					List[str]        =    self.numpy_arrow_extractor().extract_row(__SCREAMING_SNAKE_CASE       )
      _lowerCAmelCase   :					Dict        =    self.python_features_decoder.decode_row(__SCREAMING_SNAKE_CASE       )
      return self.recursive_tensorize(__SCREAMING_SNAKE_CASE       )

 def    __magic_name__       (		self :	Tuple     ,						A_ :	pa.Table       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Dict        =    self.numpy_arrow_extractor().extract_column(__SCREAMING_SNAKE_CASE       )
      _lowerCAmelCase   :					Union[str, Any]        =    self.python_features_decoder.decode_column(__SCREAMING_SNAKE_CASE     ,						pa_table.column_names[0]       )
      _lowerCAmelCase   :					Any        =    self.recursive_tensorize(__SCREAMING_SNAKE_CASE       )
      _lowerCAmelCase   :					Dict        =    self._consolidate(__SCREAMING_SNAKE_CASE       )
      return column


 def    __magic_name__       (		self :	Union[str, Any]     ,						A_ :	pa.Table       ):






      '''simple docstring'''




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









from typing import TYPE_CHECKING

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


__UpperCAmelCase    				=     {'configuration_yolos': ['YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP', 'YolosConfig', 'YolosOnnxConfig']}

try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase    				=     ['YolosFeatureExtractor']
  __UpperCAmelCase    				=     ['YolosImageProcessor']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase    				=     [
      'YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST',
      'YolosForObjectDetection',
      'YolosModel',
      'YolosPreTrainedModel',
  ]


if TYPE_CHECKING:
  from .configuration_yolos import YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP, YolosConfig, YolosOnnxConfig

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .feature_extraction_yolos import YolosFeatureExtractor
    from .image_processing_yolos import YolosImageProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_yolos import (
        YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST,
        YolosForObjectDetection,
        YolosModel,
        YolosPreTrainedModel,
    )


else:
  import sys

  __UpperCAmelCase    				=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






