





import tensorflow as tf

from ...tf_utils import shape_list
class 			UpperCamelCase_		(  tf.keras.layers.Layer					):




					def __init__(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_=1				,						lowerCAmelCase_=False				,						**lowerCAmelCase_  )		->							Any:
										super().__init__(**lowerCAmelCase_  )

										_snake_case									=				vocab_size
										_snake_case									=				d_embed
										_snake_case									=				d_proj

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

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

										_snake_case									=				[]
										_snake_case									=				[]




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_  )		->							Optional[int]:
										if self.n_clusters > 0:
															_snake_case									=				self.add_weight(
															    shape=(self.n_clusters, self.d_embed)				,						initializer='zeros'				,						trainable=lowerCAmelCase_				,						name='cluster_weight'  )
															_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:
																									_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_  )
																				_snake_case									=				self.add_weight(
																				    shape=(self.vocab_size, self.d_embed)				,						initializer='zeros'				,						trainable=lowerCAmelCase_				,						name=F'''out_layers_._{i}_._weight'''				,						)
																				_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  )  ):
																				_snake_case							,     _snake_case									=				self.cutoff_ends[i], self.cutoff_ends[i + 1]
																				_snake_case									=				self.d_embed // (self.div_val**i)

																				_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_  )
																				_snake_case									=				self.add_weight(
																				    shape=(r_idx - l_idx, d_emb_i)				,						initializer='zeros'				,						trainable=lowerCAmelCase_				,						name=F'''out_layers_._{i}_._weight'''				,						)
																				_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 	lowerCAmelCase			(		lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_=None  )		->							Dict:
										_snake_case									=				x
										if proj is not None:
															_snake_case									=				tf.einsum('ibd,ed->ibe'				,						lowerCAmelCase_				,						lowerCAmelCase_  )
										return tf.einsum('ibd,nd->ibn'				,						lowerCAmelCase_				,						lowerCAmelCase_  ) + b




					@staticmethod
					def 	lowerCAmelCase			(		lowerCAmelCase_				,						lowerCAmelCase_  )		->							Tuple:
										_snake_case									=				shape_list(lowerCAmelCase_  )
										_snake_case									=				tf.range(lp_size[0]				,						dtype=target.dtype  )
										_snake_case									=				tf.stack([r, target]				,						1  )
										return tf.gather_nd(lowerCAmelCase_				,						lowerCAmelCase_  )






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

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

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

																									_snake_case									=				self._logit(lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						self.out_projs[0]  )
																									_snake_case									=				tf.nn.log_softmax(lowerCAmelCase_  )
																									out.append(head_logprob[..., : self.cutoffs[0]]  )
																									if target is not None:
																														_snake_case									=				tf.boolean_mask(lowerCAmelCase_				,						lowerCAmelCase_  )
																														_snake_case									=				self._gather_logprob(lowerCAmelCase_				,						lowerCAmelCase_  )
																				else:
																									_snake_case									=				self._logit(lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						self.out_projs[i]  )
																									_snake_case									=				tf.nn.log_softmax(lowerCAmelCase_  )
																									_snake_case									=				self.cutoffs[0] + i - 1  # No probability for the head cluster
																									_snake_case									=				head_logprob[..., cluster_prob_idx, None] + tail_logprob
																									out.append(lowerCAmelCase_  )
																									if target is not None:
																														_snake_case									=				tf.boolean_mask(lowerCAmelCase_				,						lowerCAmelCase_  )
																														_snake_case									=				tf.boolean_mask(lowerCAmelCase_				,						lowerCAmelCase_  )
																														_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_  )  )
															_snake_case									=				tf.concat(lowerCAmelCase_				,						axis=-1  )

										if target is not None:
															if return_mean:
																				_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










from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_      							=     logging.get_logger(__name__)

UpperCAmelCase_      							=     {
    """sayakpaul/vit-msn-base""": """https://huggingface.co/sayakpaul/vit-msn-base/resolve/main/config.json""",
    # See all ViT MSN models at https://huggingface.co/models?filter=vit_msn
}
class 			UpperCamelCase_		(  _lowerCamelCase					):
					lowerCAmelCase_							=	'''vit_msn'''




					def __init__(		self				,						lowerCAmelCase_=768				,						lowerCAmelCase_=12				,						lowerCAmelCase_=12				,						lowerCAmelCase_=3072				,						lowerCAmelCase_="gelu"				,						lowerCAmelCase_=0.0				,						lowerCAmelCase_=0.0				,						lowerCAmelCase_=0.02				,						lowerCAmelCase_=1E-06				,						lowerCAmelCase_=224				,						lowerCAmelCase_=16				,						lowerCAmelCase_=3				,						lowerCAmelCase_=True				,						**lowerCAmelCase_				,						)		->							str:
										super().__init__(**lowerCAmelCase_  )

										_snake_case									=				hidden_size
										_snake_case									=				num_hidden_layers
										_snake_case									=				num_attention_heads
										_snake_case									=				intermediate_size
										_snake_case									=				hidden_act
										_snake_case									=				hidden_dropout_prob
										_snake_case									=				attention_probs_dropout_prob
										_snake_case									=				initializer_range
										_snake_case									=				layer_norm_eps
										_snake_case									=				image_size
										_snake_case									=				patch_size
										_snake_case									=				num_channels
										_snake_case									=				qkv_bias



