

'''simple docstring'''


import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class        lowercase    (	a_,       a_,       a_,       unittest.TestCase			):
							_lowerCamelCase		:   Any= StableDiffusionImageVariationPipeline
							_lowerCamelCase		:   List[Any]= IMAGE_VARIATION_PARAMS
							_lowerCamelCase		:   str= IMAGE_VARIATION_BATCH_PARAMS
							_lowerCamelCase		:   Tuple= frozenset([]			)
							# TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
							_lowerCamelCase		:   Union[str, Any]= frozenset([]			)


							def     _snake_case				(      self)				->					str:
								torch.manual_seed(0)
								UpperCAmelCase_		:   List[Any]													=       UNetaDConditionModel(
								    block_out_channels=(32, 64)     ,      layers_per_block=2     ,      sample_size=32     ,      in_channels=4     ,      out_channels=4     ,      down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')     ,      up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')     ,      cross_attention_dim=32     ,      )
								UpperCAmelCase_		:   Any													=       PNDMScheduler(skip_prk_steps=_snake_case)
								torch.manual_seed(0)
								UpperCAmelCase_		:   Optional[Any]													=       AutoencoderKL(
								    block_out_channels=[32, 64]     ,      in_channels=3     ,      out_channels=3     ,      down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']     ,      up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']     ,      latent_channels=4     ,      )
								torch.manual_seed(0)
								UpperCAmelCase_		:   Optional[int]													=       CLIPVisionConfig(
								    hidden_size=32     ,      projection_dim=32     ,      intermediate_size=37     ,      layer_norm_eps=1e-05     ,      num_attention_heads=4     ,      num_hidden_layers=5     ,      image_size=32     ,      patch_size=4     ,      )
								UpperCAmelCase_		:   int													=       CLIPVisionModelWithProjection(_snake_case)
								UpperCAmelCase_		:   List[Any]													=       CLIPImageProcessor(crop_size=32     ,      size=32)

								UpperCAmelCase_		:   List[Any]													=       {
								    'unet': unet,
								    'scheduler': scheduler,
								    'vae': vae,
								    'image_encoder': image_encoder,
								    'feature_extractor': feature_extractor,
								    'safety_checker': None,
								}
								return components


							def     _snake_case				(      self     ,      _snake_case     ,      _snake_case=0)				->					Any:
								UpperCAmelCase_		:   Any													=       floats_tensor((1, 3, 32, 32)     ,      rng=random.Random(_snake_case))
								UpperCAmelCase_		:   Dict													=       image.cpu().permute(0     ,      2     ,      3     ,      1)[0]
								UpperCAmelCase_		:   List[Any]													=       Image.fromarray(np.uinta(_snake_case)).convert('RGB').resize((32, 32))
								if str(_snake_case).startswith('mps'):
									UpperCAmelCase_		:   Any													=       torch.manual_seed(_snake_case)
								else:
									UpperCAmelCase_		:   int													=       torch.Generator(device=_snake_case).manual_seed(_snake_case)
								UpperCAmelCase_		:   Union[str, Any]													=       {
								    'image': image,
								    'generator': generator,
								    'num_inference_steps': 2,
								    'guidance_scale': 6.0,
								    'output_type': 'numpy',
								}
								return inputs


							def     _snake_case				(      self)				->					Tuple:
								UpperCAmelCase_		:   Any													=       'cpu'  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_		:   int													=       self.get_dummy_components()
								UpperCAmelCase_		:   List[Any]													=       StableDiffusionImageVariationPipeline(**_snake_case)
								UpperCAmelCase_		:   Optional[Any]													=       sd_pipe.to(_snake_case)
								sd_pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:   Union[str, Any]													=       self.get_dummy_inputs(_snake_case)
								UpperCAmelCase_		:   Optional[Any]													=       sd_pipe(**_snake_case).images
								UpperCAmelCase_		:   List[str]													=       image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)
								UpperCAmelCase_		:   int													=       np.array([0.5_239, 0.5_723, 0.4_796, 0.5_049, 0.5_550, 0.4_685, 0.5_329, 0.4_891, 0.4_921])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3


							def     _snake_case				(      self)				->					Union[str, Any]:
								UpperCAmelCase_		:   Dict													=       'cpu'  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase_		:   Dict													=       self.get_dummy_components()
								UpperCAmelCase_		:   List[Any]													=       StableDiffusionImageVariationPipeline(**_snake_case)
								UpperCAmelCase_		:   Union[str, Any]													=       sd_pipe.to(_snake_case)
								sd_pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:   List[Any]													=       self.get_dummy_inputs(_snake_case)
								UpperCAmelCase_		:   Optional[int]													=       2 * [inputs['image']]
								UpperCAmelCase_		:   List[str]													=       sd_pipe(**_snake_case)

								UpperCAmelCase_		:   Dict													=       output.images

								UpperCAmelCase_		:   str													=       image[-1, -3:, -3:, -1]

								assert image.shape == (2, 64, 64, 3)
								UpperCAmelCase_		:   Tuple													=       np.array([0.6_892, 0.5_637, 0.5_836, 0.5_771, 0.6_254, 0.6_409, 0.5_580, 0.5_569, 0.5_289])

								assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3





							def     _snake_case				(      self)				->					List[Any]:
								super().test_inference_batch_single_identical(expected_max_diff=3e-3)



@slow
@require_torch_gpu
class        lowercase    (	unittest.TestCase			):


							def     _snake_case				(      self)				->					Any:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()


							def     _snake_case				(      self     ,      _snake_case     ,      _snake_case="cpu"     ,      _snake_case=torch.floataa     ,      _snake_case=0)				->					List[Any]:
								UpperCAmelCase_		:   Optional[int]													=       torch.Generator(device=_snake_case).manual_seed(_snake_case)
								UpperCAmelCase_		:   Any													=       load_image(
								    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
								    '/stable_diffusion_imgvar/input_image_vermeer.png')
								UpperCAmelCase_		:   Optional[Any]													=       np.random.RandomState(_snake_case).standard_normal((1, 4, 64, 64))
								UpperCAmelCase_		:   str													=       torch.from_numpy(_snake_case).to(device=_snake_case     ,      dtype=_snake_case)
								UpperCAmelCase_		:   int													=       {
								    'image': init_image,
								    'latents': latents,
								    'generator': generator,
								    'num_inference_steps': 3,
								    'guidance_scale': 7.5,
								    'output_type': 'numpy',
								}
								return inputs


							def     _snake_case				(      self)				->					List[Any]:
								UpperCAmelCase_		:   str													=       StableDiffusionImageVariationPipeline.from_pretrained(
								    'lambdalabs/sd-image-variations-diffusers'     ,      safety_checker=_snake_case)
								UpperCAmelCase_		:   Tuple													=       sd_pipe.to(_snake_case)
								sd_pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:   str													=       self.get_inputs(_snake_case)
								UpperCAmelCase_		:   str													=       sd_pipe(**_snake_case).images
								UpperCAmelCase_		:   Union[str, Any]													=       image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)
								UpperCAmelCase_		:   Any													=       np.array([0.84_491, 0.90_789, 0.75_708, 0.78_734, 0.83_485, 0.70_099, 0.66_938, 0.68_727, 0.61_379])
								assert np.abs(image_slice - expected_slice).max() < 6e-3


							def     _snake_case				(      self)				->					List[Any]:
								UpperCAmelCase_		:   Any													=       0

								def callback_fn(_snake_case     ,      _snake_case     ,      _snake_case) -> None:
									UpperCAmelCase_		:   List[str]													=       True
									nonlocal number_of_steps
									number_of_steps += 1
									if step == 1:
										UpperCAmelCase_		:   Union[str, Any]													=       latents.detach().cpu().numpy()
										assert latents.shape == (1, 4, 64, 64)
										UpperCAmelCase_		:   str													=       latents[0, -3:, -3:, -1]
										UpperCAmelCase_		:   Union[str, Any]													=       np.array(
										    [-0.1_621, 0.2_837, -0.7_979, -0.1_221, -1.3_057, 0.7_681, -2.1_191, 0.0_464, 1.6_309])

										assert np.abs(latents_slice.flatten() - expected_slice).max() < 5e-2
									elif step == 2:
										UpperCAmelCase_		:   List[str]													=       latents.detach().cpu().numpy()
										assert latents.shape == (1, 4, 64, 64)
										UpperCAmelCase_		:   Any													=       latents[0, -3:, -3:, -1]
										UpperCAmelCase_		:   Optional[Any]													=       np.array([0.6_299, 1.7_500, 1.1_992, -2.1_582, -1.8_994, 0.7_334, -0.7_090, 1.0_137, 1.5_273])

										assert np.abs(latents_slice.flatten() - expected_slice).max() < 5e-2

								UpperCAmelCase_		:   Any													=       False

								UpperCAmelCase_		:   List[str]													=       StableDiffusionImageVariationPipeline.from_pretrained(
								    'fusing/sd-image-variations-diffusers'     ,      safety_checker=_snake_case     ,      torch_dtype=torch.floataa     ,      )
								pipe.to(_snake_case)
								pipe.set_progress_bar_config(disable=_snake_case)
								pipe.enable_attention_slicing()

								UpperCAmelCase_		:   Tuple													=       self.get_inputs(_snake_case     ,      dtype=torch.floataa)
								pipe(**_snake_case     ,      callback=_snake_case     ,      callback_steps=1)
								assert callback_fn.has_been_called
								assert number_of_steps == inputs["num_inference_steps"]





							def     _snake_case				(      self)				->					str:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

								UpperCAmelCase_		:   int													=       'fusing/sd-image-variations-diffusers'
								UpperCAmelCase_		:   Tuple													=       StableDiffusionImageVariationPipeline.from_pretrained(
								    _snake_case     ,      safety_checker=_snake_case     ,      torch_dtype=torch.floataa)
								UpperCAmelCase_		:   str													=       pipe.to(_snake_case)
								pipe.set_progress_bar_config(disable=_snake_case)
								pipe.enable_attention_slicing(1)
								pipe.enable_sequential_cpu_offload()

								UpperCAmelCase_		:   Tuple													=       self.get_inputs(_snake_case     ,      dtype=torch.floataa)
								UpperCAmelCase_		:   str													=       pipe(**_snake_case)

								UpperCAmelCase_		:   Optional[int]													=       torch.cuda.max_memory_allocated()
								# make sure that less than 2.6 GB is allocated
								assert mem_bytes < 2.6 * 10**9



@nightly
@require_torch_gpu
class        lowercase    (	unittest.TestCase			):


							def     _snake_case				(      self)				->					Optional[Any]:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()


							def     _snake_case				(      self     ,      _snake_case     ,      _snake_case="cpu"     ,      _snake_case=torch.floataa     ,      _snake_case=0)				->					str:
								UpperCAmelCase_		:   str													=       torch.Generator(device=_snake_case).manual_seed(_snake_case)
								UpperCAmelCase_		:   Any													=       load_image(
								    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
								    '/stable_diffusion_imgvar/input_image_vermeer.png')
								UpperCAmelCase_		:   Dict													=       np.random.RandomState(_snake_case).standard_normal((1, 4, 64, 64))
								UpperCAmelCase_		:   Any													=       torch.from_numpy(_snake_case).to(device=_snake_case     ,      dtype=_snake_case)
								UpperCAmelCase_		:   Dict													=       {
								    'image': init_image,
								    'latents': latents,
								    'generator': generator,
								    'num_inference_steps': 50,
								    'guidance_scale': 7.5,
								    'output_type': 'numpy',
								}
								return inputs


							def     _snake_case				(      self)				->					int:
								UpperCAmelCase_		:   Dict													=       StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers')
								sd_pipe.to(_snake_case)
								sd_pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:   int													=       self.get_inputs(_snake_case)
								UpperCAmelCase_		:   str													=       sd_pipe(**_snake_case).images[0]

								UpperCAmelCase_		:   int													=       load_numpy(
								    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
								    '/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy')
								UpperCAmelCase_		:   Dict													=       np.abs(expected_image - image).max()
								assert max_diff < 1e-3





							def     _snake_case				(      self)				->					Union[str, Any]:
								UpperCAmelCase_		:   int													=       StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers')
								UpperCAmelCase_		:   List[Any]													=       DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config)
								sd_pipe.to(_snake_case)
								sd_pipe.set_progress_bar_config(disable=_snake_case)

								UpperCAmelCase_		:   Dict													=       self.get_inputs(_snake_case)
								UpperCAmelCase_		:   Optional[int]													=       25
								UpperCAmelCase_		:   List[Any]													=       sd_pipe(**_snake_case).images[0]

								UpperCAmelCase_		:   int													=       load_numpy(
								    'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
								    '/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy')
								UpperCAmelCase_		:   Tuple													=       np.abs(expected_image - image).max()
								assert max_diff < 1e-3









'''simple docstring'''


import inspect
import unittest

from transformers import ResNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
						import jax
						import jax.numpy as jnp

						from transformers.models.resnet.modeling_flax_resnet import FlaxResNetForImageClassification, FlaxResNetModel

if is_vision_available():
						from PIL import Image

						from transformers import AutoImageProcessor



class        lowercase    (	unittest.TestCase			):


							def __init__(      self     ,      _snake_case     ,      _snake_case=3     ,      _snake_case=32     ,      _snake_case=3     ,      _snake_case=10     ,      _snake_case=[10, 20, 30, 40]     ,      _snake_case=[1, 1, 2, 1]     ,      _snake_case=True     ,      _snake_case=True     ,      _snake_case="relu"     ,      _snake_case=3     ,      _snake_case=None     ,      )				->					Optional[Any]:
								UpperCAmelCase_		:   List[Any]													=       parent
								UpperCAmelCase_		:   Union[str, Any]													=       batch_size
								UpperCAmelCase_		:   List[Any]													=       image_size
								UpperCAmelCase_		:   List[Any]													=       num_channels
								UpperCAmelCase_		:   Dict													=       embeddings_size
								UpperCAmelCase_		:   str													=       hidden_sizes
								UpperCAmelCase_		:   List[Any]													=       depths
								UpperCAmelCase_		:   Optional[Any]													=       is_training
								UpperCAmelCase_		:   str													=       use_labels
								UpperCAmelCase_		:   List[str]													=       hidden_act
								UpperCAmelCase_		:   Dict													=       num_labels
								UpperCAmelCase_		:   List[Any]													=       scope
								UpperCAmelCase_		:   Tuple													=       len(_snake_case)


							def     _snake_case				(      self)				->					int:
								UpperCAmelCase_		:   int													=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

								UpperCAmelCase_		:   Optional[int]													=       self.get_config()

								return config, pixel_values


							def     _snake_case				(      self)				->					Union[str, Any]:
								return ResNetConfig(
								    num_channels=self.num_channels     ,      embeddings_size=self.embeddings_size     ,      hidden_sizes=self.hidden_sizes     ,      depths=self.depths     ,      hidden_act=self.hidden_act     ,      num_labels=self.num_labels     ,      image_size=self.image_size     ,      )


							def     _snake_case				(      self     ,      _snake_case     ,      _snake_case)				->					str:
								UpperCAmelCase_		:   Any													=       FlaxResNetModel(config=_snake_case)
								UpperCAmelCase_		:   Optional[Any]													=       model(_snake_case)

								# Output shape (b, c, h, w)
								self.parent.assertEqual(
								    result.last_hidden_state.shape     ,      (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)     ,      )


							def     _snake_case				(      self     ,      _snake_case     ,      _snake_case)				->					List[str]:
								UpperCAmelCase_		:   Union[str, Any]													=       self.num_labels
								UpperCAmelCase_		:   Optional[Any]													=       FlaxResNetForImageClassification(config=_snake_case)
								UpperCAmelCase_		:   Optional[Any]													=       model(_snake_case)
								self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.num_labels))





							def     _snake_case				(      self)				->					str:
								UpperCAmelCase_		:   List[Any]													=       self.prepare_config_and_inputs()
								UpperCAmelCase_						,     UpperCAmelCase_		:   List[Any]													=       config_and_inputs
								UpperCAmelCase_		:   str													=       {'pixel_values': pixel_values}
								return config, inputs_dict



@require_flax
class        lowercase    (	a_,       unittest.TestCase			):
							_lowerCamelCase		:   Dict= (FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

							_lowerCamelCase		:   Any= False
							_lowerCamelCase		:   List[Any]= False
							_lowerCamelCase		:   str= False


							def     _snake_case				(      self)				->					None:
								UpperCAmelCase_		:   Any													=       FlaxResNetModelTester(self)
								UpperCAmelCase_		:   Optional[int]													=       ConfigTester(self     ,      config_class=_snake_case     ,      has_text_modality=_snake_case)


							def     _snake_case				(      self)				->					Tuple:
								self.create_and_test_config_common_properties()
								self.config_tester.create_and_test_config_to_json_string()
								self.config_tester.create_and_test_config_to_json_file()
								self.config_tester.create_and_test_config_from_and_save_pretrained()
								self.config_tester.create_and_test_config_with_num_labels()
								self.config_tester.check_config_can_be_init_without_params()
								self.config_tester.check_config_arguments_init()


							def     _snake_case				(      self)				->					int:
								return


							def     _snake_case				(      self)				->					Optional[int]:
								UpperCAmelCase_		:   Optional[Any]													=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case)


							def     _snake_case				(      self)				->					Dict:
								UpperCAmelCase_		:   Union[str, Any]													=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_snake_case)


							@unittest.skip(reason='ResNet does not use inputs_embeds')
							def     _snake_case				(      self)				->					Optional[Any]:
								pass


							@unittest.skip(reason='ResNet does not support input and output embeddings')
							def     _snake_case				(      self)				->					Union[str, Any]:
								pass


							def     _snake_case				(      self)				->					Dict:
								UpperCAmelCase_						,     UpperCAmelCase_		:   List[Any]													=       self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									UpperCAmelCase_		:   Any													=       model_class(_snake_case)
									UpperCAmelCase_		:   str													=       inspect.signature(model.__call__)
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									UpperCAmelCase_		:   Tuple													=       [*signature.parameters.keys()]

									UpperCAmelCase_		:   List[str]													=       ['pixel_values']
									self.assertListEqual(arg_names[:1]     ,      _snake_case)


							def     _snake_case				(      self)				->					List[Any]:
								def check_hidden_states_output(_snake_case     ,      _snake_case     ,      _snake_case):
									UpperCAmelCase_		:   Union[str, Any]													=       model_class(_snake_case)

									UpperCAmelCase_		:   Any													=       model(**self._prepare_for_class(_snake_case     ,      _snake_case))

									UpperCAmelCase_		:   List[Any]													=       outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

									UpperCAmelCase_		:   int													=       self.model_tester.num_stages
									self.assertEqual(len(_snake_case)     ,      expected_num_stages + 1)


							@unittest.skip(reason='ResNet does not use feedforward chunking')
							def     _snake_case				(      self)				->					List[Any]:
								pass





							def     _snake_case				(      self)				->					Tuple:
								UpperCAmelCase_						,     UpperCAmelCase_		:   Any													=       self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									with self.subTest(model_class.__name__):
										UpperCAmelCase_		:   str													=       self._prepare_for_class(_snake_case     ,      _snake_case)
										UpperCAmelCase_		:   str													=       model_class(_snake_case)

										@jax.jit
										def model_jitted(_snake_case     ,      **_snake_case):
											return model(pixel_values=_snake_case     ,      **_snake_case)

										with self.subTest('JIT Enabled'):
											UpperCAmelCase_		:   Optional[Any]													=       model_jitted(**_snake_case).to_tuple()

										with self.subTest('JIT Disabled'):
											with jax.disable_jit():
												UpperCAmelCase_		:   Union[str, Any]													=       model_jitted(**_snake_case).to_tuple()

										self.assertEqual(len(_snake_case)     ,      len(_snake_case))
										for jitted_output, output in zip(_snake_case     ,      _snake_case):
											self.assertEqual(jitted_output.shape     ,      output.shape)







def     SCREAMING_SNAKE_CASE(	)     ->     Dict:
	UpperCAmelCase_		:   Tuple													=       Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
	return image







@require_flax
class        lowercase    (	unittest.TestCase			):


							@cached_property
							def     _snake_case				(      self)				->					List[str]:
								return AutoImageProcessor.from_pretrained('microsoft/resnet-50') if is_vision_available() else None





							@slow
							def     _snake_case				(      self)				->					Union[str, Any]:
								UpperCAmelCase_		:   Union[str, Any]													=       FlaxResNetForImageClassification.from_pretrained('microsoft/resnet-50')

								UpperCAmelCase_		:   List[str]													=       self.default_image_processor
								UpperCAmelCase_		:   Union[str, Any]													=       prepare_img()
								UpperCAmelCase_		:   int													=       image_processor(images=_snake_case     ,      return_tensors='np')

								UpperCAmelCase_		:   Any													=       model(**_snake_case)

								# verify the logits
								UpperCAmelCase_		:   Optional[int]													=       (1, 1000)
								self.assertEqual(outputs.logits.shape     ,      _snake_case)

								UpperCAmelCase_		:   Tuple													=       jnp.array([-11.1_069, -9.7_877, -8.3_777])

								self.assertTrue(jnp.allclose(outputs.logits[0, :3]     ,      _snake_case     ,      atol=1e-4))






