






import torch
from torch import nn
class     lowerCamelCase_					(   nn.Module			):







				def __init__(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase=1							,       __lowerCAmelCase=False  ):
								"""simple docstring"""
								super().__init__()

								__magic_name__		:Union[str, Any]   =   n_token
								__magic_name__		:Union[str, Any]   =   d_embed
								__magic_name__		:int   =   d_proj

								__magic_name__		:List[Any]   =   cutoffs + [n_token]
								__magic_name__		:str   =   [0] + self.cutoffs
								__magic_name__		:int   =   div_val

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

								if self.n_clusters > 0:
												__magic_name__		:str   =   nn.Parameter(torch.zeros(self.n_clusters							,       self.d_embed  )  )
												__magic_name__		:Tuple   =   nn.Parameter(torch.zeros(self.n_clusters  )  )

								__magic_name__		:Union[str, Any]   =   nn.ModuleList()
								__magic_name__		: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(__lowerCAmelCase							,       __lowerCAmelCase  )  )  )
																else:
																				self.out_projs.append(__lowerCAmelCase  )

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

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

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

								__magic_name__		:List[str]   =   keep_order







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  ):
								"""simple docstring"""
								if proj is None:
												__magic_name__		:Any   =   nn.functional.linear(__lowerCAmelCase							,       __lowerCAmelCase							,       bias=__lowerCAmelCase  )
								else:
												# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
												__magic_name__		:Any   =   nn.functional.linear(__lowerCAmelCase							,       proj.t().contiguous()  )
												__magic_name__		:str   =   nn.functional.linear(__lowerCAmelCase							,       __lowerCAmelCase							,       bias=__lowerCAmelCase  )
												# else:
												#     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
												#     if bias is not None:
												#         logit = logit + bias

								return logit







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase=None							,       __lowerCAmelCase=False  ):
								"""simple docstring"""

								if labels is not None:
												# Shift so that tokens < n predict n
												__magic_name__		:List[str]   =   hidden[..., :-1, :].contiguous()
												__magic_name__		:Dict   =   labels[..., 1:].contiguous()
												__magic_name__		:Any   =   hidden.view(-1							,       hidden.size(-1  )  )
												__magic_name__		: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:
												__magic_name__		:Union[str, Any]   =   hidden.view(-1							,       hidden.size(-1  )  )

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

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

																weights.append(__lowerCAmelCase  )
																biases.append(__lowerCAmelCase  )

												__magic_name__  ,					__magic_name__  ,					__magic_name__		:int   =   weights[0], biases[0], self.out_projs[0]

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

												if labels is None:
																__magic_name__		:str   =   hidden.new_empty((head_logit.size(0  ), self.n_token)  )
												else:
																__magic_name__		:int   =   torch.zeros_like(__lowerCAmelCase							,       dtype=hidden.dtype							,       device=hidden.device  )

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

																if labels is not None:
																				__magic_name__		:Tuple   =   (labels >= l_idx) & (labels < r_idx)
																				__magic_name__		:Optional[Any]   =   mask_i.nonzero().squeeze()

																				if indices_i.numel() == 0:
																								continue

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

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

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

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

								return out







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

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

																weights.append(__lowerCAmelCase  )
																biases.append(__lowerCAmelCase  )

												__magic_name__  ,					__magic_name__  ,					__magic_name__		:str   =   weights[0], biases[0], self.out_projs[0]
												__magic_name__		:Dict   =   self._compute_logit(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  )

												__magic_name__		:Optional[int]   =   hidden.new_empty((head_logit.size(0  ), self.n_token)  )
												__magic_name__		:Tuple   =   nn.functional.log_softmax(__lowerCAmelCase							,       dim=1  )

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

																if i == 0:
																				__magic_name__		:Tuple   =   head_logprob[:, : self.cutoffs[0]]
																else:
																				__magic_name__  ,					__magic_name__  ,					__magic_name__		:Any   =   weights[i], biases[i], self.out_projs[i]

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

																				__magic_name__		:Any   =   head_logprob[:, -i] + tail_logprob_i
																				__magic_name__		:Union[str, Any]   =   logprob_i

												return out














import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
					import jax.numpy as jnp

if is_tf_available():
					import tensorflow as tf

if is_torch_available():
					import torch
class     lowerCamelCase_					(   unittest.TestCase			):







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   {
								    '''task_specific_params''': {
								        '''summarization''': {'''length_penalty''': 1.0, '''max_length''': 1_2_8, '''min_length''': 1_2, '''num_beams''': 4},
								        '''summarization_cnn''': {'''length_penalty''': 2.0, '''max_length''': 1_4_2, '''min_length''': 5_6, '''num_beams''': 4},
								        '''summarization_xsum''': {'''length_penalty''': 1.0, '''max_length''': 6_2, '''min_length''': 1_1, '''num_beams''': 6},
								    }
								}
								__magic_name__		:List[str]   =   {
								    '''task_specific_params.summarization.length_penalty''': 1.0,
								    '''task_specific_params.summarization.max_length''': 1_2_8,
								    '''task_specific_params.summarization.min_length''': 1_2,
								    '''task_specific_params.summarization.num_beams''': 4,
								    '''task_specific_params.summarization_cnn.length_penalty''': 2.0,
								    '''task_specific_params.summarization_cnn.max_length''': 1_4_2,
								    '''task_specific_params.summarization_cnn.min_length''': 5_6,
								    '''task_specific_params.summarization_cnn.num_beams''': 4,
								    '''task_specific_params.summarization_xsum.length_penalty''': 1.0,
								    '''task_specific_params.summarization_xsum.max_length''': 6_2,
								    '''task_specific_params.summarization_xsum.min_length''': 1_1,
								    '''task_specific_params.summarization_xsum.num_beams''': 6,
								}

								self.assertEqual(flatten_dict(__lowerCAmelCase  )							,       __lowerCAmelCase  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]   =   np.random.randn(3							,       4  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase  )							,       x.transpose()  )  )

								__magic_name__		:List[Any]   =   np.random.randn(3							,       4							,       5  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )							,       x.transpose((1, 2, 0)  )  )  )







				@require_torch
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]   =   np.random.randn(3							,       4  )
								__magic_name__		:Tuple   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase  )							,       transpose(__lowerCAmelCase  ).numpy()  )  )

								__magic_name__		:int   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:Union[str, Any]   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )							,       transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  ).numpy()  )  )







				@require_tf
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:int   =   np.random.randn(3							,       4  )
								__magic_name__		:Optional[Any]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase  )							,       transpose(__lowerCAmelCase  ).numpy()  )  )

								__magic_name__		:List[str]   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:Optional[Any]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )							,       transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  ).numpy()  )  )







				@require_flax
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:int   =   np.random.randn(3							,       4  )
								__magic_name__		:Dict   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase  )							,       np.asarray(transpose(__lowerCAmelCase  )  )  )  )

								__magic_name__		:Dict   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:Dict   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )							,       np.asarray(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )  )  )  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Any   =   np.random.randn(3							,       4  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (4, 3)  )							,       np.reshape(__lowerCAmelCase							,       (4, 3)  )  )  )

								__magic_name__		:Union[str, Any]   =   np.random.randn(3							,       4							,       5  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (1_2, 5)  )							,       np.reshape(__lowerCAmelCase							,       (1_2, 5)  )  )  )







				@require_torch
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   np.random.randn(3							,       4  )
								__magic_name__		:Tuple   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (4, 3)  )							,       reshape(__lowerCAmelCase							,       (4, 3)  ).numpy()  )  )

								__magic_name__		:Union[str, Any]   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:List[str]   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (1_2, 5)  )							,       reshape(__lowerCAmelCase							,       (1_2, 5)  ).numpy()  )  )







				@require_tf
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   np.random.randn(3							,       4  )
								__magic_name__		:Union[str, Any]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (4, 3)  )							,       reshape(__lowerCAmelCase							,       (4, 3)  ).numpy()  )  )

								__magic_name__		:List[Any]   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:Optional[int]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (1_2, 5)  )							,       reshape(__lowerCAmelCase							,       (1_2, 5)  ).numpy()  )  )







				@require_flax
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[str]   =   np.random.randn(3							,       4  )
								__magic_name__		:Any   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (4, 3)  )							,       np.asarray(reshape(__lowerCAmelCase							,       (4, 3)  )  )  )  )

								__magic_name__		:List[Any]   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:List[str]   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (1_2, 5)  )							,       np.asarray(reshape(__lowerCAmelCase							,       (1_2, 5)  )  )  )  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   np.random.randn(1							,       3							,       4  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase  )							,       np.squeeze(__lowerCAmelCase  )  )  )

								__magic_name__		:Optional[Any]   =   np.random.randn(1							,       4							,       1							,       5  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,       axis=2  )							,       np.squeeze(__lowerCAmelCase							,       axis=2  )  )  )







				@require_torch
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   np.random.randn(1							,       3							,       4  )
								__magic_name__		:List[Any]   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase  )							,       squeeze(__lowerCAmelCase  ).numpy()  )  )

								__magic_name__		:List[str]   =   np.random.randn(1							,       4							,       1							,       5  )
								__magic_name__		:str   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,       axis=2  )							,       squeeze(__lowerCAmelCase							,       axis=2  ).numpy()  )  )







				@require_tf
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:int   =   np.random.randn(1							,       3							,       4  )
								__magic_name__		:Tuple   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase  )							,       squeeze(__lowerCAmelCase  ).numpy()  )  )

								__magic_name__		:Tuple   =   np.random.randn(1							,       4							,       1							,       5  )
								__magic_name__		:Optional[int]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,       axis=2  )							,       squeeze(__lowerCAmelCase							,       axis=2  ).numpy()  )  )







				@require_flax
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   np.random.randn(1							,       3							,       4  )
								__magic_name__		:Optional[Any]   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase  )							,       np.asarray(squeeze(__lowerCAmelCase  )  )  )  )

								__magic_name__		:List[Any]   =   np.random.randn(1							,       4							,       1							,       5  )
								__magic_name__		:Optional[Any]   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,       axis=2  )							,       np.asarray(squeeze(__lowerCAmelCase							,       axis=2  )  )  )  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Any   =   np.random.randn(3							,       4  )
								self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,       axis=1  )							,       np.expand_dims(__lowerCAmelCase							,       axis=1  )  )  )







				@require_torch
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   np.random.randn(3							,       4  )
								__magic_name__		:Any   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,       axis=1  )							,       expand_dims(__lowerCAmelCase							,       axis=1  ).numpy()  )  )







				@require_tf
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Union[str, Any]   =   np.random.randn(3							,       4  )
								__magic_name__		:Union[str, Any]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,       axis=1  )							,       expand_dims(__lowerCAmelCase							,       axis=1  ).numpy()  )  )







				@require_flax
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[str]   =   np.random.randn(3							,       4  )
								__magic_name__		:Tuple   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,       axis=1  )							,       np.asarray(expand_dims(__lowerCAmelCase							,       axis=1  )  )  )  )






