





import tensorflow as tf

from ...tf_utils import shape_list

class 	UpperCamelCase__				(   tf.keras.layers.Layer     ):



     '''simple docstring'''

     def __init__(							self		:					str		,lowerCamelCase__		:					Any		,lowerCamelCase__		:					Union[str, Any]		,lowerCamelCase__		:					Dict		,lowerCamelCase__		:					Union[str, Any]		,lowerCamelCase__		:					Any=1		,lowerCamelCase__		:					Tuple=False		,**lowerCamelCase__		:					Dict )  ->	Any:




           '''simple docstring'''


           super().__init__(**lowerCamelCase__ )

           SCREAMING_SNAKE_CASE	=		vocab_size
           SCREAMING_SNAKE_CASE	=		d_embed
           SCREAMING_SNAKE_CASE	=		d_proj

           SCREAMING_SNAKE_CASE	=		cutoffs + [vocab_size]
           SCREAMING_SNAKE_CASE	=		[0] + self.cutoffs
           SCREAMING_SNAKE_CASE	=		div_val

           SCREAMING_SNAKE_CASE	=		self.cutoffs[0]
           SCREAMING_SNAKE_CASE	=		len(self.cutoffs ) - 1
           SCREAMING_SNAKE_CASE	=		self.shortlist_size + self.n_clusters
           SCREAMING_SNAKE_CASE	=		keep_order

           SCREAMING_SNAKE_CASE	=		[]
           SCREAMING_SNAKE_CASE	=		[]






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[Any]		,lowerCamelCase__		:					Optional[Any] )  ->	Dict:




           '''simple docstring'''


           if self.n_clusters > 0:
                 SCREAMING_SNAKE_CASE	=		self.add_weight(
                     shape=(self.n_clusters, self.d_embed)		,initializer="""zeros"""		,trainable=lowerCamelCase__		,name="""cluster_weight""" )
                 SCREAMING_SNAKE_CASE	=		self.add_weight(
                     shape=(self.n_clusters,)		,initializer="""zeros"""		,trainable=lowerCamelCase__		,name="""cluster_bias""" )

           if self.div_val == 1:
                 for i in range(len(self.cutoffs ) ):
                       if self.d_proj != self.d_embed:
                             SCREAMING_SNAKE_CASE	=		self.add_weight(
                                 shape=(self.d_embed, self.d_proj)		,initializer="""zeros"""		,trainable=lowerCamelCase__		,name=F"""out_projs_._{i}"""		,)
                             self.out_projs.append(lowerCamelCase__ )
                       else:
                             self.out_projs.append(lowerCamelCase__ )
                       SCREAMING_SNAKE_CASE	=		self.add_weight(
                           shape=(self.vocab_size, self.d_embed)		,initializer="""zeros"""		,trainable=lowerCamelCase__		,name=F"""out_layers_._{i}_._weight"""		,)
                       SCREAMING_SNAKE_CASE	=		self.add_weight(
                           shape=(self.vocab_size,)		,initializer="""zeros"""		,trainable=lowerCamelCase__		,name=F"""out_layers_._{i}_._bias"""		,)
                       self.out_layers.append((weight, bias) )
           else:
                 for i in range(len(self.cutoffs ) ):
                       SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self.cutoff_ends[i], self.cutoff_ends[i + 1]
                       SCREAMING_SNAKE_CASE	=		self.d_embed // (self.div_val**i)

                       SCREAMING_SNAKE_CASE	=		self.add_weight(
                           shape=(d_emb_i, self.d_proj)		,initializer="""zeros"""		,trainable=lowerCamelCase__		,name=F"""out_projs_._{i}""" )
                       self.out_projs.append(lowerCamelCase__ )
                       SCREAMING_SNAKE_CASE	=		self.add_weight(
                           shape=(r_idx - l_idx, d_emb_i)		,initializer="""zeros"""		,trainable=lowerCamelCase__		,name=F"""out_layers_._{i}_._weight"""		,)
                       SCREAMING_SNAKE_CASE	=		self.add_weight(
                           shape=(r_idx - l_idx,)		,initializer="""zeros"""		,trainable=lowerCamelCase__		,name=F"""out_layers_._{i}_._bias"""		,)
                       self.out_layers.append((weight, bias) )
           super().build(lowerCamelCase__ )






     @staticmethod
     def        SCREAMING_SNAKE_CASE__						(							lowerCamelCase__		:					Optional[int]		,lowerCamelCase__		:					Optional[Any]		,lowerCamelCase__		:					Optional[Any]		,lowerCamelCase__		:					Union[str, Any]=None )  ->	int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		x
           if proj is not None:
                 SCREAMING_SNAKE_CASE	=		tf.einsum("""ibd,ed->ibe"""		,lowerCamelCase__		,lowerCamelCase__ )
           return tf.einsum("""ibd,nd->ibn"""		,lowerCamelCase__		,lowerCamelCase__ ) + b






     @staticmethod
     def        SCREAMING_SNAKE_CASE__						(							lowerCamelCase__		:					int		,lowerCamelCase__		:					Dict )  ->	Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		shape_list(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		tf.range(lp_size[0]		,dtype=target.dtype )
           SCREAMING_SNAKE_CASE	=		tf.stack([r, target]		,1 )
           return tf.gather_nd(lowerCamelCase__		,lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					Any		,lowerCamelCase__		:					Union[str, Any]		,lowerCamelCase__		:					str		,lowerCamelCase__		:					Union[str, Any]=True		,lowerCamelCase__		:					Tuple=False )  ->	List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		0
           if self.n_clusters == 0:
                 SCREAMING_SNAKE_CASE	=		self._logit(lowerCamelCase__		,self.out_layers[0][0]		,self.out_layers[0][1]		,self.out_projs[0] )
                 if target is not None:
                       SCREAMING_SNAKE_CASE	=		tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowerCamelCase__		,logits=lowerCamelCase__ )
                 SCREAMING_SNAKE_CASE	=		tf.nn.log_softmax(lowerCamelCase__		,axis=-1 )
           else:
                 SCREAMING_SNAKE_CASE	=		shape_list(lowerCamelCase__ )
                 SCREAMING_SNAKE_CASE	=		[]
                 SCREAMING_SNAKE_CASE	=		tf.zeros(hidden_sizes[:2] )
                 for i in range(len(self.cutoffs ) ):
                       SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self.cutoff_ends[i], self.cutoff_ends[i + 1]
                       if target is not None:
                             SCREAMING_SNAKE_CASE	=		(target >= l_idx) & (target < r_idx)
                             SCREAMING_SNAKE_CASE	=		tf.where(lowerCamelCase__ )
                             SCREAMING_SNAKE_CASE	=		tf.boolean_mask(lowerCamelCase__		,lowerCamelCase__ ) - l_idx

                       if self.div_val == 1:
                             SCREAMING_SNAKE_CASE	=		self.out_layers[0][0][l_idx:r_idx]
                             SCREAMING_SNAKE_CASE	=		self.out_layers[0][1][l_idx:r_idx]
                       else:
                             SCREAMING_SNAKE_CASE	=		self.out_layers[i][0]
                             SCREAMING_SNAKE_CASE	=		self.out_layers[i][1]

                       if i == 0:
                             SCREAMING_SNAKE_CASE	=		tf.concat([cur_W, self.cluster_weight]		,0 )
                             SCREAMING_SNAKE_CASE	=		tf.concat([cur_b, self.cluster_bias]		,0 )

                             SCREAMING_SNAKE_CASE	=		self._logit(lowerCamelCase__		,lowerCamelCase__		,lowerCamelCase__		,self.out_projs[0] )
                             SCREAMING_SNAKE_CASE	=		tf.nn.log_softmax(lowerCamelCase__ )
                             out.append(head_logprob[..., : self.cutoffs[0]] )
                             if target is not None:
                                   SCREAMING_SNAKE_CASE	=		tf.boolean_mask(lowerCamelCase__		,lowerCamelCase__ )
                                   SCREAMING_SNAKE_CASE	=		self._gather_logprob(lowerCamelCase__		,lowerCamelCase__ )
                       else:
                             SCREAMING_SNAKE_CASE	=		self._logit(lowerCamelCase__		,lowerCamelCase__		,lowerCamelCase__		,self.out_projs[i] )
                             SCREAMING_SNAKE_CASE	=		tf.nn.log_softmax(lowerCamelCase__ )
                             SCREAMING_SNAKE_CASE	=		self.cutoffs[0] + i - 1  # No probability for the head cluster
                             SCREAMING_SNAKE_CASE	=		head_logprob[..., cluster_prob_idx, None] + tail_logprob
                             out.append(lowerCamelCase__ )
                             if target is not None:
                                   SCREAMING_SNAKE_CASE	=		tf.boolean_mask(lowerCamelCase__		,lowerCamelCase__ )
                                   SCREAMING_SNAKE_CASE	=		tf.boolean_mask(lowerCamelCase__		,lowerCamelCase__ )
                                   SCREAMING_SNAKE_CASE	=		self._gather_logprob(lowerCamelCase__		,lowerCamelCase__ )
                                   cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
                       if target is not None:
                             loss += tf.scatter_nd(lowerCamelCase__		,-cur_logprob		,shape_list(lowerCamelCase__ ) )
                 SCREAMING_SNAKE_CASE	=		tf.concat(lowerCamelCase__		,axis=-1 )

           if target is not None:
                 if return_mean:
                       SCREAMING_SNAKE_CASE	=		tf.reduce_mean(lowerCamelCase__ )
                 # Add the training-time loss value to the layer using `self.add_loss()`.
                 self.add_loss(lowerCamelCase__ )

                 # Log the loss as a metric (we could log arbitrary metrics,
                 # including different metrics for training and inference.
                 self.add_metric(lowerCamelCase__		,name=self.name		,aggregation="""mean""" if return_mean else """""" )

           return out









import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

class 	UpperCamelCase__				(   lowerCAmelCase_							,   lowerCAmelCase_							,   unittest.TestCase     ):



     '''simple docstring'''

     __snake_case					:			str				=    VQModel
     __snake_case					:			Optional[Any]				=    "sample"






     @property
     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple		,lowerCamelCase__		:					Optional[Any]=(32, 32) )  ->	Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		4
           SCREAMING_SNAKE_CASE	=		3

           SCREAMING_SNAKE_CASE	=		floats_tensor((batch_size, num_channels) + sizes ).to(lowerCamelCase__ )

           return {"sample": image}






     @property
     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[Any] )  ->	List[str]:




           '''simple docstring'''


           return (3, 32, 32)






     @property
     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple )  ->	List[str]:




           '''simple docstring'''


           return (3, 32, 32)






     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[int] )  ->	Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		{
               """block_out_channels""": [32, 64],
               """in_channels""": 3,
               """out_channels""": 3,
               """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
               """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
               """latent_channels""": 3,
           }
           SCREAMING_SNAKE_CASE	=		self.dummy_input
           return init_dict, inputs_dict






     def        SCREAMING_SNAKE_CASE__						(							self		:					int )  ->	str:




           '''simple docstring'''


           pass






     def        SCREAMING_SNAKE_CASE__						(							self		:					Union[str, Any] )  ->	List[str]:




           '''simple docstring'''


           pass






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[str] )  ->	Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		VQModel.from_pretrained("""fusing/vqgan-dummy"""		,output_loading_info=lowerCamelCase__ )
           self.assertIsNotNone(lowerCamelCase__ )
           self.assertEqual(len(loading_info["""missing_keys"""] )		,0 )

           model.to(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		model(**self.dummy_input )

           assert image is not None, "Make sure output is not None"






     def        SCREAMING_SNAKE_CASE__						(							self		:					str )  ->	List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		VQModel.from_pretrained("""fusing/vqgan-dummy""" )
           model.to(lowerCamelCase__ ).eval()

           torch.manual_seed(0 )
           if torch.cuda.is_available():
                 torch.cuda.manual_seed_all(0 )

           SCREAMING_SNAKE_CASE	=		torch.randn(1		,model.config.in_channels		,model.config.sample_size		,model.config.sample_size )
           SCREAMING_SNAKE_CASE	=		image.to(lowerCamelCase__ )
           with torch.no_grad():
                 SCREAMING_SNAKE_CASE	=		model(lowerCamelCase__ ).sample

           SCREAMING_SNAKE_CASE	=		output[0, -1, -3:, -3:].flatten().cpu()
           # fmt: off
           SCREAMING_SNAKE_CASE	=		torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143] )
           # fmt: on
           self.assertTrue(torch.allclose(lowerCamelCase__		,lowerCamelCase__		,atol=1e-3 ) )


