"""simple docstring"""






import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest
class 							snake_case			(					__UpperCAmelCase							):




					"""simple docstring"""

					snake_case__							       =					(VQDiffusionScheduler,)







					def 						__lowerCAmelCase		(							self  :      Optional[Any]       ,**lowerCamelCase__  :      str     ):
									UpperCAmelCase__			 =				{
									    'num_vec_classes': 4_097,
									    'num_train_timesteps': 100,
									}

									config.update(**lowerCamelCase__     )
									return config







					def 						__lowerCAmelCase		(							self  :      Dict       ,lowerCamelCase__  :      str     ):
									UpperCAmelCase__			 =				4
									UpperCAmelCase__			 =				8
									UpperCAmelCase__			 =				8

									UpperCAmelCase__			 =				torch.randint(0       ,lowerCamelCase__       ,(batch_size, height * width)     )

									return sample







					@property
					def 						__lowerCAmelCase		(							self  :      Tuple     ):
									assert False







					def 						__lowerCAmelCase		(							self  :      str       ,lowerCamelCase__  :      str     ):
									def model(lowerCamelCase__  :      Optional[Any]       ,lowerCamelCase__  :      List[Any]       ,*lowerCamelCase__  :      Any     ):
													UpperCAmelCase__ ,      UpperCAmelCase__			 =				sample.shape
													UpperCAmelCase__			 =				torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)     )
													UpperCAmelCase__			 =				F.log_softmax(logits.double()       ,dim=1     ).float()
													return return_value

									return model







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







					def 						__lowerCAmelCase		(							self  :      Optional[int]     ):
									for num_vec_classes in [5, 100, 1_000, 4_000]:
													self.check_over_configs(num_vec_classes=lowerCamelCase__     )







					def 						__lowerCAmelCase		(							self  :      Optional[Any]     ):
									for t in [0, 50, 99]:
													self.check_over_forward(time_step=lowerCamelCase__     )







					def 						__lowerCAmelCase		(							self  :      Union[str, Any]     ):
									pass




"""simple docstring"""






import os
import unittest

import pytest
from parameterized import parameterized

from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device


# skipping in unittest tests

lowerCAmelCase__    :    Union[str, Any]			=      [(1,)]



def     a_		(   ):
				UpperCAmelCase__			 =				bool(os.getenv('RUN_SLOW'   ,     0			)			)
				if run_slow:
								assert True
				else:
								assert False, "should have been skipped"



def     a_		(   ):
				UpperCAmelCase__			 =				bool(os.getenv('RUN_SLOW'   ,     0			)			)
				if run_slow and torch_device == "cuda":
								assert True
				else:
								assert False, "should have been skipped"
@require_torch
class 							snake_case			(					unittest.TestCase							):




					"""simple docstring"""

					@slow
					@require_torch_gpu
					def 						__lowerCAmelCase		(							self  :      Optional[Any]     ):
									check_slow_torch_cuda()







					@require_torch_gpu
					@slow
					def 						__lowerCAmelCase		(							self  :      Any     ):
									check_slow_torch_cuda()







					@parameterized.expand(lowerCamelCase__     )
					@slow
					def 						__lowerCAmelCase		(							self  :      Union[str, Any]       ,lowerCamelCase__  :      str=None     ):
									check_slow()






@slow
@require_torch_gpu
def     a_		(   ):
				check_slow_torch_cuda()



@require_torch_gpu
@slow
def     a_		(   ):
				check_slow_torch_cuda()



@slow
@pytest.mark.parametrize('param'   ,     [1]			)
def     a_		(   lowerCamelCase			):
				check_slow()







@pytest.mark.parametrize('param'   ,     [1]			)
@slow
def     a_		(   lowerCamelCase			):
				check_slow()



