



import importlib
import math
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, Tuple, Union

import flax
import jax.numpy as jnp

from ..utils import BaseOutput


_lowerCAmelCase	:   Tuple    =   "scheduler_config.json"


class 	__snake_case   ( SCREAMING_SNAKE_CASE ):
		SCREAMING_SNAKE_CASE__								=	1
		SCREAMING_SNAKE_CASE__								=	2
		SCREAMING_SNAKE_CASE__								=	3
		SCREAMING_SNAKE_CASE__								=	4
		SCREAMING_SNAKE_CASE__								=	5



@dataclass
class 	__snake_case   ( SCREAMING_SNAKE_CASE ):
		SCREAMING_SNAKE_CASE__								=	42



class 	__snake_case   :
		SCREAMING_SNAKE_CASE__								=	SCHEDULER_CONFIG_NAME
		SCREAMING_SNAKE_CASE__								=	['dtype']
		SCREAMING_SNAKE_CASE__								=	[]
		SCREAMING_SNAKE_CASE__								=	True
		@classmethod
		def 						SCREAMING_SNAKE_CASE_					(							cls    ,a_ = None    ,a_ = None    ,a_=False    ,**a_    ,):







						"""simple docstring"""
						lowerCAmelCase__  ,   lowerCAmelCase__		  =  cls.load_config(
						    pretrained_model_name_or_path=a_    ,subfolder=a_    ,return_unused_kwargs=a_    ,**a_    ,)
						lowerCAmelCase__  ,   lowerCAmelCase__		  =  cls.from_config(a_    ,return_unused_kwargs=a_    ,**a_						)

						if hasattr(a_    ,'create_state'						) and getattr(a_    ,'has_state'    ,a_						):
										lowerCAmelCase__		  =  scheduler.create_state()

						if return_unused_kwargs:
										return scheduler, state, unused_kwargs

						return scheduler, state
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_ = False    ,**a_						):







						"""simple docstring"""
						self.save_config(save_directory=a_    ,push_to_hub=a_    ,**a_						)
		@property
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						return self._get_compatibles()



		@classmethod
		def 						SCREAMING_SNAKE_CASE_					(							cls						):







						"""simple docstring"""
						lowerCAmelCase__		  =  list(set([cls.__name__] + cls._compatibles						)						)
						lowerCAmelCase__		  =  importlib.import_module(__name__.split('.'						)[0]						)
						lowerCAmelCase__		  =  [
						    getattr(a_    ,a_						) for c in compatible_classes_str if hasattr(a_    ,a_						)
						]
						return compatible_classes







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__					)	->     jnp.ndarray:



				"""simple docstring"""





				assert len(snake_case__					) >= x.ndim
				return jnp.broadcast_to(x.reshape(x.shape + (1,) * (len(snake_case__					) - x.ndim)					)    ,			snake_case__					)







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__=0.999    ,			snake_case__=jnp.floataa					)	->     jnp.ndarray:



				"""simple docstring"""





				def alpha_bar(snake_case__					):
								return math.cos((time_step + 0.008) / 1.008 * math.pi / 2					) ** 2

				lowerCAmelCase__		  =  []
				for i in range(snake_case__					):
								lowerCAmelCase__		  =  i / num_diffusion_timesteps
								lowerCAmelCase__		  =  (i + 1) / num_diffusion_timesteps
								betas.append(min(1 - alpha_bar(snake_case__					) / alpha_bar(snake_case__					)    ,			snake_case__					)					)
				return jnp.array(snake_case__    ,			dtype=snake_case__					)


@flax.struct.dataclass
class 	__snake_case   :
		SCREAMING_SNAKE_CASE__								=	42
		SCREAMING_SNAKE_CASE__								=	42
		SCREAMING_SNAKE_CASE__								=	42
		@classmethod
		def 						SCREAMING_SNAKE_CASE_					(							cls    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  scheduler.config

						if config.trained_betas is not None:
										lowerCAmelCase__		  =  jnp.asarray(config.trained_betas    ,dtype=scheduler.dtype						)
						elif config.beta_schedule == "linear":
										lowerCAmelCase__		  =  jnp.linspace(config.beta_start    ,config.beta_end    ,config.num_train_timesteps    ,dtype=scheduler.dtype						)
						elif config.beta_schedule == "scaled_linear":
										# this schedule is very specific to the latent diffusion model.
										lowerCAmelCase__		  =  (
										    jnp.linspace(
										        config.beta_start**0.5    ,config.beta_end**0.5    ,config.num_train_timesteps    ,dtype=scheduler.dtype						)
										    ** 2
										)
						elif config.beta_schedule == "squaredcos_cap_v2":
										# Glide cosine schedule
										lowerCAmelCase__		  =  betas_for_alpha_bar(config.num_train_timesteps    ,dtype=scheduler.dtype						)
						else:
										raise NotImplementedError(
										    f'beta_schedule {config.beta_schedule} is not implemented for scheduler {scheduler.__class__.__name__}'						)

						lowerCAmelCase__		  =  1.0 - betas

						lowerCAmelCase__		  =  jnp.cumprod(a_    ,axis=0						)

						return cls(
						    alphas=a_    ,betas=a_    ,alphas_cumprod=a_    ,)







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__					)	->     List[Any]:



				"""simple docstring"""





				lowerCAmelCase__		  =  state.alphas_cumprod

				lowerCAmelCase__		  =  alphas_cumprod[timesteps] ** 0.5
				lowerCAmelCase__		  =  sqrt_alpha_prod.flatten()
				lowerCAmelCase__		  =  broadcast_to_shape_from_left(snake_case__    ,			original_samples.shape					)

				lowerCAmelCase__		  =  (1 - alphas_cumprod[timesteps]) ** 0.5
				lowerCAmelCase__		  =  sqrt_one_minus_alpha_prod.flatten()
				lowerCAmelCase__		  =  broadcast_to_shape_from_left(snake_case__    ,			original_samples.shape					)

				return sqrt_alpha_prod, sqrt_one_minus_alpha_prod







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__					)	->     Optional[int]:



				"""simple docstring"""





				lowerCAmelCase__  ,   lowerCAmelCase__		  =  get_sqrt_alpha_prod(snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__					)
				lowerCAmelCase__		  =  sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
				return noisy_samples







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__					)	->     Optional[Any]:



				"""simple docstring"""





				lowerCAmelCase__  ,   lowerCAmelCase__		  =  get_sqrt_alpha_prod(snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__					)
				lowerCAmelCase__		  =  sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
				return velocity





import inspect
import unittest
from typing import List

import numpy as np

from transformers import EfficientFormerConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
	import tensorflow as tf

	from transformers import (
	    TFEfficientFormerForImageClassification,
	    TFEfficientFormerForImageClassificationWithTeacher,
	    TFEfficientFormerModel,
	)
	from transformers.models.efficientformer.modeling_tf_efficientformer import (
	    TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
	)


if is_vision_available():
	from PIL import Image

	from transformers import EfficientFormerImageProcessor


class 	__snake_case   :
		def __init__(							self    ,a_    ,a_ = 13    ,a_ = 64    ,a_ = 2    ,a_ = 3    ,a_ = 3    ,a_ = True    ,a_ = True    ,a_ = 128    ,a_=[16, 32, 64, 128]    ,a_ = 7    ,a_ = 4    ,a_ = 37    ,a_ = "gelu"    ,a_ = 0.1    ,a_ = 0.1    ,a_ = 10    ,a_ = 0.02    ,a_ = 2    ,a_ = 1    ,a_ = 128    ,a_ = [2, 2, 2, 2]    ,a_ = 2    ,a_ = 2    ,):







						"""simple docstring"""
						lowerCAmelCase__		  =  parent
						lowerCAmelCase__		  =  batch_size
						lowerCAmelCase__		  =  image_size
						lowerCAmelCase__		  =  patch_size
						lowerCAmelCase__		  =  num_channels
						lowerCAmelCase__		  =  is_training
						lowerCAmelCase__		  =  use_labels
						lowerCAmelCase__		  =  hidden_size
						lowerCAmelCase__		  =  num_hidden_layers
						lowerCAmelCase__		  =  num_attention_heads
						lowerCAmelCase__		  =  intermediate_size
						lowerCAmelCase__		  =  hidden_act
						lowerCAmelCase__		  =  hidden_dropout_prob
						lowerCAmelCase__		  =  attention_probs_dropout_prob
						lowerCAmelCase__		  =  type_sequence_label_size
						lowerCAmelCase__		  =  initializer_range
						lowerCAmelCase__		  =  encoder_stride
						lowerCAmelCase__		  =  num_attention_outputs
						lowerCAmelCase__		  =  embed_dim
						lowerCAmelCase__		  =  embed_dim + 1
						lowerCAmelCase__		  =  resolution
						lowerCAmelCase__		  =  depths
						lowerCAmelCase__		  =  hidden_sizes
						lowerCAmelCase__		  =  dim
						lowerCAmelCase__		  =  mlp_expansion_ratio
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

						lowerCAmelCase__		  =  None
						if self.use_labels:
										lowerCAmelCase__		  =  ids_tensor([self.batch_size]    ,self.type_sequence_label_size						)

						lowerCAmelCase__		  =  self.get_config()

						return config, pixel_values, labels
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						return EfficientFormerConfig(
						    image_size=self.image_size    ,patch_size=self.patch_size    ,num_channels=self.num_channels    ,hidden_size=self.hidden_size    ,num_hidden_layers=self.num_hidden_layers    ,num_attention_heads=self.num_attention_heads    ,intermediate_size=self.intermediate_size    ,hidden_act=self.hidden_act    ,hidden_dropout_prob=self.hidden_dropout_prob    ,attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,is_decoder=a_    ,initializer_range=self.initializer_range    ,encoder_stride=self.encoder_stride    ,resolution=self.resolution    ,depths=self.depths    ,hidden_sizes=self.hidden_sizes    ,dim=self.dim    ,mlp_expansion_ratio=self.mlp_expansion_ratio    ,)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  TFEfficientFormerModel(config=a_						)
						lowerCAmelCase__		  =  model(a_    ,training=a_						)
						self.parent.assertEqual(result.last_hidden_state.shape    ,(self.batch_size, self.seq_length, self.hidden_size)						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.type_sequence_label_size
						lowerCAmelCase__		  =  TFEfficientFormerForImageClassification(a_						)
						lowerCAmelCase__		  =  model(a_    ,labels=a_    ,training=a_						)
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.type_sequence_label_size)						)

						# test greyscale images
						lowerCAmelCase__		  =  1
						lowerCAmelCase__		  =  TFEfficientFormerForImageClassification(a_						)

						lowerCAmelCase__		  =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
						lowerCAmelCase__		  =  model(a_    ,labels=a_						)
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.type_sequence_label_size)						)



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.prepare_config_and_inputs()
						lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__		  =  config_and_inputs
						lowerCAmelCase__		  =  {'pixel_values': pixel_values}
						return config, inputs_dict



@require_tf
class 	__snake_case   ( SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, unittest.TestCase ):
		SCREAMING_SNAKE_CASE__								=	(
		    (
		        TFEfficientFormerModel,
		        TFEfficientFormerForImageClassificationWithTeacher,
		        TFEfficientFormerForImageClassification,
		    )
		    if is_tf_available()
		    else ()
		)
		SCREAMING_SNAKE_CASE__								=	(
		    {
		        'feature-extraction': TFEfficientFormerModel,
		        'image-classification': (
		            TFEfficientFormerForImageClassification,
		            TFEfficientFormerForImageClassificationWithTeacher,
		        ),
		    }
		    if is_tf_available()
		    else {}
		)

		SCREAMING_SNAKE_CASE__								=	False

		SCREAMING_SNAKE_CASE__								=	False
		SCREAMING_SNAKE_CASE__								=	False
		SCREAMING_SNAKE_CASE__								=	False
		SCREAMING_SNAKE_CASE__								=	False
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  TFEfficientFormerModelTester(self						)
						lowerCAmelCase__		  =  ConfigTester(
						    self    ,config_class=a_    ,has_text_modality=a_    ,hidden_size=37						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						self.config_tester.run_common_tests()
		@unittest.skip(reason='EfficientFormer does not use inputs_embeds'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass
		@unittest.skip(reason='EfficientFormer does not support input and output embeddings'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__  ,   lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										lowerCAmelCase__		  =  model_class(a_						)
										lowerCAmelCase__		  =  inspect.signature(model.call						)
										# signature.parameters is an OrderedDict => so arg_names order is deterministic
										lowerCAmelCase__		  =  [*signature.parameters.keys()]

										lowerCAmelCase__		  =  ['pixel_values']
										self.assertListEqual(arg_names[:1]    ,a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						def check_hidden_states_output(a_    ,a_    ,a_						):
										lowerCAmelCase__		  =  model_class(a_						)

										lowerCAmelCase__		  =  model(**self._prepare_for_class(a_    ,a_						)    ,training=a_						)
										lowerCAmelCase__		  =  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

										lowerCAmelCase__		  =  getattr(
										    self.model_tester    ,'expected_num_hidden_layers'    ,self.model_tester.num_hidden_layers + 1						)
										self.assertEqual(len(a_						)    ,a_						)

										if hasattr(self.model_tester    ,'encoder_seq_length'						):
														lowerCAmelCase__		  =  self.model_tester.encoder_seq_length
														if hasattr(self.model_tester    ,'chunk_length'						) and self.model_tester.chunk_length > 1:
																		lowerCAmelCase__		  =  seq_length * self.model_tester.chunk_length
										else:
														lowerCAmelCase__		  =  self.model_tester.seq_length

										self.assertListEqual(
										    list(hidden_states[-1].shape[-2:]						)    ,[seq_length, self.model_tester.hidden_size]    ,)

										if config.is_encoder_decoder:
														lowerCAmelCase__		  =  outputs.decoder_hidden_states

														self.asseretIsInstance(a_    ,(list, tuple)						)
														self.assertEqual(len(a_						)    ,a_						)
														lowerCAmelCase__		  =  getattr(self.model_tester    ,'seq_length'    ,a_						)
														lowerCAmelCase__		  =  getattr(self.model_tester    ,'decoder_seq_length'    ,a_						)

														self.assertListEqual(
														    list(hidden_states[-1].shape[-2:]						)    ,[decoder_seq_length, self.model_tester.hidden_size]    ,)

						lowerCAmelCase__  ,   lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										lowerCAmelCase__		  =  True
										check_hidden_states_output(a_    ,a_    ,a_						)

										# check that output_hidden_states also work using config
										del inputs_dict["output_hidden_states"]
										lowerCAmelCase__		  =  True

										check_hidden_states_output(a_    ,a_    ,a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_=False						):







						"""simple docstring"""
						lowerCAmelCase__		  =  super()._prepare_for_class(a_    ,a_    ,return_labels=a_						)

						if return_labels:
										if model_class.__name__ == "TFEfficientFormerForImageClassificationWithTeacher":
														del inputs_dict["labels"]

						return inputs_dict
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*a_						)
		@unittest.skip(reason='EfficientFormer does not implement masked image modeling yet'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_image_modeling(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*a_						)
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						for model_name in TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase__		  =  TFEfficientFormerModel.from_pretrained(a_						)
										self.assertIsNotNone(a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__  ,   lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs_for_common()
						lowerCAmelCase__		  =  True

						lowerCAmelCase__		  =  getattr(self.model_tester    ,'seq_length'    ,a_						)
						lowerCAmelCase__		  =  getattr(self.model_tester    ,'encoder_seq_length'    ,a_						)
						lowerCAmelCase__		  =  getattr(self.model_tester    ,'key_length'    ,a_						)
						lowerCAmelCase__		  =  getattr(self.model_tester    ,'chunk_length'    ,a_						)

						if chunk_length is not None and hasattr(self.model_tester    ,'num_hashes'						):
										lowerCAmelCase__		  =  encoder_seq_length * self.model_tester.num_hashes

						for model_class in self.all_model_classes:
										lowerCAmelCase__		  =  True
										lowerCAmelCase__		  =  False
										lowerCAmelCase__		  =  True
										lowerCAmelCase__		  =  model_class(a_						)

										lowerCAmelCase__		  =  model(**self._prepare_for_class(a_    ,a_						)    ,training=a_						)
										lowerCAmelCase__		  =  outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
										self.assertEqual(len(a_						)    ,self.model_tester.num_attention_outputs						)

										# check that output_attentions also work using config
										del inputs_dict["output_attentions"]
										lowerCAmelCase__		  =  True
										lowerCAmelCase__		  =  model_class(a_						)
										lowerCAmelCase__		  =  model(**self._prepare_for_class(a_    ,a_						)    ,training=a_						)

										lowerCAmelCase__		  =  outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
										self.assertEqual(len(a_						)    ,self.model_tester.num_attention_outputs						)

										if chunk_length is not None:
														self.assertListEqual(
														    list(attentions[0].shape[-4:]						)    ,[self.model_tester.num_attention_heads, encoder_seq_length, chunk_length, encoder_key_length]    ,)
										else:
														self.assertListEqual(
														    list(attentions[0].shape[-3:]						)    ,[self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length]    ,)



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						# We use a simplified version of this test for EfficientFormer because it requires training=False
						# and Keras refuses to let us force that during functional construction
						lowerCAmelCase__  ,   lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										# Prepare our model
										lowerCAmelCase__		  =  model_class(a_						)
										# These are maximally general inputs for the model, with multiple None dimensions
										# Hopefully this will catch any conditionals that fail for flexible shapes
										lowerCAmelCase__		  =  {
										    key: tf.keras.Input(shape=val.shape[1:]    ,dtype=val.dtype    ,name=a_						)
										    for key, val in model.input_signature.items()
										    if key in model.dummy_inputs
										}
										lowerCAmelCase__		  =  model(a_						)
										self.assertTrue(outputs_dict is not None						)







def     UpperCAmelCase_		(					)	->     List[str]:



				"""simple docstring"""





				lowerCAmelCase__		  =  Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'					)
				return image


@require_tf
@require_vision
class 	__snake_case   ( unittest.TestCase ):
		@cached_property
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						return (
						    EfficientFormerImageProcessor.from_pretrained('snap-research/efficientformer-l1-300'						)
						    if is_vision_available()
						    else None
						)
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  TFEfficientFormerForImageClassification.from_pretrained('snap-research/efficientformer-l1-300'						)
						lowerCAmelCase__		  =  self.default_image_processor
						lowerCAmelCase__		  =  prepare_img()
						lowerCAmelCase__		  =  image_processor(images=a_    ,return_tensors='tf'						)
						# forward pass
						lowerCAmelCase__		  =  model(**a_    ,training=a_						)
						# verify the logits
						lowerCAmelCase__		  =  tf.TensorShape((1, 1000)						)
						self.assertEqual(outputs.logits.shape    ,a_						)
						lowerCAmelCase__		  =  tf.constant([-0.0555, 0.4825, -0.0852]						)
						self.assertTrue(np.allclose(outputs.logits[0, :3]    ,a_    ,atol=1e-4						)						)



		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  TFEfficientFormerForImageClassificationWithTeacher.from_pretrained(
						    'snap-research/efficientformer-l1-300'						)
						lowerCAmelCase__		  =  self.default_image_processor
						lowerCAmelCase__		  =  prepare_img()
						lowerCAmelCase__		  =  image_processor(images=a_    ,return_tensors='tf'						)
						# forward pass
						lowerCAmelCase__		  =  model(**a_    ,training=a_						)
						# verify the logits
						lowerCAmelCase__		  =  tf.TensorShape((1, 1000)						)
						self.assertEqual(outputs.logits.shape    ,a_						)
						lowerCAmelCase__		  =  tf.constant([-0.1312, 0.4353, -1.0499]						)
						self.assertTrue(np.allclose(outputs.logits[0, :3]    ,a_    ,atol=1e-4						)						)
