



"""simple docstring"""







from collections import OrderedDict
from typing import Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...feature_extraction_utils import FeatureExtractionMixin
from ...onnx import OnnxConfig
from ...onnx.utils import compute_effective_axis_dimension
from ...tokenization_utils_base import PreTrainedTokenizerBase
from ...utils import TensorType, logging


lowercase__		:				str          =	logging.get_logger(__name__)

lowercase__		:				Union[str, Any]          =	{
    """deepmind/language-perceiver""": """https://huggingface.co/deepmind/language-perceiver/resolve/main/config.json""",
    # See all Perceiver models at https://huggingface.co/models?filter=perceiver
}







class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     """perceiver"""



      def __init__(      self				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=2_5_6    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=1_2_8_0    ,			SCREAMING_SNAKE_CASE_				:      str=7_6_8    ,			SCREAMING_SNAKE_CASE_				:      List[str]=1    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=2_6    ,			SCREAMING_SNAKE_CASE_				:      Any=8    ,			SCREAMING_SNAKE_CASE_				:      Dict=8    ,			SCREAMING_SNAKE_CASE_				:      List[str]=None    ,			SCREAMING_SNAKE_CASE_				:      Dict=None    ,			SCREAMING_SNAKE_CASE_				:      int="kv"    ,			SCREAMING_SNAKE_CASE_				:      Dict=1    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=1    ,			SCREAMING_SNAKE_CASE_				:      Dict="gelu"    ,			SCREAMING_SNAKE_CASE_				:      List[str]=0.1    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=0.02    ,			SCREAMING_SNAKE_CASE_				:      List[str]=1E-12    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=True    ,			SCREAMING_SNAKE_CASE_				:      str=2_6_2    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=2_0_4_8    ,			SCREAMING_SNAKE_CASE_				:      Any=5_6    ,			SCREAMING_SNAKE_CASE_				:      Tuple=[3_6_8, 4_9_6]    ,			SCREAMING_SNAKE_CASE_				:      List[str]=1_6    ,			SCREAMING_SNAKE_CASE_				:      Tuple=1_9_2_0    ,			SCREAMING_SNAKE_CASE_				:      Dict=1_6    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=[1, 1_6, 2_2_4, 2_2_4]    ,			**SCREAMING_SNAKE_CASE_				:      str    ,			):
        super().__init__(**SCREAMING_SNAKE_CASE_       )

        lowerCAmelCase_   :				str								= num_latents
        lowerCAmelCase_   :				Optional[Any]								= d_latents
        lowerCAmelCase_   :				List[Any]								= d_model
        lowerCAmelCase_   :				Optional[Any]								= num_blocks
        lowerCAmelCase_   :				int								= num_self_attends_per_block
        lowerCAmelCase_   :				Dict								= num_self_attention_heads
        lowerCAmelCase_   :				Tuple								= num_cross_attention_heads
        lowerCAmelCase_   :				int								= qk_channels
        lowerCAmelCase_   :				Optional[Any]								= v_channels
        lowerCAmelCase_   :				Any								= cross_attention_shape_for_attention
        lowerCAmelCase_   :				str								= self_attention_widening_factor
        lowerCAmelCase_   :				str								= cross_attention_widening_factor
        lowerCAmelCase_   :				Optional[int]								= hidden_act
        lowerCAmelCase_   :				Any								= attention_probs_dropout_prob
        lowerCAmelCase_   :				Optional[int]								= initializer_range
        lowerCAmelCase_   :				str								= layer_norm_eps
        lowerCAmelCase_   :				Tuple								= use_query_residual
        # masked language modeling attributes
        lowerCAmelCase_   :				int								= vocab_size
        lowerCAmelCase_   :				Optional[Any]								= max_position_embeddings
        # image classification attributes
        lowerCAmelCase_   :				Optional[int]								= image_size
        # flow attributes
        lowerCAmelCase_   :				List[Any]								= train_size
        # multimodal autoencoding attributes
        lowerCAmelCase_   :				Union[str, Any]								= num_frames
        lowerCAmelCase_   :				Any								= audio_samples_per_frame
        lowerCAmelCase_   :				int								= samples_per_patch
        lowerCAmelCase_   :				Any								= output_shape







class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      @property
      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        if self.task == "multiple-choice":
          lowerCAmelCase_   :				str								= {0: 'batch', 1: 'choice', 2: 'sequence'}
        else:
          lowerCAmelCase_   :				int								= {0: 'batch', 1: 'sequence'}
        return OrderedDict(
            [
                ('inputs', dynamic_axis),
                ('attention_mask', dynamic_axis),
            ]       )



      @property
      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]       ):
        return 1E-4





      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      int = -1    ,			SCREAMING_SNAKE_CASE_				:      bool = False    ,			SCREAMING_SNAKE_CASE_				:      Optional[TensorType] = None    ,			SCREAMING_SNAKE_CASE_				:      int = 3    ,			SCREAMING_SNAKE_CASE_				:      int = 4_0    ,			SCREAMING_SNAKE_CASE_				:      int = 4_0    ,			):
        # copied from `transformers.onnx.config.OnnxConfig` and slightly altered/simplified

        if isinstance(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ):
          # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
          lowerCAmelCase_   :				Tuple								= compute_effective_axis_dimension(
              SCREAMING_SNAKE_CASE_    ,			fixed_dimension=OnnxConfig.default_fixed_batch    ,			num_token_to_add=0       )
          # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
          lowerCAmelCase_   :				Dict								= preprocessor.num_special_tokens_to_add(SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				Tuple								= compute_effective_axis_dimension(
              SCREAMING_SNAKE_CASE_    ,			fixed_dimension=OnnxConfig.default_fixed_sequence    ,			num_token_to_add=SCREAMING_SNAKE_CASE_       )
          # Generate dummy inputs according to compute batch and sequence
          lowerCAmelCase_   :				Union[str, Any]								= [' '.join(['a']       ) * seq_length] * batch_size
          lowerCAmelCase_   :				List[Any]								= dict(preprocessor(SCREAMING_SNAKE_CASE_    ,			return_tensors=SCREAMING_SNAKE_CASE_       )       )
          lowerCAmelCase_   :				List[Any]								= inputs.pop('input_ids'       )
          return inputs
        elif isinstance(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ) and preprocessor.model_input_names[0] == "pixel_values":
          # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
          lowerCAmelCase_   :				Optional[Any]								= compute_effective_axis_dimension(SCREAMING_SNAKE_CASE_    ,			fixed_dimension=OnnxConfig.default_fixed_batch       )
          lowerCAmelCase_   :				str								= self._generate_dummy_images(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
          lowerCAmelCase_   :				str								= dict(preprocessor(images=SCREAMING_SNAKE_CASE_    ,			return_tensors=SCREAMING_SNAKE_CASE_       )       )
          lowerCAmelCase_   :				Optional[Any]								= inputs.pop('pixel_values'       )
          return inputs
        else:
          raise ValueError(
              'Unable to generate dummy inputs for the model. Please provide a tokenizer or a preprocessor.'       )










"""simple docstring"""







import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
     import tensorflow as tf
     from tensorflow.python.eager import context
     from tensorflow.python.framework import ops

     from transformers import GradientAccumulator, create_optimizer







@require_tf
class       UpperCamelCase__   (	unittest.TestCase  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      List[str]       ):
        self.assertEqual(len(SCREAMING_SNAKE_CASE_       )    ,			len(SCREAMING_SNAKE_CASE_       )       )
        for a, b in zip(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       ):
          self.assertAlmostEqual(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			delta=SCREAMING_SNAKE_CASE_       )



      def      SCREAMING_SNAKE_CASE__ (      self				:      int       ):
        lowerCAmelCase_   :				int								= GradientAccumulator()
        accumulator([tf.constant([1.0, 2.0]       )]       )
        accumulator([tf.constant([-2.0, 1.0]       )]       )
        accumulator([tf.constant([-1.0, 2.0]       )]       )
        with self.assertRaises(SCREAMING_SNAKE_CASE_       ):
          accumulator([tf.constant([1.0, 1.0]       ), tf.constant([2.0, 2.0]       )]       )
        self.assertEqual(accumulator.step    ,			3       )
        self.assertEqual(len(accumulator.gradients       )    ,			1       )
        self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()    ,			[-2.0, 5.0]    ,			tol=1E-2       )
        accumulator.reset()
        self.assertEqual(accumulator.step    ,			0       )
        self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()    ,			[0.0, 0.0]    ,			tol=1E-2       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        lowerCAmelCase_   :				Optional[int]								= None
        ops.enable_eager_execution_internal()
        lowerCAmelCase_   :				str								= tf.config.list_physical_devices('CPU'       )
        if len(SCREAMING_SNAKE_CASE_       ) == 1:
          tf.config.set_logical_device_configuration(
              physical_devices[0]    ,			[tf.config.LogicalDeviceConfiguration(), tf.config.LogicalDeviceConfiguration()]       )
        lowerCAmelCase_   :				Dict								= tf.config.list_logical_devices(device_type='CPU'       )
        lowerCAmelCase_   :				Optional[int]								= tf.distribute.MirroredStrategy(devices=devices[:2]       )

        with strategy.scope():
          lowerCAmelCase_   :				Union[str, Any]								= GradientAccumulator()
          lowerCAmelCase_   :				int								= tf.Variable([4.0, 3.0]       )
          lowerCAmelCase_						,lowerCAmelCase_   :				Optional[int]								= create_optimizer(5E-5    ,			1_0    ,			5       )
          lowerCAmelCase_   :				Union[str, Any]								= tf.Variable([0.0, 0.0]    ,			trainable=SCREAMING_SNAKE_CASE_       )

        def accumulate_on_replica(SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
          accumulator([gradient]       )

        def apply_on_replica():
          optimizer.apply_gradients(list(zip(accumulator.gradients    ,			[variable]       )       )       )

        @tf.function
        def accumulate(SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
          with strategy.scope():
            lowerCAmelCase_   :				Tuple								= strategy.experimental_local_results(SCREAMING_SNAKE_CASE_       )
            local_variables[0].assign(SCREAMING_SNAKE_CASE_       )
            local_variables[1].assign(SCREAMING_SNAKE_CASE_       )
            strategy.run(SCREAMING_SNAKE_CASE_    ,			args=(gradient_placeholder,)       )

        @tf.function
        def apply_grad():
          with strategy.scope():
            strategy.run(SCREAMING_SNAKE_CASE_       )

        def _check_local_values(SCREAMING_SNAKE_CASE_				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
          lowerCAmelCase_   :				List[Any]								= strategy.experimental_local_results(accumulator._gradients[0]       )
          self.assertListAlmostEqual(values[0].value()    ,			SCREAMING_SNAKE_CASE_    ,			tol=1E-2       )
          self.assertListAlmostEqual(values[1].value()    ,			SCREAMING_SNAKE_CASE_    ,			tol=1E-2       )

        accumulate([1.0, 2.0]    ,			[-1.0, 1.0]       )
        accumulate([3.0, -1.0]    ,			[-1.0, -1.0]       )
        accumulate([-2.0, 2.0]    ,			[3.0, -2.0]       )
        self.assertEqual(accumulator.step    ,			3       )
        _check_local_values([2.0, 3.0]    ,			[1.0, -2.0]       )
        apply_grad()
        self.assertListAlmostEqual(variable.value()    ,			[4.0, 3.0]    ,			tol=1E-2       )
        accumulator.reset()
        self.assertEqual(accumulator.step    ,			0       )
        _check_local_values([0.0, 0.0]    ,			[0.0, 0.0]       )





