






from __future__ import annotations

import math
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase) ->			list:
      """simple docstring"""

      if len(_UpperCamelCase) != 2 or len(a[0]) != 2 or len(_UpperCamelCase) != 2 or len(b[0]) != 2:
            raise Exception('Matrices are not 2x2')
      UpperCamelCase					       =  [
          [a[0][0] * b[0][0] + a[0][1] * b[1][0], a[0][0] * b[0][1] + a[0][1] * b[1][1]],
          [a[1][0] * b[0][0] + a[1][1] * b[1][0], a[1][0] * b[0][1] + a[1][1] * b[1][1]],
      ]
      return new_matrix
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase) ->			Union[str, Any]:
      """simple docstring"""

      return [
          [matrix_a[row][col] + matrix_b[row][col] for col in range(len(matrix_a[row]))]
          for row in range(len(_UpperCamelCase))
      ]
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase) ->			Optional[int]:
      """simple docstring"""

      return [
          [matrix_a[row][col] - matrix_b[row][col] for col in range(len(matrix_a[row]))]
          for row in range(len(_UpperCamelCase))
      ]
def 		lowercase__  (				_UpperCamelCase) ->			tuple[list, list, list, list]:
      """simple docstring"""

      if len(_UpperCamelCase) % 2 != 0 or len(a[0]) % 2 != 0:
            raise Exception('Odd matrices are not supported!')

      UpperCamelCase					       =  len(_UpperCamelCase)
      UpperCamelCase					       =  matrix_length // 2

      UpperCamelCase					       =  [[a[i][j] for j in range(_UpperCamelCase       ,							_UpperCamelCase)] for i in range(_UpperCamelCase)]
      UpperCamelCase					       =  [
          [a[i][j] for j in range(_UpperCamelCase       ,							_UpperCamelCase)] for i in range(_UpperCamelCase       ,							_UpperCamelCase)
      ]

      UpperCamelCase					       =  [[a[i][j] for j in range(_UpperCamelCase)] for i in range(_UpperCamelCase)]
      UpperCamelCase					       =  [[a[i][j] for j in range(_UpperCamelCase)] for i in range(_UpperCamelCase       ,							_UpperCamelCase)]

      return top_left, top_right, bot_left, bot_right
def 		lowercase__  (				_UpperCamelCase) ->			tuple[int, int]:
      """simple docstring"""

      return len(_UpperCamelCase), len(matrix[0])
def 		lowercase__  (				_UpperCamelCase) ->			None:
      """simple docstring"""

      print('\n'.join(str(_UpperCamelCase) for line in matrix))
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase) ->			list:
      """simple docstring"""

      if matrix_dimensions(_UpperCamelCase) == (2, 2):
            return default_matrix_multiplication(_UpperCamelCase       ,							_UpperCamelCase)

      UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  split_matrix(_UpperCamelCase)
      UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  split_matrix(_UpperCamelCase)

      UpperCamelCase					       =  actual_strassen(_UpperCamelCase       ,							matrix_subtraction(_UpperCamelCase       ,							_UpperCamelCase))
      UpperCamelCase					       =  actual_strassen(matrix_addition(_UpperCamelCase       ,							_UpperCamelCase)       ,							_UpperCamelCase)
      UpperCamelCase					       =  actual_strassen(matrix_addition(_UpperCamelCase       ,							_UpperCamelCase)       ,							_UpperCamelCase)
      UpperCamelCase					       =  actual_strassen(_UpperCamelCase       ,							matrix_subtraction(_UpperCamelCase       ,							_UpperCamelCase))
      UpperCamelCase					       =  actual_strassen(matrix_addition(_UpperCamelCase       ,							_UpperCamelCase)       ,							matrix_addition(_UpperCamelCase       ,							_UpperCamelCase))
      UpperCamelCase					       =  actual_strassen(matrix_subtraction(_UpperCamelCase       ,							_UpperCamelCase)       ,							matrix_addition(_UpperCamelCase       ,							_UpperCamelCase))
      UpperCamelCase					       =  actual_strassen(matrix_subtraction(_UpperCamelCase       ,							_UpperCamelCase)       ,							matrix_addition(_UpperCamelCase       ,							_UpperCamelCase))

      UpperCamelCase					       =  matrix_addition(matrix_subtraction(matrix_addition(_UpperCamelCase       ,							_UpperCamelCase)       ,							_UpperCamelCase)       ,							_UpperCamelCase)
      UpperCamelCase					       =  matrix_addition(_UpperCamelCase       ,							_UpperCamelCase)
      UpperCamelCase					       =  matrix_addition(_UpperCamelCase       ,							_UpperCamelCase)
      UpperCamelCase					       =  matrix_subtraction(matrix_subtraction(matrix_addition(_UpperCamelCase       ,							_UpperCamelCase)       ,							_UpperCamelCase)       ,							_UpperCamelCase)

      # construct the new matrix from our 4 quadrants
      UpperCamelCase					       =  []
      for i in range(len(_UpperCamelCase)):
            new_matrix.append(top_left[i] + top_right[i])
      for i in range(len(_UpperCamelCase)):
            new_matrix.append(bot_left[i] + bot_right[i])
      return new_matrix
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase) ->			list:
      """simple docstring"""

      if matrix_dimensions(_UpperCamelCase)[1] != matrix_dimensions(_UpperCamelCase)[0]:
            UpperCamelCase					       =  (
                'Unable to multiply these matrices, please check the dimensions.\n'
                F'Matrix A: {matrixa}\n'
                F'Matrix B: {matrixa}'
            )
            raise Exception(_UpperCamelCase)
      UpperCamelCase					       =  matrix_dimensions(_UpperCamelCase)
      UpperCamelCase					       =  matrix_dimensions(_UpperCamelCase)

      if dimensiona[0] == dimensiona[1] and dimensiona[0] == dimensiona[1]:
            return [matrixa, matrixa]

      UpperCamelCase					       =  max(*_UpperCamelCase       ,							*_UpperCamelCase)
      UpperCamelCase					       =  int(math.pow(2       ,							math.ceil(math.loga(_UpperCamelCase))))
      UpperCamelCase					       =  matrixa
      UpperCamelCase					       =  matrixa

      # Adding zeros to the matrices so that the arrays dimensions are the same and also
      # power of 2
      for i in range(0       ,							_UpperCamelCase):
            if i < dimensiona[0]:
                  for _ in range(dimensiona[1]       ,							_UpperCamelCase):
                        new_matrixa[i].append(0)
            else:
                  new_matrixa.append([0] * maxim)
            if i < dimensiona[0]:
                  for _ in range(dimensiona[1]       ,							_UpperCamelCase):
                        new_matrixa[i].append(0)
            else:
                  new_matrixa.append([0] * maxim)

      UpperCamelCase					       =  actual_strassen(_UpperCamelCase       ,							_UpperCamelCase)

      # Removing the additional zeros
      for i in range(0       ,							_UpperCamelCase):
            if i < dimensiona[0]:
                  for _ in range(dimensiona[1]       ,							_UpperCamelCase):
                        final_matrix[i].pop()
            else:
                  final_matrix.pop()
      return final_matrix


if __name__ == "__main__":
       __magic_name__   :  Optional[Any]        =							[
           [2, 3, 4, 5],
           [6, 4, 3, 1],
           [2, 3, 6, 7],
           [3, 1, 2, 4],
           [2, 3, 4, 5],
           [6, 4, 3, 1],
           [2, 3, 6, 7],
           [3, 1, 2, 4],
           [2, 3, 4, 5],
           [6, 2, 3, 1],
       ]
       __magic_name__   :  List[str]        =							[[0, 2, 1, 1], [16, 2, 3, 3], [2, 2, 7, 7], [13, 11, 22, 4]]
       print(strassen(matrixa, matrixa))











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



