
import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from tax import checkpoints

from transformers import MTaConfig, UMTaEncoderModel, UMTaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()





def      lowerCamelCase_					(							lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__       ):
				return params[F'{prefix}/{prefix}/relpos_bias/rel_embedding'][:, i, :]





def      lowerCamelCase_					(							lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__="attention"       ):
				lowerCamelCase_         =     lowerCamelCase_         =     np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/key/kernel'][:, i, :, :]       )
				lowerCamelCase_         =     k_tmp.reshape(k_tmp.shape[0]      ,					k_tmp.shape[1] * k_tmp.shape[2]       )
				lowerCamelCase_         =     np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/out/kernel'][:, i, :, :]       )
				lowerCamelCase_         =     o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]      ,					o_tmp.shape[2]       )
				lowerCamelCase_         =     np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/query/kernel'][:, i, :, :]       )
				lowerCamelCase_         =     q_tmp.reshape(q_tmp.shape[0]      ,					q_tmp.shape[1] * q_tmp.shape[2]       )
				lowerCamelCase_         =     np.ascontiguousarray(params[F'{prefix}/{prefix}/{layer_name}/value/kernel'][:, i, :, :]       )
				lowerCamelCase_         =     v_tmp.reshape(v_tmp.shape[0]      ,					v_tmp.shape[1] * v_tmp.shape[2]       )
				return k, o, q, v





def      lowerCamelCase_					(							lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__=False       ):

				if split_mlp_wi:
								lowerCamelCase_         =     params[F'{prefix}/{prefix}/mlp/wi_0/kernel'][:, i, :]
								lowerCamelCase_         =     params[F'{prefix}/{prefix}/mlp/wi_1/kernel'][:, i, :]
								lowerCamelCase_         =     (wi_a, wi_a)
				else:
								lowerCamelCase_         =     params[F'{prefix}/{prefix}/mlp/wi/kernel'][:, i, :]

				lowerCamelCase_         =     params[F'{prefix}/{prefix}/mlp/wo/kernel'][:, i, :]
				return wi, wo





def      lowerCamelCase_					(							lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__       ):
				return params[F'{prefix}/{prefix}/{layer_name}/scale'][:, i]





def      lowerCamelCase_					(							lowerCamelCase__      ,					*, lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__ = False       ):
				lowerCamelCase_         =     traverse_util.flatten_dict(variables["target"]       )
				lowerCamelCase_         =     {"/".join(_UpperCamelCase       ): v for k, v in old.items()}

				# v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
				lowerCamelCase_         =     "encoder/encoder/mlp/wi_0/kernel" in old
				print("Split MLP:"      ,					_UpperCamelCase       )

				lowerCamelCase_         =     collections.OrderedDict()

				# Shared embeddings.
				lowerCamelCase_         =     old["token_embedder/embedding"]

				# Encoder.
				for i in range(_UpperCamelCase       ):
								# Block i, layer 0 (Self Attention).
								lowerCamelCase_         =     tax_layer_norm_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"encoder"      ,					"pre_attention_layer_norm"       )
								lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_         =     tax_attention_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"encoder"      ,					"attention"       )
								lowerCamelCase_         =     layer_norm
								lowerCamelCase_         =     k.T
								lowerCamelCase_         =     o.T
								lowerCamelCase_         =     q.T
								lowerCamelCase_         =     v.T

								# Block i, layer 1 (MLP).
								lowerCamelCase_         =     tax_layer_norm_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"encoder"      ,					"pre_mlp_layer_norm"       )
								lowerCamelCase_      ,       lowerCamelCase_         =     tax_mlp_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"encoder"      ,					_UpperCamelCase       )
								lowerCamelCase_         =     layer_norm
								if split_mlp_wi:
												lowerCamelCase_         =     wi[0].T
												lowerCamelCase_         =     wi[1].T
								else:
												lowerCamelCase_         =     wi.T
								lowerCamelCase_         =     wo.T
								if scalable_attention:
												# convert the rel_embedding of each layer
												lowerCamelCase_         =     tax_relpos_bias_lookup(
												    _UpperCamelCase      ,					_UpperCamelCase      ,					"encoder"       ).T

				lowerCamelCase_         =     old["encoder/encoder_norm/scale"]

				if not scalable_attention:
								lowerCamelCase_         =     tax_relpos_bias_lookup(
								    _UpperCamelCase      ,					0      ,					"encoder"       ).T
								lowerCamelCase_         =     tax_relpos_bias_lookup(
								    _UpperCamelCase      ,					0      ,					"decoder"       ).T

				if not is_encoder_only:
								# Decoder.
								for i in range(_UpperCamelCase       ):
												# Block i, layer 0 (Self Attention).
												lowerCamelCase_         =     tax_layer_norm_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"decoder"      ,					"pre_self_attention_layer_norm"       )
												lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_         =     tax_attention_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"decoder"      ,					"self_attention"       )
												lowerCamelCase_         =     layer_norm
												lowerCamelCase_         =     k.T
												lowerCamelCase_         =     o.T
												lowerCamelCase_         =     q.T
												lowerCamelCase_         =     v.T

												# Block i, layer 1 (Cross Attention).
												lowerCamelCase_         =     tax_layer_norm_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"decoder"      ,					"pre_cross_attention_layer_norm"       )
												lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_         =     tax_attention_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"decoder"      ,					"encoder_decoder_attention"       )
												lowerCamelCase_         =     layer_norm
												lowerCamelCase_         =     k.T
												lowerCamelCase_         =     o.T
												lowerCamelCase_         =     q.T
												lowerCamelCase_         =     v.T

												# Block i, layer 2 (MLP).
												lowerCamelCase_         =     tax_layer_norm_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"decoder"      ,					"pre_mlp_layer_norm"       )
												lowerCamelCase_      ,       lowerCamelCase_         =     tax_mlp_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"decoder"      ,					_UpperCamelCase       )
												lowerCamelCase_         =     layer_norm
												if split_mlp_wi:
																lowerCamelCase_         =     wi[0].T
																lowerCamelCase_         =     wi[1].T
												else:
																lowerCamelCase_         =     wi.T
												lowerCamelCase_         =     wo.T

												if scalable_attention:
																# convert the rel_embedding of each layer
																lowerCamelCase_         =     tax_relpos_bias_lookup(_UpperCamelCase      ,					_UpperCamelCase      ,					"decoder"       ).T

								lowerCamelCase_         =     old["decoder/decoder_norm/scale"]

								# LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
								if "decoder/logits_dense/kernel" in old:
												lowerCamelCase_         =     old["decoder/logits_dense/kernel"].T

				return new





def      lowerCamelCase_					(							lowerCamelCase__      ,					lowerCamelCase__       ):
				lowerCamelCase_         =     collections.OrderedDict([(k, torch.from_numpy(v.copy()       )) for (k, v) in converted_params.items()]       )

				# Add what is missing.
				if "encoder.embed_tokens.weight" not in state_dict:
								lowerCamelCase_         =     state_dict["shared.weight"]

				if not is_encoder_only:
								if "decoder.embed_tokens.weight" not in state_dict:
												lowerCamelCase_         =     state_dict["shared.weight"]

								if "lm_head.weight" not in state_dict:  # For old 1.0 models.
												print("Using shared word embeddings as lm_head."       )
												lowerCamelCase_         =     state_dict["shared.weight"]

				return state_dict





def      lowerCamelCase_					(							lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__       ):
				lowerCamelCase_         =     checkpoints.load_tax_checkpoint(_UpperCamelCase       )
				lowerCamelCase_         =     convert_tax_to_pytorch(
				    _UpperCamelCase      ,					num_layers=config.num_layers      ,					is_encoder_only=_UpperCamelCase      ,					scalable_attention=_UpperCamelCase       )
				lowerCamelCase_         =     make_state_dict(_UpperCamelCase      ,					_UpperCamelCase       )
				model.load_state_dict(_UpperCamelCase      ,					strict=_UpperCamelCase       )





def      lowerCamelCase_					(							lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__      ,					lowerCamelCase__ = False      ,					lowerCamelCase__ = False      ,					):
				lowerCamelCase_         =     MTaConfig.from_json_file(_UpperCamelCase       )
				print(F'Building PyTorch model from configuration: {config}'       )
				# Non-v1.1 checkpoints could also use T5Model, but this works for all.
				# The v1.0 checkpoints will simply have an LM head that is the word embeddings.
				if is_encoder_only:
								lowerCamelCase_         =     UMTaEncoderModel(_UpperCamelCase       )
				else:
								lowerCamelCase_         =     UMTaForConditionalGeneration(_UpperCamelCase       )

				# Load weights from tf checkpoint
				load_tax_weights_in_ta(_UpperCamelCase      ,					_UpperCamelCase      ,					_UpperCamelCase      ,					_UpperCamelCase      ,					_UpperCamelCase       )

				# Save pytorch-model
				print(F'Save PyTorch model to {pytorch_dump_path}'       )
				model.save_pretrained(_UpperCamelCase       )

				# Verify that we can load the checkpoint.
				model.from_pretrained(_UpperCamelCase       )
				print("Done"       )


if __name__ == "__main__":
			__A							=argparse.ArgumentParser(description='''Converts a native T5X checkpoint into a PyTorch checkpoint.''')
			# Required parameters
			parser.add_argument(
			    '''--t5x_checkpoint_path''', default=None, type=str, required=True, help='''Path to the T5X checkpoint.'''
			)
			parser.add_argument(
			    '''--config_file''',
			    default=None,
			    type=str,
			    required=True,
			    help='''The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.''',
			)
			parser.add_argument(
			    '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
			)
			parser.add_argument(
			    '''--is_encoder_only''', action='''store_true''', help='''Check if the model is encoder-decoder model''', default=False
			)
			parser.add_argument(
			    '''--scalable_attention''',
			    action='''store_true''',
			    help='''Whether the model uses scaled attention (umt5 model)''',
			    default=False,
			)
			__A							=parser.parse_args()
			convert_tax_checkpoint_to_pytorch(
			    args.tax_checkpoint_path,
			    args.config_file,
			    args.pytorch_dump_path,
			    args.is_encoder_only,
			    args.scalable_attention,
			)










import torch
from torch import nn





class 					A__     (     nn.Module ):

     '''simple docstring'''






     def __init__( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=1   ,						_SCREAMING_SNAKE_CASE      :    List[str]=False ):






           """simple docstring"""




           super().__init__()

           UpperCamelCase					       =  n_token
           UpperCamelCase					       =  d_embed
           UpperCamelCase					       =  d_proj

           UpperCamelCase					       =  cutoffs + [n_token]
           UpperCamelCase					       =  [0] + self.cutoffs
           UpperCamelCase					       =  div_val

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

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

           UpperCamelCase					       =  nn.ModuleList()
           UpperCamelCase					       =  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 ) ):
                       UpperCamelCase       ,  UpperCamelCase					       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
                       UpperCamelCase					       =  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 ) )

           UpperCamelCase					       =  keep_order







     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




           if proj is None:
                 UpperCamelCase					       =  nn.functional.linear(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						bias=_SCREAMING_SNAKE_CASE )
           else:
                 # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
                 UpperCamelCase					       =  nn.functional.linear(_SCREAMING_SNAKE_CASE   ,						proj.t().contiguous() )
                 UpperCamelCase					       =  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      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    Dict=False ):






           """simple docstring"""




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

           if self.n_clusters == 0:
                 UpperCamelCase					       =  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:
                       UpperCamelCase					       =  labels != -100
                       UpperCamelCase					       =  torch.zeros_like(_SCREAMING_SNAKE_CASE   ,						dtype=hidden.dtype   ,						device=hidden.device )
                       UpperCamelCase					       =  (
                           -nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=-1 )[mask].gather(1   ,						labels[mask].unsqueeze(1 ) ).squeeze(1 )
                       )
                 else:
                       UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=-1 )
           else:
                 # construct weights and biases
                 UpperCamelCase       ,  UpperCamelCase					       =  [], []
                 for i in range(len(self.cutoffs ) ):
                       if self.div_val == 1:
                             UpperCamelCase       ,  UpperCamelCase					       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
                             UpperCamelCase					       =  self.out_layers[0].weight[l_idx:r_idx]
                             UpperCamelCase					       =  self.out_layers[0].bias[l_idx:r_idx]
                       else:
                             UpperCamelCase					       =  self.out_layers[i].weight
                             UpperCamelCase					       =  self.out_layers[i].bias

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

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

                 UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  weights[0], biases[0], self.out_projs[0]

                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

                 if labels is None:
                       UpperCamelCase					       =  hidden.new_empty((head_logit.size(0 ), self.n_token) )
                 else:
                       UpperCamelCase					       =  torch.zeros_like(_SCREAMING_SNAKE_CASE   ,						dtype=hidden.dtype   ,						device=hidden.device )

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

                       if labels is not None:
                             UpperCamelCase					       =  (labels >= l_idx) & (labels < r_idx)
                             UpperCamelCase					       =  mask_i.nonzero().squeeze()

                             if indices_i.numel() == 0:
                                   continue

                             UpperCamelCase					       =  labels.index_select(0   ,						_SCREAMING_SNAKE_CASE ) - l_idx
                             UpperCamelCase					       =  head_logprob.index_select(0   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  hidden.index_select(0   ,						_SCREAMING_SNAKE_CASE )
                       else:
                             UpperCamelCase					       =  hidden

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

                             UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )
                             UpperCamelCase					       =  self.cutoffs[0] + i - 1  # No probability for the head cluster
                             if labels is not None:
                                   UpperCamelCase					       =  head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                                       1   ,						target_i[:, None] ).squeeze(1 )
                             else:
                                   UpperCamelCase					       =  head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                                   UpperCamelCase					       =  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      :    int   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           if self.n_clusters == 0:
                 UpperCamelCase					       =  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
                 UpperCamelCase       ,  UpperCamelCase					       =  [], []
                 for i in range(len(self.cutoffs ) ):
                       if self.div_val == 1:
                             UpperCamelCase       ,  UpperCamelCase					       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
                             UpperCamelCase					       =  self.out_layers[0].weight[l_idx:r_idx]
                             UpperCamelCase					       =  self.out_layers[0].bias[l_idx:r_idx]
                       else:
                             UpperCamelCase					       =  self.out_layers[i].weight
                             UpperCamelCase					       =  self.out_layers[i].bias

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

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

                 UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  weights[0], biases[0], self.out_projs[0]
                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

                 UpperCamelCase					       =  hidden.new_empty((head_logit.size(0 ), self.n_token) )
                 UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

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

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

                             UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

                             UpperCamelCase					       =  head_logprob[:, -i] + tail_logprob_i
                             UpperCamelCase					       =  logprob_i

                 return out



