


import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()


class        __lowerCamelCase      (				unittest.TestCase ):


				"""simple docstring"""







				def        lowerCAmelCase__			(				self     :    List[Any]					):
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()






				@property
				def        lowerCAmelCase__			(				self     :    Any					):
								_lowerCAmelCase  						=1
								_lowerCAmelCase  						=3
								_lowerCAmelCase  						=(32, 32)

								_lowerCAmelCase  						=floats_tensor((batch_size, num_channels) + sizes	,							rng=random.Random(0					)					).to(__UpperCAmelCase					)
								return image






				@property
				def        lowerCAmelCase__			(				self     :    List[str]					):
								torch.manual_seed(0					)
								_lowerCAmelCase  						=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	,							)
								return model






				@property
				def        lowerCAmelCase__			(				self     :    Optional[int]					):
								torch.manual_seed(0					)
								_lowerCAmelCase  						=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	,							)
								return model






				@property
				def        lowerCAmelCase__			(				self     :    str					):
								torch.manual_seed(0					)
								_lowerCAmelCase  						=CLIPTextConfig(
								    bos_token_id=0	,							eos_token_id=2	,							hidden_size=32	,							intermediate_size=37	,							layer_norm_eps=1e-05	,							num_attention_heads=4	,							num_hidden_layers=5	,							pad_token_id=1	,							vocab_size=1000	,							)
								return CLIPTextModel(__UpperCAmelCase					)






				@property
				def        lowerCAmelCase__			(				self     :    Any					):



								def extract(*lowerCamelCase_     :    Tuple	,							**lowerCamelCase_     :    Tuple					):


												class        __lowerCamelCase      :


																"""simple docstring"""







																def __init__(				self     :    str					):
																				_lowerCAmelCase  						=torch.ones([0]					)






																def        lowerCAmelCase__			(				self     :    str	,							lowerCamelCase_     :    Tuple					):
																				self.pixel_values.to(__UpperCAmelCase					)
																				return self
												return Out()

								return extract






				def        lowerCAmelCase__			(				self     :    int					):
								_lowerCAmelCase  						='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								_lowerCAmelCase  						=self.dummy_cond_unet
								_lowerCAmelCase  						=DDIMScheduler(
								    beta_start=0.0_0085	,							beta_end=0.012	,							beta_schedule="""scaled_linear"""	,							clip_sample=__UpperCAmelCase	,							set_alpha_to_one=__UpperCAmelCase	,							)

								_lowerCAmelCase  						=self.dummy_vae
								_lowerCAmelCase  						=self.dummy_text_encoder
								_lowerCAmelCase  						=CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

								# make sure here that pndm scheduler skips prk
								_lowerCAmelCase  						=StableDiffusionPipeline(
								    unet=__UpperCAmelCase	,							scheduler=__UpperCAmelCase	,							vae=__UpperCAmelCase	,							text_encoder=__UpperCAmelCase	,							tokenizer=__UpperCAmelCase	,							safety_checker=__UpperCAmelCase	,							feature_extractor=self.dummy_extractor	,							)
								_lowerCAmelCase  						=sd_pipe.to(__UpperCAmelCase					)
								sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

								_lowerCAmelCase  						='''A painting of a squirrel eating a burger'''

								_lowerCAmelCase  						=torch.Generator(device=__UpperCAmelCase					).manual_seed(0					)
								_lowerCAmelCase  						=sd_pipe([prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=6.0	,							num_inference_steps=2	,							output_type="""np"""					)
								_lowerCAmelCase  						=output.images

								_lowerCAmelCase  						=torch.Generator(device=__UpperCAmelCase					).manual_seed(0					)
								_lowerCAmelCase  						=sd_pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=6.0	,							num_inference_steps=2	,							output_type="""np"""	,							return_dict=__UpperCAmelCase	,							)[0]

								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)
								_lowerCAmelCase  						=np.array([0.5753, 0.6114, 0.5001, 0.5034, 0.5470, 0.4729, 0.4971, 0.4867, 0.4867]					)

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






				def        lowerCAmelCase__			(				self     :    Optional[int]					):
								_lowerCAmelCase  						='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								_lowerCAmelCase  						=self.dummy_cond_unet
								_lowerCAmelCase  						=PNDMScheduler(skip_prk_steps=__UpperCAmelCase					)
								_lowerCAmelCase  						=self.dummy_vae
								_lowerCAmelCase  						=self.dummy_text_encoder
								_lowerCAmelCase  						=CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

								# make sure here that pndm scheduler skips prk
								_lowerCAmelCase  						=StableDiffusionPipeline(
								    unet=__UpperCAmelCase	,							scheduler=__UpperCAmelCase	,							vae=__UpperCAmelCase	,							text_encoder=__UpperCAmelCase	,							tokenizer=__UpperCAmelCase	,							safety_checker=__UpperCAmelCase	,							feature_extractor=self.dummy_extractor	,							)
								_lowerCAmelCase  						=sd_pipe.to(__UpperCAmelCase					)
								sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

								_lowerCAmelCase  						='''A painting of a squirrel eating a burger'''
								_lowerCAmelCase  						=torch.Generator(device=__UpperCAmelCase					).manual_seed(0					)
								_lowerCAmelCase  						=sd_pipe([prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=6.0	,							num_inference_steps=2	,							output_type="""np"""					)

								_lowerCAmelCase  						=output.images

								_lowerCAmelCase  						=torch.Generator(device=__UpperCAmelCase					).manual_seed(0					)
								_lowerCAmelCase  						=sd_pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=6.0	,							num_inference_steps=2	,							output_type="""np"""	,							return_dict=__UpperCAmelCase	,							)[0]

								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)
								_lowerCAmelCase  						=np.array([0.5122, 0.5712, 0.4825, 0.5053, 0.5646, 0.4769, 0.5179, 0.4894, 0.4994]					)

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






				def        lowerCAmelCase__			(				self     :    int					):
								_lowerCAmelCase  						=StableDiffusionPipeline.from_pretrained(
								    """hf-internal-testing/tiny-stable-diffusion-lms-pipe"""	,							safety_checker=__UpperCAmelCase					)
								assert isinstance(__UpperCAmelCase	,							__UpperCAmelCase					)
								assert isinstance(pipe.scheduler	,							__UpperCAmelCase					)
								assert pipe.safety_checker is None

								_lowerCAmelCase  						=pipe("""example prompt"""	,							num_inference_steps=2					).images[0]
								assert image is not None

								# check that there's no error when saving a pipeline with one of the models being None
								with tempfile.TemporaryDirectory() as tmpdirname:
												pipe.save_pretrained(__UpperCAmelCase					)
												_lowerCAmelCase  						=StableDiffusionPipeline.from_pretrained(__UpperCAmelCase					)

								# sanity check that the pipeline still works
								assert pipe.safety_checker is None
								_lowerCAmelCase  						=pipe("""example prompt"""	,							num_inference_steps=2					).images[0]
								assert image is not None






				@unittest.skipIf(torch_device != """cuda"""	,							"""This test requires a GPU"""					)
				def        lowerCAmelCase__			(				self     :    List[str]					):
								_lowerCAmelCase  						=self.dummy_cond_unet
								_lowerCAmelCase  						=PNDMScheduler(skip_prk_steps=__UpperCAmelCase					)
								_lowerCAmelCase  						=self.dummy_vae
								_lowerCAmelCase  						=self.dummy_text_encoder
								_lowerCAmelCase  						=CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

								# put models in fp16
								_lowerCAmelCase  						=unet.half()
								_lowerCAmelCase  						=vae.half()
								_lowerCAmelCase  						=bert.half()

								# make sure here that pndm scheduler skips prk
								_lowerCAmelCase  						=StableDiffusionPipeline(
								    unet=__UpperCAmelCase	,							scheduler=__UpperCAmelCase	,							vae=__UpperCAmelCase	,							text_encoder=__UpperCAmelCase	,							tokenizer=__UpperCAmelCase	,							safety_checker=__UpperCAmelCase	,							feature_extractor=self.dummy_extractor	,							)
								_lowerCAmelCase  						=sd_pipe.to(__UpperCAmelCase					)
								sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

								_lowerCAmelCase  						='''A painting of a squirrel eating a burger'''
								_lowerCAmelCase  						=sd_pipe([prompt]	,							num_inference_steps=2	,							output_type="""np"""					).images

								assert image.shape == (1, 64, 64, 3)


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


				"""simple docstring"""







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






				def        lowerCAmelCase__			(				self     :    List[Any]					):
								_lowerCAmelCase  						='''cuda'''
								_lowerCAmelCase  						=StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""					)
								_lowerCAmelCase  						=pipe.to(__UpperCAmelCase					)
								pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

								_lowerCAmelCase  						='''a photo of a cat'''
								_lowerCAmelCase  						={
								    '''editing_prompt''': ['''sunglasses'''],
								    '''reverse_editing_direction''': [False],
								    '''edit_warmup_steps''': 10,
								    '''edit_guidance_scale''': 6,
								    '''edit_threshold''': 0.95,
								    '''edit_momentum_scale''': 0.5,
								    '''edit_mom_beta''': 0.6,
								}

								_lowerCAmelCase  						=3
								_lowerCAmelCase  						=7

								# no sega enabled
								_lowerCAmelCase  						=torch.Generator(__UpperCAmelCase					)
								generator.manual_seed(__UpperCAmelCase					)
								_lowerCAmelCase  						=pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=__UpperCAmelCase	,							num_inference_steps=50	,							output_type="""np"""	,							width=512	,							height=512	,							)

								_lowerCAmelCase  						=output.images
								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=[
								    0.3467_3113,
								    0.3849_2733,
								    0.3759_7352,
								    0.3408_6335,
								    0.3565_0748,
								    0.3557_9205,
								    0.338_4763,
								    0.3434_0236,
								    0.357_3271,
								]

								assert image.shape == (1, 512, 512, 3)

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

								# with sega enabled
								# generator = torch.manual_seed(seed)
								generator.manual_seed(__UpperCAmelCase					)
								_lowerCAmelCase  						=pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=__UpperCAmelCase	,							num_inference_steps=50	,							output_type="""np"""	,							width=512	,							height=512	,							**__UpperCAmelCase	,							)

								_lowerCAmelCase  						=output.images
								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=[
								    0.4188_7826,
								    0.3772_8766,
								    0.3013_8272,
								    0.4141_6335,
								    0.4166_4985,
								    0.3628_3392,
								    0.3619_1246,
								    0.4336_4465,
								    0.4300_1732,
								]

								assert image.shape == (1, 512, 512, 3)

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






				def        lowerCAmelCase__			(				self     :    Optional[int]					):
								_lowerCAmelCase  						='''cuda'''
								_lowerCAmelCase  						=StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""					)
								_lowerCAmelCase  						=pipe.to(__UpperCAmelCase					)
								pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

								_lowerCAmelCase  						='''an image of a crowded boulevard, realistic, 4k'''
								_lowerCAmelCase  						={
								    '''editing_prompt''': '''crowd, crowded, people''',
								    '''reverse_editing_direction''': True,
								    '''edit_warmup_steps''': 10,
								    '''edit_guidance_scale''': 8.3,
								    '''edit_threshold''': 0.9,
								    '''edit_momentum_scale''': 0.5,
								    '''edit_mom_beta''': 0.6,
								}

								_lowerCAmelCase  						=9
								_lowerCAmelCase  						=7

								# no sega enabled
								_lowerCAmelCase  						=torch.Generator(__UpperCAmelCase					)
								generator.manual_seed(__UpperCAmelCase					)
								_lowerCAmelCase  						=pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=__UpperCAmelCase	,							num_inference_steps=50	,							output_type="""np"""	,							width=512	,							height=512	,							)

								_lowerCAmelCase  						=output.images
								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=[
								    0.4349_7998,
								    0.9181_4065,
								    0.754_0739,
								    0.5558_0205,
								    0.846_7265,
								    0.538_9691,
								    0.6257_4506,
								    0.5889_7763,
								    0.5092_6757,
								]

								assert image.shape == (1, 512, 512, 3)

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

								# with sega enabled
								# generator = torch.manual_seed(seed)
								generator.manual_seed(__UpperCAmelCase					)
								_lowerCAmelCase  						=pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=__UpperCAmelCase	,							num_inference_steps=50	,							output_type="""np"""	,							width=512	,							height=512	,							**__UpperCAmelCase	,							)

								_lowerCAmelCase  						=output.images
								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=[
								    0.308_9719,
								    0.3050_0144,
								    0.2901_6042,
								    0.3063_0964,
								    0.32_5687,
								    0.2941_9225,
								    0.290_8091,
								    0.2872_3598,
								    0.2769_6294,
								]

								assert image.shape == (1, 512, 512, 3)

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






				def        lowerCAmelCase__			(				self     :    Tuple					):
								_lowerCAmelCase  						='''cuda'''
								_lowerCAmelCase  						=StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""					)
								_lowerCAmelCase  						=pipe.to(__UpperCAmelCase					)
								pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

								_lowerCAmelCase  						='''a castle next to a river'''
								_lowerCAmelCase  						={
								    '''editing_prompt''': ['''boat on a river, boat''', '''monet, impression, sunrise'''],
								    '''reverse_editing_direction''': False,
								    '''edit_warmup_steps''': [15, 18],
								    '''edit_guidance_scale''': 6,
								    '''edit_threshold''': [0.9, 0.8],
								    '''edit_momentum_scale''': 0.5,
								    '''edit_mom_beta''': 0.6,
								}

								_lowerCAmelCase  						=48
								_lowerCAmelCase  						=7

								# no sega enabled
								_lowerCAmelCase  						=torch.Generator(__UpperCAmelCase					)
								generator.manual_seed(__UpperCAmelCase					)
								_lowerCAmelCase  						=pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=__UpperCAmelCase	,							num_inference_steps=50	,							output_type="""np"""	,							width=512	,							height=512	,							)

								_lowerCAmelCase  						=output.images
								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=[
								    0.7516_3555,
								    0.7603_7145,
								    0.6_1785,
								    0.918_9673,
								    0.862_7701,
								    0.8518_9694,
								    0.851_2813,
								    0.8701_2076,
								    0.831_2857,
								]

								assert image.shape == (1, 512, 512, 3)

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

								# with sega enabled
								# generator = torch.manual_seed(seed)
								generator.manual_seed(__UpperCAmelCase					)
								_lowerCAmelCase  						=pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=__UpperCAmelCase	,							num_inference_steps=50	,							output_type="""np"""	,							width=512	,							height=512	,							**__UpperCAmelCase	,							)

								_lowerCAmelCase  						=output.images
								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=[
								    0.7355_3365,
								    0.753_7271,
								    0.7434_1905,
								    0.6648_0356,
								    0.647_2925,
								    0.6303_9416,
								    0.6481_2905,
								    0.674_9717,
								    0.651_7102,
								]

								assert image.shape == (1, 512, 512, 3)

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






				def        lowerCAmelCase__			(				self     :    Any					):
								_lowerCAmelCase  						='''cuda'''
								_lowerCAmelCase  						=StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""	,							torch_dtype=torch.floataa					)
								_lowerCAmelCase  						=pipe.to(__UpperCAmelCase					)
								pipe.set_progress_bar_config(disable=__UpperCAmelCase					)

								_lowerCAmelCase  						='''a photo of a cat'''
								_lowerCAmelCase  						={
								    '''editing_prompt''': ['''sunglasses'''],
								    '''reverse_editing_direction''': [False],
								    '''edit_warmup_steps''': 10,
								    '''edit_guidance_scale''': 6,
								    '''edit_threshold''': 0.95,
								    '''edit_momentum_scale''': 0.5,
								    '''edit_mom_beta''': 0.6,
								}

								_lowerCAmelCase  						=3
								_lowerCAmelCase  						=7

								# no sega enabled
								_lowerCAmelCase  						=torch.Generator(__UpperCAmelCase					)
								generator.manual_seed(__UpperCAmelCase					)
								_lowerCAmelCase  						=pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=__UpperCAmelCase	,							num_inference_steps=50	,							output_type="""np"""	,							width=512	,							height=512	,							)

								_lowerCAmelCase  						=output.images
								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=[
								    0.3488_7695,
								    0.387_6953,
								    0.375,
								    0.3442_3828,
								    0.358_1543,
								    0.3571_7773,
								    0.338_3789,
								    0.3457_0312,
								    0.35_9375,
								]

								assert image.shape == (1, 512, 512, 3)

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

								# with sega enabled
								# generator = torch.manual_seed(seed)
								generator.manual_seed(__UpperCAmelCase					)
								_lowerCAmelCase  						=pipe(
								    [prompt]	,							generator=__UpperCAmelCase	,							guidance_scale=__UpperCAmelCase	,							num_inference_steps=50	,							output_type="""np"""	,							width=512	,							height=512	,							**__UpperCAmelCase	,							)

								_lowerCAmelCase  						=output.images
								_lowerCAmelCase  						=image[0, -3:, -3:, -1]
								_lowerCAmelCase  						=[
								    0.4228_5156,
								    0.3691_4062,
								    0.2907_7148,
								    0.4204_1016,
								    0.4191_8945,
								    0.3549_8047,
								    0.361_8164,
								    0.442_3828,
								    0.4311_5234,
								]

								assert image.shape == (1, 512, 512, 3)

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


"""simple docstring"""




import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
					import torch

					from transformers import GPTaLMHeadModel

@require_torch
@require_sentencepiece
@require_tokenizers
class        __SCREAMING_SNAKE_CASE     (   unittest.TestCase							):







				@slow
				@unittest.skip('''Model is not available.'''  )
				def   lowercase_		(    self    :Union[str, Any]  )      ->					Optional[int]:




									"""simple docstring"""



									lowerCamelCase__		:   Any  				=					'''nvidia/megatron-gpt2-345m/'''
									if "MYDIR" in os.environ:
														lowerCamelCase__		:   str  				=					os.path.join(os.environ['''MYDIR''']      ,__UpperCAmelCase  )
									lowerCamelCase__		:   Dict  				=					GPTaLMHeadModel.from_pretrained(__UpperCAmelCase  )
									model.to(__UpperCAmelCase  )
									model.half()

									lowerCamelCase__		:   Dict  				=					torch.tensor(
									    [[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]      ,device=__UpperCAmelCase      ,dtype=torch.long      ,)

									with torch.no_grad():
														lowerCamelCase__		:   Tuple  				=					model(__UpperCAmelCase  ).logits

									lowerCamelCase__		:   List[str]  				=					torch.Size((1, 9, 5_02_57)  )
									self.assertEqual(output.shape      ,__UpperCAmelCase  )

									lowerCamelCase__		:   Tuple  				=					torch.tensor(
									    [
									        4.9_414,
									        -0.2_920,
									        -1.2_148,
									        -4.0_273,
									        -0.5_161,
									        -5.2_109,
									        -1.2_412,
									        -1.8_301,
									        -1.7_734,
									        -4.7_148,
									        -0.2_317,
									        -1.0_811,
									        -2.1_777,
									        0.4_141,
									        -3.7_969,
									        -4.0_586,
									        -2.5_332,
									        -3.3_809,
									        4.3_867,
									    ]      ,device=__UpperCAmelCase      ,dtype=torch.half      ,)

									for i in range(19  ):
														lowerCamelCase__				,						lowerCamelCase__		:   Dict  				=					8 * i // 17, 27_92 * i  # along the diagonal
														lowerCamelCase__				,						lowerCamelCase__		:   Optional[int]  				=					output[0, r, c], expected_diag[i]
														lowerCamelCase__		:   Dict  				=					F"""row={r} col={c} computed={computed} expected={expected}"""
														self.assertAlmostEqual(__UpperCAmelCase      ,__UpperCAmelCase      ,delta=1E-4      ,msg=__UpperCAmelCase  )






