


import torch
from torch import nn

class   A__ (   nn.Module     ):







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case=1      , __snake_case=False   ):
         super().__init__()

         snake_case					=    n_token
         snake_case					=    d_embed
         snake_case					=    d_proj

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

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

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

         snake_case					=    nn.ModuleList()
         snake_case					=    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(__snake_case      , __snake_case   )   )   )
                   else:
                        self.out_projs.append(__snake_case   )

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

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

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

         snake_case					=    keep_order
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         if proj is None:
              snake_case					=    nn.functional.linear(__snake_case      , __snake_case      , bias=__snake_case   )
         else:
              # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
              snake_case					=    nn.functional.linear(__snake_case      , proj.t().contiguous()   )
              snake_case					=    nn.functional.linear(__snake_case      , __snake_case      , bias=__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 		a_      (	self      , __snake_case      , __snake_case=None      , __snake_case=False   ):

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

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

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

                   weights.append(__snake_case   )
                   biases.append(__snake_case   )

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

              snake_case					=    self._compute_logit(__snake_case      , __snake_case      , __snake_case      , __snake_case   )
              snake_case					=    nn.functional.log_softmax(__snake_case      , dim=1   )

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

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

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

                        if indices_i.numel() == 0:
                             continue

                        snake_case					=    labels.index_select(0      , __snake_case   ) - l_idx
                        snake_case					=    head_logprob.index_select(0      , __snake_case   )
                        snake_case					=    hidden.index_select(0      , __snake_case   )
                   else:
                        snake_case					=    hidden

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

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

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

         return out




    def 		a_      (	self      , __snake_case   ):
         if self.n_clusters == 0:
              snake_case					=    self._compute_logit(__snake_case      , self.out_layers[0].weight      , self.out_layers[0].bias      , self.out_projs[0]   )
              return nn.functional.log_softmax(__snake_case      , dim=-1   )
         else:
              # construct weights and biases
              snake_case			,  snake_case					=    [], []
              for i in range(len(self.cutoffs   )   ):
                   if self.div_val == 1:
                        snake_case			,  snake_case					=    self.cutoff_ends[i], self.cutoff_ends[i + 1]
                        snake_case					=    self.out_layers[0].weight[l_idx:r_idx]
                        snake_case					=    self.out_layers[0].bias[l_idx:r_idx]
                   else:
                        snake_case					=    self.out_layers[i].weight
                        snake_case					=    self.out_layers[i].bias

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

                   weights.append(__snake_case   )
                   biases.append(__snake_case   )

              snake_case			,  snake_case			,  snake_case					=    weights[0], biases[0], self.out_projs[0]
              snake_case					=    self._compute_logit(__snake_case      , __snake_case      , __snake_case      , __snake_case   )

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

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

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

                        snake_case					=    self._compute_logit(__snake_case      , __snake_case      , __snake_case      , __snake_case   )
                        snake_case					=    nn.functional.log_softmax(__snake_case      , dim=1   )

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

              return out








from __future__ import annotations

import time

_SCREAMING_SNAKE_CASE      :			List[Any]      				=   list[tuple[int, int]]

_SCREAMING_SNAKE_CASE      :			Any      				=   [
    [0, 0, 0, 0, 0, 0, 0],
    [0, 1, 0, 0, 0, 0, 0],  # 0 are free path whereas 1's are obstacles
    [0, 0, 0, 0, 0, 0, 0],
    [0, 0, 1, 0, 0, 0, 0],
    [1, 0, 1, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 1, 0, 0],
]

_SCREAMING_SNAKE_CASE      :			Any      				=   [[-1, 0], [0, -1], [1, 0], [0, 1]]  # up, left, down, right

class   A__ :







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    pos_x
         snake_case					=    pos_y
         snake_case					=    (pos_y, pos_x)
         snake_case					=    goal_x
         snake_case					=    goal_y
         snake_case					=    parent






class   A__ :







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case   ):
         snake_case					=    Node(start[1]      , start[0]      , goal[1]      , goal[0]      , __snake_case   )
         snake_case					=    Node(goal[1]      , goal[0]      , goal[1]      , goal[0]      , __snake_case   )

         snake_case					=    [self.start]
         snake_case					=    False
    def 		a_      (	self   ):
         while self.node_queue:
              snake_case					=    self.node_queue.pop(0   )

              if current_node.pos == self.target.pos:
                   snake_case					=    True
                   return self.retrace_path(__snake_case   )

              snake_case					=    self.get_successors(__snake_case   )

              for node in successors:
                   self.node_queue.append(__snake_case   )

         if not self.reached:
              return [self.start.pos]
         return None
    def 		a_      (	self      , __snake_case   ):
         snake_case					=    []
         for action in delta:
              snake_case					=    parent.pos_x + action[1]
              snake_case					=    parent.pos_y + action[0]
              if not (0 <= pos_x <= len(grid[0]   ) - 1 and 0 <= pos_y <= len(__snake_case   ) - 1):
                   continue

              if grid[pos_y][pos_x] != 0:
                   continue

              successors.append(
                  Node(__snake_case      , __snake_case      , self.target.pos_y      , self.target.pos_x      , __snake_case   )   )
         return successors




    def 		a_      (	self      , __snake_case   ):
         snake_case					=    node
         snake_case					=    []
         while current_node is not None:
              path.append((current_node.pos_y, current_node.pos_x)   )
              snake_case					=    current_node.parent
         path.reverse()
         return path






class   A__ :







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case   ):
         snake_case					=    BreadthFirstSearch(__snake_case      , __snake_case   )
         snake_case					=    BreadthFirstSearch(__snake_case      , __snake_case   )
         snake_case					=    False
    def 		a_      (	self   ):
         while self.fwd_bfs.node_queue or self.bwd_bfs.node_queue:
              snake_case					=    self.fwd_bfs.node_queue.pop(0   )
              snake_case					=    self.bwd_bfs.node_queue.pop(0   )

              if current_bwd_node.pos == current_fwd_node.pos:
                   snake_case					=    True
                   return self.retrace_bidirectional_path(
                       __snake_case      , __snake_case   )

              snake_case					=    current_bwd_node
              snake_case					=    current_fwd_node

              snake_case					=    {
                  self.fwd_bfs: self.fwd_bfs.get_successors(__snake_case   ),
                  self.bwd_bfs: self.bwd_bfs.get_successors(__snake_case   ),
              }

              for bfs in [self.fwd_bfs, self.bwd_bfs]:
                   for node in successors[bfs]:
                        bfs.node_queue.append(__snake_case   )

         if not self.reached:
              return [self.fwd_bfs.start.pos]
         return None




    def 		a_      (	self      , __snake_case      , __snake_case   ):
         snake_case					=    self.fwd_bfs.retrace_path(__snake_case   )
         snake_case					=    self.bwd_bfs.retrace_path(__snake_case   )
         bwd_path.pop()
         bwd_path.reverse()
         snake_case					=    fwd_path + bwd_path
         return path


if __name__ == "__main__":
 # all coordinates are given in format [y,x]
 import doctest

 doctest.testmod()
 _SCREAMING_SNAKE_CASE      :			Optional[Any]      				=   (0, 0)
 _SCREAMING_SNAKE_CASE      :			List[Any]      				=   (len(grid) - 1, len(grid[0]) - 1)
 for elem in grid:
  print(elem)

 _SCREAMING_SNAKE_CASE      :			List[Any]      				=   time.time()
 _SCREAMING_SNAKE_CASE      :			List[Any]      				=   BreadthFirstSearch(init, goal)
 _SCREAMING_SNAKE_CASE      :			List[str]      				=   bfs.search()
 _SCREAMING_SNAKE_CASE      :			int      				=   time.time() - start_bfs_time

 print("Unidirectional BFS computation time : ", bfs_time)

 _SCREAMING_SNAKE_CASE      :			Any      				=   time.time()
 _SCREAMING_SNAKE_CASE      :			Union[str, Any]      				=   BidirectionalBreadthFirstSearch(init, goal)
 _SCREAMING_SNAKE_CASE      :			Union[str, Any]      				=   bd_bfs.search()
 _SCREAMING_SNAKE_CASE      :			Tuple      				=   time.time() - start_bd_bfs_time

 print("Bidirectional BFS computation time : ", bd_bfs_time)




