

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

__lowerCamelCase										=		get_logger()

__lowerCamelCase										=		None
class        _lowercase					(     TensorFormatter[Mapping, '''jax.Array''', Mapping]     ):







     def __init__(      self ,						UpperCamelCase_=None ,						UpperCamelCase_=None ,						**UpperCamelCase_       ):
      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`.'''       )
      __magic_name__     =      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:
       __magic_name__     =      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]       )}.'''       )
       __magic_name__     =      str(jax.devices()[0]       )
      __magic_name__     =      jnp_array_kwargs







     @staticmethod
     def 	lowerCAmelCase__    (      ):
      import jax

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







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      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 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      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()

      __magic_name__     =      {}

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

       if isinstance(SCREAMING_SNAKE_CASE_ ,						PIL.Image.Image       ):
        __magic_name__     =      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:
       __magic_name__     =      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 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      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       ):
       __magic_name__     =      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 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      return map_nested(self._recursive_tensorize ,						SCREAMING_SNAKE_CASE_ ,						map_list=SCREAMING_SNAKE_CASE_       )







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      __magic_name__     =      self.numpy_arrow_extractor().extract_row(SCREAMING_SNAKE_CASE_       )
      __magic_name__     =      self.python_features_decoder.decode_row(SCREAMING_SNAKE_CASE_       )
      return self.recursive_tensorize(SCREAMING_SNAKE_CASE_       )







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      __magic_name__     =      self.numpy_arrow_extractor().extract_column(SCREAMING_SNAKE_CASE_       )
      __magic_name__     =      self.python_features_decoder.decode_column(SCREAMING_SNAKE_CASE_ ,						pa_table.column_names[0]       )
      __magic_name__     =      self.recursive_tensorize(SCREAMING_SNAKE_CASE_       )
      __magic_name__     =      self._consolidate(SCREAMING_SNAKE_CASE_       )
      return column







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      __magic_name__     =      self.numpy_arrow_extractor().extract_batch(SCREAMING_SNAKE_CASE_       )
      __magic_name__     =      self.python_features_decoder.decode_batch(SCREAMING_SNAKE_CASE_       )
      __magic_name__     =      self.recursive_tensorize(SCREAMING_SNAKE_CASE_       )
      for column_name in batch:
       __magic_name__     =      self._consolidate(batch[column_name]       )
      return batch








"""simple docstring"""






from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase										=		logging.get_logger(__name__)

__lowerCamelCase										=		{
    "google/fnet-base": "https://huggingface.co/google/fnet-base/resolve/main/config.json",
    "google/fnet-large": "https://huggingface.co/google/fnet-large/resolve/main/config.json"
    # See all FNet models at https://huggingface.co/models?filter=fnet
}
class        _lowercase					(     __UpperCAmelCase     ):
     _lowerCamelCase     =       '''fnet'''







     def __init__(      self ,						UpperCamelCase_=3_2000 ,						UpperCamelCase_=768 ,						UpperCamelCase_=12 ,						UpperCamelCase_=3072 ,						UpperCamelCase_="gelu_new" ,						UpperCamelCase_=0.1 ,						UpperCamelCase_=512 ,						UpperCamelCase_=4 ,						UpperCamelCase_=0.0_2 ,						UpperCamelCase_=1E-1_2 ,						UpperCamelCase_=False ,						UpperCamelCase_=512 ,						UpperCamelCase_=3 ,						UpperCamelCase_=1 ,						UpperCamelCase_=2 ,						**UpperCamelCase_ ,						):
      super().__init__(pad_token_id=UpperCamelCase_ ,						bos_token_id=UpperCamelCase_ ,						eos_token_id=UpperCamelCase_ ,						**UpperCamelCase_       )

      __magic_name__     =      vocab_size
      __magic_name__     =      max_position_embeddings
      __magic_name__     =      hidden_size
      __magic_name__     =      num_hidden_layers
      __magic_name__     =      intermediate_size
      __magic_name__     =      hidden_act
      __magic_name__     =      hidden_dropout_prob
      __magic_name__     =      initializer_range
      __magic_name__     =      type_vocab_size
      __magic_name__     =      layer_norm_eps
      __magic_name__     =      use_tpu_fourier_optimizations
      __magic_name__     =      tpu_short_seq_length

