



"""simple docstring"""






import pickle

import numpy as np
from matplotlib import pyplot as plt



class     A__     :




	'''simple docstring'''





	def __init__(  self:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	List[str] ,    _SCREAMING_SNAKE_CASE:	Dict ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	List[Any] ,    _SCREAMING_SNAKE_CASE:	List[str] ,    _SCREAMING_SNAKE_CASE:	Dict=0.2 ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=0.2)       -> int:

				"""simple docstring"""

				__lowerCAmelCase   :						Union[str, Any]				= bp_numa
				__lowerCAmelCase   :						Tuple				= bp_numa
				__lowerCAmelCase   :						Optional[Any]				= bp_numa
				__lowerCAmelCase   :						str				= conva_get[:2]
				__lowerCAmelCase   :						Union[str, Any]				= conva_get[2]
				__lowerCAmelCase   :						int				= size_pa
				__lowerCAmelCase   :						Optional[Any]				= rate_w
				__lowerCAmelCase   :						Any				= rate_t
				__lowerCAmelCase   :						Optional[Any]				= [
				    np.mat(-1 * np.random.rand(self.conva[0] ,    self.conva[0]) + 0.5)
				    for i in range(self.conva[1])
				]
				__lowerCAmelCase   :						Optional[Any]				= np.mat(-1 * np.random.rand(self.num_bpa ,    self.num_bpa) + 0.5)
				__lowerCAmelCase   :						List[Any]				= np.mat(-1 * np.random.rand(self.num_bpa ,    self.num_bpa) + 0.5)
				__lowerCAmelCase   :						Tuple				= -2 * np.random.rand(self.conva[1]) + 1
				__lowerCAmelCase   :						Tuple				= -2 * np.random.rand(self.num_bpa) + 1
				__lowerCAmelCase   :						str				= -2 * np.random.rand(self.num_bpa) + 1

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any] ,    _SCREAMING_SNAKE_CASE:	Dict)       -> List[str]:

				"""simple docstring"""

				__lowerCAmelCase   :						Any				= {
				    "num_bp1": self.num_bpa,
				    "num_bp2": self.num_bpa,
				    "num_bp3": self.num_bpa,
				    "conv1": self.conva,
				    "step_conv1": self.step_conva,
				    "size_pooling1": self.size_poolinga,
				    "rate_weight": self.rate_weight,
				    "rate_thre": self.rate_thre,
				    "w_conv1": self.w_conva,
				    "wkj": self.wkj,
				    "vji": self.vji,
				    "thre_conv1": self.thre_conva,
				    "thre_bp2": self.thre_bpa,
				    "thre_bp3": self.thre_bpa,
				}
				with open(_SCREAMING_SNAKE_CASE ,    "wb") as f:
							pickle.dump(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)

				print(F"""Model saved： {save_path}""")

	@classmethod
	def 							_SCREAMING_SNAKE_CASE     (  cls:	int ,    _SCREAMING_SNAKE_CASE:	Any)       -> Tuple:

				"""simple docstring"""

				with open(_SCREAMING_SNAKE_CASE ,    "rb") as f:
							__lowerCAmelCase   :						Any				= pickle.load(_SCREAMING_SNAKE_CASE)  # noqa: S301

				__lowerCAmelCase   :						Optional[int]				= model_dic.get("conv1")
				conv_get.append(model_dic.get("step_conv1"))
				__lowerCAmelCase   :						Any				= model_dic.get("size_pooling1")
				__lowerCAmelCase   :						List[Any]				= model_dic.get("num_bp1")
				__lowerCAmelCase   :						List[Any]				= model_dic.get("num_bp2")
				__lowerCAmelCase   :						Optional[int]				= model_dic.get("num_bp3")
				__lowerCAmelCase   :						Any				= model_dic.get("rate_weight")
				__lowerCAmelCase   :						int				= model_dic.get("rate_thre")
				# create model instance
				__lowerCAmelCase   :						Any				= CNN(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
				# modify model parameter
				__lowerCAmelCase   :						Tuple				= model_dic.get("w_conv1")
				__lowerCAmelCase   :						Union[str, Any]				= model_dic.get("wkj")
				__lowerCAmelCase   :						Union[str, Any]				= model_dic.get("vji")
				__lowerCAmelCase   :						Any				= model_dic.get("thre_conv1")
				__lowerCAmelCase   :						List[Any]				= model_dic.get("thre_bp2")
				__lowerCAmelCase   :						Union[str, Any]				= model_dic.get("thre_bp3")
				return conv_ins

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

				"""simple docstring"""

				return 1 / (1 + np.exp(-1 * x))

	def 							_SCREAMING_SNAKE_CASE     (  self:	Tuple ,    _SCREAMING_SNAKE_CASE:	Dict)       -> Optional[int]:

				"""simple docstring"""

				return round(_SCREAMING_SNAKE_CASE ,    3)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Any ,    _SCREAMING_SNAKE_CASE:	Any ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	List[str] ,    _SCREAMING_SNAKE_CASE:	List[Any] ,    _SCREAMING_SNAKE_CASE:	Any)       -> Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase   :						Optional[int]				= convs[0]
				__lowerCAmelCase   :						int				= convs[1]
				__lowerCAmelCase   :						Union[str, Any]				= np.shape(_SCREAMING_SNAKE_CASE)[0]
				# get the data slice of original image data, data_focus
				__lowerCAmelCase   :						List[str]				= []
				for i_focus in range(0 ,    size_data - size_conv + 1 ,    _SCREAMING_SNAKE_CASE):
							for j_focus in range(0 ,    size_data - size_conv + 1 ,    _SCREAMING_SNAKE_CASE):
										__lowerCAmelCase   :						Dict				= data[
										    i_focus : i_focus + size_conv, j_focus : j_focus + size_conv
										]
										data_focus.append(_SCREAMING_SNAKE_CASE)
        # calculate the feature map of every single kernel, and saved as list of matrix
				__lowerCAmelCase   :						List[str]				= []
				__lowerCAmelCase   :						Any				= int((size_data - size_conv) / conv_step + 1)
				for i_map in range(_SCREAMING_SNAKE_CASE):
							__lowerCAmelCase   :						Optional[int]				= []
							for i_focus in range(len(_SCREAMING_SNAKE_CASE)):
										__lowerCAmelCase   :						Optional[int]				= (
										    np.sum(np.multiply(data_focus[i_focus] ,    w_convs[i_map]))
										    - thre_convs[i_map]
										)
										featuremap.append(self.sig(_SCREAMING_SNAKE_CASE))
							__lowerCAmelCase   :						Dict				= np.asmatrix(_SCREAMING_SNAKE_CASE).reshape(
							    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
							data_featuremap.append(_SCREAMING_SNAKE_CASE)

				# expanding the data slice to One dimenssion
				__lowerCAmelCase   :						Optional[Any]				= []
				for each_focus in data_focus:
							focusa_list.extend(self.Expand_Mat(_SCREAMING_SNAKE_CASE))
				__lowerCAmelCase   :						List[Any]				= np.asarray(_SCREAMING_SNAKE_CASE)
				return focus_list, data_featuremap

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any] ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	Tuple="average_pool")       -> Union[str, Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						int				= len(featuremaps[0])
				__lowerCAmelCase   :						List[str]				= int(size_map / size_pooling)
				__lowerCAmelCase   :						List[Any]				= []
				for i_map in range(len(_SCREAMING_SNAKE_CASE)):
							__lowerCAmelCase   :						List[str]				= featuremaps[i_map]
							__lowerCAmelCase   :						List[Any]				= []
							for i_focus in range(0 ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE):
										for j_focus in range(0 ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE):
													__lowerCAmelCase   :						Tuple				= feature_map[
													    i_focus : i_focus + size_pooling,
													    j_focus : j_focus + size_pooling,
													]
													if pooling_type == "average_pool":
																# average pooling
																map_pooled.append(np.average(_SCREAMING_SNAKE_CASE))
													elif pooling_type == "max_pooling":
																# max pooling
																map_pooled.append(np.max(_SCREAMING_SNAKE_CASE))
							__lowerCAmelCase   :						Tuple				= np.asmatrix(_SCREAMING_SNAKE_CASE).reshape(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
							featuremap_pooled.append(_SCREAMING_SNAKE_CASE)
				return featuremap_pooled

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	Optional[int])       -> Dict:

				"""simple docstring"""

				__lowerCAmelCase   :						Any				= []
				for i in range(len(_SCREAMING_SNAKE_CASE)):
							__lowerCAmelCase   :						List[str]				= np.shape(data[i])
							__lowerCAmelCase   :						List[str]				= data[i].reshape(1 ,    shapes[0] * shapes[1])
							__lowerCAmelCase   :						Optional[Any]				= data_listed.getA().tolist()[0]
							data_expanded.extend(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Tuple				= np.asarray(_SCREAMING_SNAKE_CASE)
				return data_expanded

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

				"""simple docstring"""

				__lowerCAmelCase   :						List[Any]				= np.asarray(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Optional[Any]				= np.shape(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						List[Any]				= data_mat.reshape(1 ,    shapes[0] * shapes[1])
				return data_expanded

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

				"""simple docstring"""

				__lowerCAmelCase   :						List[Any]				= []
				__lowerCAmelCase   :						Optional[int]				= 0
				for i_map in range(_SCREAMING_SNAKE_CASE):
							__lowerCAmelCase   :						Optional[int]				= np.ones((size_map, size_map))
							for i in range(0 ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE):
										for j in range(0 ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE):
													__lowerCAmelCase   :						List[str]				= pd_pool[
													    i_pool
													]
													__lowerCAmelCase   :						Tuple				= i_pool + 1
							__lowerCAmelCase   :						str				= np.multiply(
							    _SCREAMING_SNAKE_CASE ,    np.multiply(out_map[i_map] ,    (1 - out_map[i_map])))
							pd_all.append(_SCREAMING_SNAKE_CASE)
				return pd_all

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict ,    _SCREAMING_SNAKE_CASE:	Any ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	Union[str, Any] ,    _SCREAMING_SNAKE_CASE:	int ,    _SCREAMING_SNAKE_CASE:	Any=bool)       -> Any:

				"""simple docstring"""

				print("----------------------Start Training-------------------------")
				print((" - - Shape: Train_Data  ", np.shape(_SCREAMING_SNAKE_CASE)))
				print((" - - Shape: Teach_Data  ", np.shape(_SCREAMING_SNAKE_CASE)))
				__lowerCAmelCase   :						Dict				= 0
				__lowerCAmelCase   :						int				= []
				__lowerCAmelCase   :						Union[str, Any]				= 1_0000
				while rp < n_repeat and mse >= error_accuracy:
							__lowerCAmelCase   :						Tuple				= 0
							print(F"""-------------Learning Time {rp}--------------""")
							for p in range(len(_SCREAMING_SNAKE_CASE)):
										# print('------------Learning Image: %d--------------'%p)
										__lowerCAmelCase   :						Optional[Any]				= np.asmatrix(datas_train[p])
										__lowerCAmelCase   :						Dict				= np.asarray(datas_teach[p])
										__lowerCAmelCase		,			__lowerCAmelCase   :						int				= self.convolute(
										    _SCREAMING_SNAKE_CASE ,    self.conva ,    self.w_conva ,    self.thre_conva ,    conv_step=self.step_conva ,    )
										__lowerCAmelCase   :						Optional[Any]				= self.pooling(_SCREAMING_SNAKE_CASE ,    self.size_poolinga)
										__lowerCAmelCase   :						List[Any]				= np.shape(_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase   :						Union[str, Any]				= self._expand(_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase   :						List[Any]				= data_bp_input

										__lowerCAmelCase   :						List[str]				= np.dot(_SCREAMING_SNAKE_CASE ,    self.vji.T) - self.thre_bpa
										__lowerCAmelCase   :						int				= self.sig(_SCREAMING_SNAKE_CASE)
										__lowerCAmelCase   :						List[str]				= np.dot(_SCREAMING_SNAKE_CASE ,    self.wkj.T) - self.thre_bpa
										__lowerCAmelCase   :						Union[str, Any]				= self.sig(_SCREAMING_SNAKE_CASE)

										# --------------Model Leaning ------------------------
										# calculate error and gradient---------------
										__lowerCAmelCase   :						Dict				= np.multiply(
										    (data_teach - bp_outa) ,    np.multiply(_SCREAMING_SNAKE_CASE ,    (1 - bp_outa)))
										__lowerCAmelCase   :						List[Any]				= np.multiply(
										    np.dot(_SCREAMING_SNAKE_CASE ,    self.wkj) ,    np.multiply(_SCREAMING_SNAKE_CASE ,    (1 - bp_outa)))
										__lowerCAmelCase   :						Optional[int]				= np.dot(_SCREAMING_SNAKE_CASE ,    self.vji)

										__lowerCAmelCase   :						Dict				= pd_i_all / (self.size_poolinga * self.size_poolinga)
										__lowerCAmelCase   :						Union[str, Any]				= pd_conva_pooled.T.getA().tolist()
										__lowerCAmelCase   :						Tuple				= self._calculate_gradient_from_pool(
										    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    shape_featuremapa[0] ,    shape_featuremapa[1] ,    self.size_poolinga ,    )
										# weight and threshold learning process---------
										# convolution layer
										for k_conv in range(self.conva[1]):
													__lowerCAmelCase   :						List[Any]				= self._expand_mat(pd_conva_all[k_conv])
													__lowerCAmelCase   :						List[str]				= self.rate_weight * np.dot(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)

													__lowerCAmelCase   :						Optional[int]				= self.w_conva[k_conv] + delta_w.reshape(
													    (self.conva[0], self.conva[0]))

													__lowerCAmelCase   :						List[str]				= (
													    self.thre_conva[k_conv]
													    - np.sum(pd_conva_all[k_conv]) * self.rate_thre
													)
										# all connected layer
										__lowerCAmelCase   :						Any				= self.wkj + pd_k_all.T * bp_outa * self.rate_weight
										__lowerCAmelCase   :						Any				= self.vji + pd_j_all.T * bp_outa * self.rate_weight
										__lowerCAmelCase   :						List[str]				= self.thre_bpa - pd_k_all * self.rate_thre
										__lowerCAmelCase   :						Tuple				= self.thre_bpa - pd_j_all * self.rate_thre
										# calculate the sum error of all single image
										__lowerCAmelCase   :						int				= np.sum(abs(data_teach - bp_outa))
										error_count += errors
										# print('   ----Teach      ',data_teach)
										# print('   ----BP_output  ',bp_out3)
							__lowerCAmelCase   :						Union[str, Any]				= rp + 1
							__lowerCAmelCase   :						Any				= error_count / patterns
							all_mse.append(_SCREAMING_SNAKE_CASE)

				def draw_error():
							__lowerCAmelCase   :						Optional[int]				= [error_accuracy for i in range(int(n_repeat * 1.2))]
							plt.plot(_SCREAMING_SNAKE_CASE ,    "+-")
							plt.plot(_SCREAMING_SNAKE_CASE ,    "r--")
							plt.xlabel("Learning Times")
							plt.ylabel("All_mse")
							plt.grid(_SCREAMING_SNAKE_CASE ,    alpha=0.5)
							plt.show()

				print("------------------Training Complished---------------------")
				print((" - - Training epoch: ", rp, F"""     - - Mse: {mse:.6f}"""))
				if draw_e:
							draw_error()
				return mse

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[int] ,    _SCREAMING_SNAKE_CASE:	List[str])       -> Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase   :						List[Any]				= []
				print("-------------------Start Testing-------------------------")
				print((" - - Shape: Test_Data  ", np.shape(_SCREAMING_SNAKE_CASE)))
				for p in range(len(_SCREAMING_SNAKE_CASE)):
							__lowerCAmelCase   :						Optional[Any]				= np.asmatrix(datas_test[p])
							__lowerCAmelCase		,			__lowerCAmelCase   :						List[str]				= self.convolute(
							    _SCREAMING_SNAKE_CASE ,    self.conva ,    self.w_conva ,    self.thre_conva ,    conv_step=self.step_conva ,    )
							__lowerCAmelCase   :						int				= self.pooling(_SCREAMING_SNAKE_CASE ,    self.size_poolinga)
							__lowerCAmelCase   :						Optional[Any]				= self._expand(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase   :						int				= data_bp_input
							__lowerCAmelCase   :						Optional[Any]				= bp_outa * self.vji.T - self.thre_bpa
							__lowerCAmelCase   :						Any				= self.sig(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase   :						Tuple				= bp_outa * self.wkj.T - self.thre_bpa
							__lowerCAmelCase   :						Optional[Any]				= self.sig(_SCREAMING_SNAKE_CASE)
							produce_out.extend(bp_outa.getA().tolist())
				__lowerCAmelCase   :						Optional[Any]				= [list(map(self.do_round ,    _SCREAMING_SNAKE_CASE)) for each in produce_out]
				return np.asarray(_SCREAMING_SNAKE_CASE)







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

				"""simple docstring"""

				__lowerCAmelCase   :						Optional[int]				= np.asmatrix(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase		,			__lowerCAmelCase   :						List[Any]				= self.convolute(
				    _SCREAMING_SNAKE_CASE ,    self.conva ,    self.w_conva ,    self.thre_conva ,    conv_step=self.step_conva ,    )
				__lowerCAmelCase   :						Optional[int]				= self.pooling(_SCREAMING_SNAKE_CASE ,    self.size_poolinga)

				return data_conveda, data_pooleda


if __name__ == "__main__":
	pass




"""simple docstring"""






import torch
from torch import nn



class     A__     (						nn.Module							):




	'''simple docstring'''





	def __init__(  self:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	List[Any] ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	List[str] ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=1 ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=False)       -> Dict:

				"""simple docstring"""

				super().__init__()

				__lowerCAmelCase   :						Optional[int]				= n_token
				__lowerCAmelCase   :						int				= d_embed
				__lowerCAmelCase   :						str				= d_proj

				__lowerCAmelCase   :						Optional[Any]				= cutoffs + [n_token]
				__lowerCAmelCase   :						Any				= [0] + self.cutoffs
				__lowerCAmelCase   :						Tuple				= div_val

				__lowerCAmelCase   :						Optional[Any]				= self.cutoffs[0]
				__lowerCAmelCase   :						Union[str, Any]				= len(self.cutoffs) - 1
				__lowerCAmelCase   :						List[str]				= self.shortlist_size + self.n_clusters

				if self.n_clusters > 0:
							__lowerCAmelCase   :						Any				= nn.Parameter(torch.zeros(self.n_clusters ,    self.d_embed))
							__lowerCAmelCase   :						Dict				= nn.Parameter(torch.zeros(self.n_clusters))

				__lowerCAmelCase   :						List[Any]				= nn.ModuleList()
				__lowerCAmelCase   :						Any				= nn.ParameterList()

				if div_val == 1:
							for i in range(len(self.cutoffs)):
										if d_proj != d_embed:
													self.out_projs.append(nn.Parameter(torch.FloatTensor(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)))
										else:
													self.out_projs.append(_SCREAMING_SNAKE_CASE)

							self.out_layers.append(nn.Linear(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE))
				else:
							for i in range(len(self.cutoffs)):
										__lowerCAmelCase		,			__lowerCAmelCase   :						Tuple				= self.cutoff_ends[i], self.cutoff_ends[i + 1]
										__lowerCAmelCase   :						str				= d_embed // (div_val**i)

										self.out_projs.append(nn.Parameter(torch.FloatTensor(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)))

										self.out_layers.append(nn.Linear(_SCREAMING_SNAKE_CASE ,    r_idx - l_idx))

				__lowerCAmelCase   :						str				= keep_order

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	Any ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	int ,    _SCREAMING_SNAKE_CASE:	Dict)       -> int:

				"""simple docstring"""

				if proj is None:
							__lowerCAmelCase   :						Dict				= nn.functional.linear(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    bias=_SCREAMING_SNAKE_CASE)
				else:
							# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
							__lowerCAmelCase   :						Optional[Any]				= nn.functional.linear(_SCREAMING_SNAKE_CASE ,    proj.t().contiguous())
							__lowerCAmelCase   :						List[Any]				= nn.functional.linear(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    bias=_SCREAMING_SNAKE_CASE)
							# else:
							#     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
							#     if bias is not None:
							#         logit = logit + bias

				return logit

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	Tuple=None ,    _SCREAMING_SNAKE_CASE:	int=False)       -> List[str]:

				"""simple docstring"""

				if labels is not None:
							# Shift so that tokens < n predict n
							__lowerCAmelCase   :						List[Any]				= hidden[..., :-1, :].contiguous()
							__lowerCAmelCase   :						Tuple				= labels[..., 1:].contiguous()
							__lowerCAmelCase   :						Union[str, Any]				= hidden.view(-1 ,    hidden.size(-1))
							__lowerCAmelCase   :						Optional[int]				= labels.view(-1)
							if hidden.size(0) != labels.size(0):
										raise RuntimeError("Input and labels should have the same size in the batch dimension.")
				else:
							__lowerCAmelCase   :						List[Any]				= hidden.view(-1 ,    hidden.size(-1))

				if self.n_clusters == 0:
							__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    self.out_layers[0].weight ,    self.out_layers[0].bias ,    self.out_projs[0])
							if labels is not None:
										__lowerCAmelCase   :						Any				= labels != -100
										__lowerCAmelCase   :						Tuple				= torch.zeros_like(_SCREAMING_SNAKE_CASE ,    dtype=hidden.dtype ,    device=hidden.device)
										__lowerCAmelCase   :						List[Any]				= (
										    -nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=-1)[mask].gather(1 ,    labels[mask].unsqueeze(1)).squeeze(1)
										)
							else:
										__lowerCAmelCase   :						List[Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=-1)
				else:
							# construct weights and biases
							__lowerCAmelCase		,			__lowerCAmelCase   :						Optional[Any]				= [], []
							for i in range(len(self.cutoffs)):
										if self.div_val == 1:
													__lowerCAmelCase		,			__lowerCAmelCase   :						int				= self.cutoff_ends[i], self.cutoff_ends[i + 1]
													__lowerCAmelCase   :						Optional[Any]				= self.out_layers[0].weight[l_idx:r_idx]
													__lowerCAmelCase   :						Optional[int]				= self.out_layers[0].bias[l_idx:r_idx]
										else:
													__lowerCAmelCase   :						int				= self.out_layers[i].weight
													__lowerCAmelCase   :						List[Any]				= self.out_layers[i].bias

										if i == 0:
													__lowerCAmelCase   :						int				= torch.cat([weight_i, self.cluster_weight] ,    dim=0)
													__lowerCAmelCase   :						Union[str, Any]				= torch.cat([bias_i, self.cluster_bias] ,    dim=0)

										weights.append(_SCREAMING_SNAKE_CASE)
										biases.append(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						Dict				= weights[0], biases[0], self.out_projs[0]

							__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
							__lowerCAmelCase   :						Optional[Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)

							if labels is None:
										__lowerCAmelCase   :						Optional[Any]				= hidden.new_empty((head_logit.size(0), self.n_token))
							else:
										__lowerCAmelCase   :						int				= torch.zeros_like(_SCREAMING_SNAKE_CASE ,    dtype=hidden.dtype ,    device=hidden.device)

							__lowerCAmelCase   :						str				= 0
							__lowerCAmelCase   :						List[str]				= [0] + self.cutoffs
							for i in range(len(_SCREAMING_SNAKE_CASE) - 1):
										__lowerCAmelCase		,			__lowerCAmelCase   :						List[Any]				= cutoff_values[i], cutoff_values[i + 1]

										if labels is not None:
													__lowerCAmelCase   :						List[str]				= (labels >= l_idx) & (labels < r_idx)
													__lowerCAmelCase   :						Any				= mask_i.nonzero().squeeze()

													if indices_i.numel() == 0:
																continue

													__lowerCAmelCase   :						Optional[Any]				= labels.index_select(0 ,    _SCREAMING_SNAKE_CASE) - l_idx
													__lowerCAmelCase   :						List[Any]				= head_logprob.index_select(0 ,    _SCREAMING_SNAKE_CASE)
													__lowerCAmelCase   :						Tuple				= hidden.index_select(0 ,    _SCREAMING_SNAKE_CASE)
										else:
													__lowerCAmelCase   :						List[str]				= hidden

										if i == 0:
													if labels is not None:
																__lowerCAmelCase   :						Optional[int]				= head_logprob_i.gather(1 ,    target_i[:, None]).squeeze(1)
													else:
																__lowerCAmelCase   :						Optional[Any]				= head_logprob[:, : self.cutoffs[0]]
										else:
													__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						List[str]				= weights[i], biases[i], self.out_projs[i]

													__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
													__lowerCAmelCase   :						Tuple				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)
													__lowerCAmelCase   :						int				= self.cutoffs[0] + i - 1  # No probability for the head cluster
													if labels is not None:
																__lowerCAmelCase   :						List[str]				= head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
																    1 ,    target_i[:, None]).squeeze(1)
													else:
																__lowerCAmelCase   :						Optional[int]				= head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
																__lowerCAmelCase   :						Optional[int]				= logprob_i

										if labels is not None:
													if (hasattr(self ,    "keep_order") and self.keep_order) or keep_order:
																out.index_copy_(0 ,    _SCREAMING_SNAKE_CASE ,    -logprob_i)
													else:
																out[offset : offset + logprob_i.size(0)].copy_(-logprob_i)
													offset += logprob_i.size(0)

				return out







	def 							_SCREAMING_SNAKE_CASE     (  self:	str ,    _SCREAMING_SNAKE_CASE:	List[Any])       -> Optional[Any]:

				"""simple docstring"""

				if self.n_clusters == 0:
							__lowerCAmelCase   :						int				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    self.out_layers[0].weight ,    self.out_layers[0].bias ,    self.out_projs[0])
							return nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=-1)
				else:
							# construct weights and biases
							__lowerCAmelCase		,			__lowerCAmelCase   :						Dict				= [], []
							for i in range(len(self.cutoffs)):
										if self.div_val == 1:
													__lowerCAmelCase		,			__lowerCAmelCase   :						int				= self.cutoff_ends[i], self.cutoff_ends[i + 1]
													__lowerCAmelCase   :						Optional[Any]				= self.out_layers[0].weight[l_idx:r_idx]
													__lowerCAmelCase   :						Union[str, Any]				= self.out_layers[0].bias[l_idx:r_idx]
										else:
													__lowerCAmelCase   :						Dict				= self.out_layers[i].weight
													__lowerCAmelCase   :						List[str]				= self.out_layers[i].bias

										if i == 0:
													__lowerCAmelCase   :						int				= torch.cat([weight_i, self.cluster_weight] ,    dim=0)
													__lowerCAmelCase   :						int				= torch.cat([bias_i, self.cluster_bias] ,    dim=0)

										weights.append(_SCREAMING_SNAKE_CASE)
										biases.append(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						List[Any]				= weights[0], biases[0], self.out_projs[0]
							__lowerCAmelCase   :						str				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)

							__lowerCAmelCase   :						Optional[Any]				= hidden.new_empty((head_logit.size(0), self.n_token))
							__lowerCAmelCase   :						Union[str, Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)

							__lowerCAmelCase   :						Dict				= [0] + self.cutoffs
							for i in range(len(_SCREAMING_SNAKE_CASE) - 1):
										__lowerCAmelCase		,			__lowerCAmelCase   :						Optional[Any]				= cutoff_values[i], cutoff_values[i + 1]

										if i == 0:
													__lowerCAmelCase   :						Union[str, Any]				= head_logprob[:, : self.cutoffs[0]]
										else:
													__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						Dict				= weights[i], biases[i], self.out_projs[i]

													__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
													__lowerCAmelCase   :						List[Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)

													__lowerCAmelCase   :						int				= head_logprob[:, -i] + tail_logprob_i
													__lowerCAmelCase   :						Tuple				= logprob_i

							return out