
'''simple docstring'''







import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest


class       lowerCAmelCase			(					__UpperCAmelCase      ):
		a				:       Dict							=       (VQDiffusionScheduler,)




		def      lowercase							(					self       ,    **UpperCamelCase   ):
								_SCREAMING_SNAKE_CASE					=			{
								    "num_vec_classes": 4_097,
								    "num_train_timesteps": 100,
								}

								config.update(**UpperCamelCase   )
								return config




		def      lowercase							(					self       ,    UpperCamelCase   ):
								_SCREAMING_SNAKE_CASE					=			4
								_SCREAMING_SNAKE_CASE					=			8
								_SCREAMING_SNAKE_CASE					=			8

								_SCREAMING_SNAKE_CASE					=			torch.randint(0       ,    UpperCamelCase       ,    (batch_size, height * width)   )

								return sample




		@property
		def      lowercase							(					self   ):
								assert False




		def      lowercase							(					self       ,    UpperCamelCase   ):
								def model(UpperCamelCase       ,    UpperCamelCase       ,    *UpperCamelCase   ):
														_SCREAMING_SNAKE_CASE,				_SCREAMING_SNAKE_CASE					=			sample.shape
														_SCREAMING_SNAKE_CASE					=			torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)   )
														_SCREAMING_SNAKE_CASE					=			F.log_softmax(logits.double()       ,    dim=1   ).float()
														return return_value

								return model




		def      lowercase							(					self   ):
								for timesteps in [2, 5, 100, 1_000]:
														self.check_over_configs(num_train_timesteps=UpperCamelCase   )




		def      lowercase							(					self   ):
								for num_vec_classes in [5, 100, 1_000, 4_000]:
														self.check_over_configs(num_vec_classes=UpperCamelCase   )




		def      lowercase							(					self   ):
								for t in [0, 50, 99]:
														self.check_over_forward(time_step=UpperCamelCase   )




		def      lowercase							(					self   ):
								pass

'''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       lowerCAmelCase			(					unittest.TestCase      ):




		@slow
		@unittest.skip("Model is not available."   )
		def      lowercase							(					self   ):
								_SCREAMING_SNAKE_CASE					=			"nvidia/megatron-gpt2-345m/"
								if "MYDIR" in os.environ:
														_SCREAMING_SNAKE_CASE					=			os.path.join(os.environ["MYDIR"]       ,    UpperCamelCase   )
								_SCREAMING_SNAKE_CASE					=			GPTaLMHeadModel.from_pretrained(UpperCamelCase   )
								model.to(UpperCamelCase   )
								model.half()

								_SCREAMING_SNAKE_CASE					=			torch.tensor(
								    [[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]       ,    device=UpperCamelCase       ,    dtype=torch.long       ,    )

								with torch.no_grad():
														_SCREAMING_SNAKE_CASE					=			model(UpperCamelCase   ).logits

								_SCREAMING_SNAKE_CASE					=			torch.Size((1, 9, 50_257)   )
								self.assertEqual(output.shape       ,    UpperCamelCase   )

								_SCREAMING_SNAKE_CASE					=			torch.tensor(
								    [
								        4.94_14,
								        -0.29_20,
								        -1.21_48,
								        -4.02_73,
								        -0.51_61,
								        -5.21_09,
								        -1.24_12,
								        -1.83_01,
								        -1.77_34,
								        -4.71_48,
								        -0.23_17,
								        -1.08_11,
								        -2.17_77,
								        0.41_41,
								        -3.79_69,
								        -4.05_86,
								        -2.53_32,
								        -3.38_09,
								        4.38_67,
								    ]       ,    device=UpperCamelCase       ,    dtype=torch.half       ,    )

								for i in range(19   ):
														_SCREAMING_SNAKE_CASE,				_SCREAMING_SNAKE_CASE					=			8 * i // 17, 2_792 * i  # along the diagonal
														_SCREAMING_SNAKE_CASE,				_SCREAMING_SNAKE_CASE					=			output[0, r, c], expected_diag[i]
														_SCREAMING_SNAKE_CASE					=			F'row={r} col={c} computed={computed} expected={expected}'
														self.assertAlmostEqual(UpperCamelCase       ,    UpperCamelCase       ,    delta=1e-4       ,    msg=UpperCamelCase   )