


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

__lowercase    =				get_logger()

__lowercase    =				None






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







    """simple docstring"""




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






    @staticmethod
    def      __snake_case							( ):
      import jax

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






    def      __snake_case							( self    :   List[Any]    ,			__UpperCAmelCase    :   Union[str, 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    :   Dict    ,			__UpperCAmelCase    :   Optional[int]):
      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()

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

        if isinstance(__UpperCAmelCase    ,			PIL.Image.Image):
          a						:		List[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:
        a						:		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    :   Optional[int]    ,			__UpperCAmelCase    :   Optional[Any]):
      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):
        a						:		List[str]											=      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    :   Optional[int]    ,			__UpperCAmelCase    :   dict):
      return map_nested(self._recursive_tensorize    ,			__UpperCAmelCase    ,			map_list=__UpperCAmelCase)






    def      __snake_case							( self    :   str    ,			__UpperCAmelCase    :   pa.Table):
      a						:		Tuple											=      self.numpy_arrow_extractor().extract_row(__UpperCAmelCase)
      a						:		Any											=      self.python_features_decoder.decode_row(__UpperCAmelCase)
      return self.recursive_tensorize(__UpperCAmelCase)






    def      __snake_case							( self    :   int    ,			__UpperCAmelCase    :   pa.Table):
      a						:		str											=      self.numpy_arrow_extractor().extract_column(__UpperCAmelCase)
      a						:		Any											=      self.python_features_decoder.decode_column(__UpperCAmelCase    ,			pa_table.column_names[0])
      a						:		Dict											=      self.recursive_tensorize(__UpperCAmelCase)
      a						:		Optional[int]											=      self._consolidate(__UpperCAmelCase)
      return column






    def      __snake_case							( self    :   Union[str, Any]    ,			__UpperCAmelCase    :   pa.Table):
      a						:		Tuple											=      self.numpy_arrow_extractor().extract_batch(__UpperCAmelCase)
      a						:		Union[str, Any]											=      self.python_features_decoder.decode_batch(__UpperCAmelCase)
      a						:		List[Any]											=      self.recursive_tensorize(__UpperCAmelCase)
      for column_name in batch:
        a						:		List[Any]											=      self._consolidate(batch[column_name])
      return batch





"""simple docstring"""







__lowercase    =				frozenset(
    [
        """prompt""",
        """height""",
        """width""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
        """cross_attention_kwargs""",
    ]
)

__lowercase    =				frozenset(["""prompt""", """negative_prompt"""])

__lowercase    =				frozenset([])

__lowercase    =				frozenset(["""image"""])

__lowercase    =				frozenset(
    [
        """image""",
        """height""",
        """width""",
        """guidance_scale""",
    ]
)

__lowercase    =				frozenset(["""image"""])

__lowercase    =				frozenset(
    [
        """prompt""",
        """image""",
        """height""",
        """width""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
    ]
)

__lowercase    =				frozenset(["""prompt""", """image""", """negative_prompt"""])

__lowercase    =				frozenset(
    [
        # Text guided image variation with an image mask
        """prompt""",
        """image""",
        """mask_image""",
        """height""",
        """width""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
    ]
)

__lowercase    =				frozenset(["""prompt""", """image""", """mask_image""", """negative_prompt"""])

__lowercase    =				frozenset(
    [
        # image variation with an image mask
        """image""",
        """mask_image""",
        """height""",
        """width""",
        """guidance_scale""",
    ]
)

__lowercase    =				frozenset(["""image""", """mask_image"""])

__lowercase    =				frozenset(
    [
        """example_image""",
        """image""",
        """mask_image""",
        """height""",
        """width""",
        """guidance_scale""",
    ]
)

__lowercase    =				frozenset(["""example_image""", """image""", """mask_image"""])

__lowercase    =				frozenset(["""class_labels"""])

__lowercase    =				frozenset(["""class_labels"""])

__lowercase    =				frozenset(["""batch_size"""])

__lowercase    =				frozenset([])

__lowercase    =				frozenset(["""batch_size"""])

__lowercase    =				frozenset([])

__lowercase    =				frozenset(
    [
        """prompt""",
        """audio_length_in_s""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
        """cross_attention_kwargs""",
    ]
)

__lowercase    =				frozenset(["""prompt""", """negative_prompt"""])
__lowercase    =				frozenset(["""input_tokens"""])

__lowercase    =				frozenset(["""input_tokens"""])






