




import torch

from transformers import PreTrainedModel

from .custom_configuration import CustomConfig, NoSuperInitConfig




class    snake_case__							(_UpperCamelCase       ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_	:							Union[str, Any]										= CustomConfig





				def __init__(				self    :					Any ,				__lowerCamelCase    :					Tuple							) ->					Any:
							super().__init__(__lowerCamelCase							)
							a   =      torch.nn.Linear(config.hidden_size ,				config.hidden_size							)





				def 							__UpperCAmelCase							(				self    :					Optional[int] ,				__lowerCamelCase    :					List[str]							) ->					Optional[int]:
							return self.linear(__lowerCamelCase							)





				def 							__UpperCAmelCase							(				self    :					int ,				__lowerCamelCase    :					List[str]							) ->					Any:
							pass




class    snake_case__							(_UpperCamelCase       ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_	:							Dict										= NoSuperInitConfig





				def __init__(				self    :					List[Any] ,				__lowerCamelCase    :					Optional[Any]							) ->					List[str]:
							super().__init__(__lowerCamelCase							)
							a   =      torch.nn.Linear(config.attribute ,				config.attribute							)





				def 							__UpperCAmelCase							(				self    :					Optional[int] ,				__lowerCamelCase    :					Union[str, Any]							) ->					str:
							return self.linear(__lowerCamelCase							)





				def 							__UpperCAmelCase							(				self    :					int ,				__lowerCamelCase    :					Optional[int]							) ->					List[Any]:
							pass












import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest




class    snake_case__							(_UpperCamelCase       ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_	:							Dict										= (VQDiffusionScheduler,)





				def 							__UpperCAmelCase							(				self    :					Tuple ,				**__lowerCamelCase    :					Optional[int]							) ->					Optional[Any]:
							a   =      {
							    "num_vec_classes": 40_97,
							    "num_train_timesteps": 1_00,
							}

							config.update(**__lowerCamelCase							)
							return config





				def 							__UpperCAmelCase							(				self    :					List[str] ,				__lowerCamelCase    :					Any							) ->					Union[str, Any]:
							a   =      4
							a   =      8
							a   =      8

							a   =      torch.randint(0 ,				__lowerCamelCase ,				(batch_size, height * width)							)

							return sample





				@property
				def 							__UpperCAmelCase							(				self    :					List[str]							) ->					Optional[Any]:
							assert False





				def 							__UpperCAmelCase							(				self    :					Any ,				__lowerCamelCase    :					Optional[int]							) ->					int:
							def model(__lowerCamelCase    :					Tuple ,				__lowerCamelCase    :					Any ,				*__lowerCamelCase    :					Optional[Any]							):
										a   ,		a   =      sample.shape
										a   =      torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)							)
										a   =      F.log_softmax(logits.double() ,				dim=1							).float()
										return return_value

							return model





				def 							__UpperCAmelCase							(				self    :					Optional[Any]							) ->					Any:
							for timesteps in [2, 5, 1_00, 10_00]:
										self.check_over_configs(num_train_timesteps=__lowerCamelCase							)





				def 							__UpperCAmelCase							(				self    :					Any							) ->					List[Any]:
							for num_vec_classes in [5, 1_00, 10_00, 40_00]:
										self.check_over_configs(num_vec_classes=__lowerCamelCase							)





				def 							__UpperCAmelCase							(				self    :					Optional[Any]							) ->					List[Any]:
							for t in [0, 50, 99]:
										self.check_over_forward(time_step=__lowerCamelCase							)





				def 							__UpperCAmelCase							(				self    :					List[Any]							) ->					Optional[Any]:
							pass






