# -*- coding: utf-8 -*-

import math
import torch
import torch.nn as nn
from torch import Tensor


# pylint: disable=arguments-differ
class MultiHeadedAttention(nn.Module):
    """
    Multi-Head Attention module from "Attention is All You Need"
    Implementation modified from OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """

    def __init__(self, num_heads: int, size: int, dropout: float = 0.1):
        """
        Create a multi-headed attention layer.
        :param num_heads: the number of heads
        :param size: model size (must be divisible by num_heads)
        :param dropout: probability of dropping a unit
        """
        super(MultiHeadedAttention, self).__init__()

        assert size % num_heads == 0

        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads

        self.k_layer = nn.Linear(size, num_heads * head_size)
        self.v_layer = nn.Linear(size, num_heads * head_size)
        self.q_layer = nn.Linear(size, num_heads * head_size)

        self.output_layer = nn.Linear(size, size)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k: Tensor, v: Tensor, q: Tensor, mask: Tensor = None):
        """
        Computes multi-headed attention.
        :param k: keys   [B, M, D] with M being the sentence length.
        :param v: values [B, M, D]
        :param q: query  [B, M, D]
        :param mask: optional mask [B, 1, M]
        :return:
        """
        batch_size = k.size(0)
        num_heads = self.num_heads

        # project the queries (q), keys (k), and values (v)
        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)

        # reshape q, k, v for our computation to [batch_size, num_heads, ..]
        k = k.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        v = v.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)

        # compute scores
        q = q / math.sqrt(self.head_size)

        # batch x num_heads x query_len x key_len
        scores = torch.matmul(q, k.transpose(2, 3))
        # apply the mask (if we have one)
        # we add a dimension for the heads to it beslow: [B, 1, 1, M]
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))

        # apply attention dropout and compute context vectors.
        attention = self.softmax(scores)
        attention = self.dropout(attention)

        # get context vector (select values with attention) and reshape
        # back to [B, M, D]
        context = torch.matmul(attention, v)
        context = context.transpose(1, 2).contiguous().view(
            batch_size, -1, num_heads * self.head_size)

        output = self.output_layer(context)

        return output

class ContMultiHeadedAttention(nn.Module):
    """
    Multi-Head Attention module from "Attention is All You Need"
    Implementation modified from OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """

    def __init__(self, num_heads: int, size: int, size_v: int, dropout: float = 0.1):
        """
        Create a multi-headed attention layer.
        :param num_heads: the number of heads
        :param size: model size (must be divisible by num_heads)
        :param dropout: probability of dropping a unit
        """
        super(ContMultiHeadedAttention, self).__init__()

        assert size % num_heads == 0

        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads

        self.k_layer = nn.Linear(size, num_heads * head_size)
        self.v_layer = nn.Linear(size_v, num_heads * head_size)
        self.q_layer = nn.Linear(size, num_heads * head_size)

        self.output_layer = nn.Linear(size, size_v)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k: Tensor, v: Tensor, q: Tensor, mask: Tensor = None):
        """
        Computes multi-headed attention.
        :param k: keys   [B, M, D] with M being the sentence length.
        :param v: values [B, M, D]
        :param q: query  [B, M, D]
        :param mask: optional mask [B, 1, M]
        :return:
        """
        batch_size = k.size(0)
        num_heads = self.num_heads

        # project the queries (q), keys (k), and values (v)
        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)

        # reshape q, k, v for our computation to [batch_size, num_heads, ..]
        k = k.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        v = v.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)

        # compute scores
        q = q / math.sqrt(self.head_size)

        # batch x num_heads x query_len x key_len
        scores = torch.matmul(q, k.transpose(2, 3))

        # apply the mask (if we have one)
        # we add a dimension for the heads to it beslow: [B, 1, 1, M]
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))

        # apply attention dropout and compute context vectors.
        attention = self.softmax(scores)
        attention = self.dropout(attention)

        # get context vector (select values with attention) and reshape
        # back to [B, M, D]
        try:
            context = torch.matmul(attention, v)
        except:
            print(attention.shape)
            print(v.shape)
        context = context.transpose(1, 2).contiguous().view(
            batch_size, -1, num_heads * self.head_size)

        output = self.output_layer(context)

        return output
    
class Custom_Attention(nn.Module):
    def __init__(self, num_heads: int, size: int, dropout: float = 0.1, sent_length: int = 128):
        super(Custom_Attention, self).__init__()

        assert size % num_heads == 0

        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads

        self.k_layer = nn.Linear(size//4, num_heads * head_size//4)
        self.v_layer = nn.Linear(size, num_heads * head_size)
        self.q_layer = nn.Linear(size//4, num_heads * head_size//4)
    
        self.confident_layer = nn.Sequential(nn.Linear(sent_length, sent_length),nn.ReLU())
        
        self.output_layer = nn.Linear(size, size)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k: Tensor, v: Tensor, q: Tensor, mask: Tensor = None, xy_pdf_score: Tensor = None):
        batch_size = k.size(0)
        num_heads = self.num_heads

        # project the queries (q), keys (k), and values (v)
        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)

        # reshape q, k, v for our computation to [batch_size, num_heads, ..]
        k = k.view(batch_size, -1, num_heads, self.head_size//4).transpose(1, 2)
        v = v.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, num_heads, self.head_size//4).transpose(1, 2)

        # compute scores
        q = q / math.sqrt(self.head_size)

        # batch x num_heads x query_len x key_len
        scores = torch.matmul(q, k.transpose(2, 3))
        # apply the mask (if we have one)
        # we add a dimension for the heads to it below: [B, 1, 1, M]
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))
#         print(scores.shape) #xy_pdf_score
        # apply attention dropout and compute context vectors.
        if xy_pdf_score is None:
            attention = self.softmax(scores)
        else:
            xy_pdf_score = self.confident_layer(xy_pdf_score)
#             xy_pdf_score = (xy_pdf_score - xy_pdf_score.mean(1).unsqueeze(-1)) + 1.
            xy_pdf_score = xy_pdf_score.view(batch_size,1,1,-1)
    
            scores_exp = scores.exp()
            new_scores = scores_exp * xy_pdf_score
            new_scores_sum = new_scores.sum(-1)
            attention = new_scores/new_scores_sum.unsqueeze(-1)
        
        attention = self.dropout(attention)

        # get context vector (select values with attention) and reshape
        # back to [B, M, D]
        context = torch.matmul(attention, v)
        context = context.transpose(1, 2).contiguous().view(
            batch_size, -1, num_heads * self.head_size)

        output = self.output_layer(context)

        return output

class GELU(nn.Module):
    """
    Paper Section 3.4, last paragraph notice that BERT used the GELU instead of RELU
    """

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))


# pylint: disable=arguments-differ
class PositionwiseFeedForward(nn.Module):
    """
    Position-wise Feed-forward layer
    Projects to ff_size and then back down to input_size.
    """

    def __init__(self, input_size, ff_size, dropout=0.1):
        """
        Initializes position-wise feed-forward layer.
        :param input_size: dimensionality of the input.
        :param ff_size: dimensionality of intermediate representation
        :param dropout:
        """
        super(PositionwiseFeedForward, self).__init__()
        self.layer_norm = nn.LayerNorm(input_size, eps=1e-6)
        self.pwff_layer = nn.Sequential(
            nn.Linear(input_size, ff_size),
            GELU(),
            nn.Dropout(dropout),
            nn.Linear(ff_size, input_size),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        x_norm = self.layer_norm(x)
        return self.pwff_layer(x_norm) + x


# pylint: disable=arguments-differ
class PositionalEncoding(nn.Module):
    """
    Pre-compute position encodings (PE).
    In forward pass, this adds the position-encodings to the
    input for as many time steps as necessary.
    Implementation based on OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """
    def __init__(self,
                 size: int = 0,
                 max_len: int = 5000):
        """
        Positional Encoding with maximum length max_len
        :param size:
        :param max_len:
        :param dropout:
        """
        if size % 2 != 0:
            raise ValueError("Cannot use sin/cos positional encoding with "
                             "odd dim (got dim={:d})".format(size))
        pe = torch.zeros(max_len, size)
        pe.require_grad = False

        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp((torch.arange(0, size, 2, dtype=torch.float) *
                              -(math.log(10000.0) / size)))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)  # shape: [1, size, max_len]
        super(PositionalEncoding, self).__init__()
        self.register_buffer('pe', pe)
        self.dim = size

    def forward(self, emb):
        """Embed inputs.
        Args:
            emb (FloatTensor): Sequence of word vectors
                ``(seq_len, batch_size, self.dim)``
        """
        # Add position encodings
        return emb + self.pe[:, :emb.size(1)]


class TransformerEncoderLayer(nn.Module):
    """
    One Transformer encoder layer has a Multi-head attention layer plus
    a position-wise feed-forward layer.
    """

    def __init__(self,
                 size: int = 0,
                 ff_size: int = 0,
                 num_heads: int = 0,
                 dropout: float = 0.1):
        """
        A single Transformer layer.
        :param size:
        :param ff_size:
        :param num_heads:
        :param dropout:
        """
        super(TransformerEncoderLayer, self).__init__()

        self.layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.src_src_att = MultiHeadedAttention(num_heads, size,
                                                dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    # pylint: disable=arguments-differ
    def forward(self, x: Tensor, mask: Tensor) -> Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        """
        x_norm = self.layer_norm(x)
        h = self.src_src_att(x_norm, x_norm, x_norm, mask)
        h = self.dropout(h) + x
        o = self.feed_forward(h)
        return o


class TransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self,
                 size: int = 0,
                 ff_size: int = 0,
                 num_heads: int = 0,
                 dropout: float = 0.1,
                 src_trg_att: bool = True):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerDecoderLayer, self).__init__()
        self.size = size
        self.src_trg_att = src_trg_att
        self.trg_trg_att = MultiHeadedAttention(num_heads, size,
                                                dropout=dropout)
        if src_trg_att:
            self.src_trg_att = MultiHeadedAttention(num_heads, size,
                                                dropout=dropout)

        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)

        self.x_layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.dec_layer_norm = nn.LayerNorm(size, eps=1e-6)

        self.dropout = nn.Dropout(dropout)

    # pylint: disable=arguments-differ
    def forward(self,
                x: Tensor = None,
                memory: Tensor = None,
                src_mask: Tensor = None,
                trg_mask: Tensor = None) -> Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        # decoder/target self-attention

        x_norm = self.x_layer_norm(x)
#         print(x_norm)
#         trg_mask()
#         print(trg_mask)

        h1 = self.trg_trg_att(x_norm, x_norm, x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + x


        if self.src_trg_att:
            # source-target attention
            h1_norm = self.dec_layer_norm(h1)
            h2 = self.src_trg_att(memory, memory, h1_norm, mask=src_mask)
            # final position-wise feed-forward layer
            o = self.feed_forward(self.dropout(h2) + h1)
        else:
            o = self.feed_forward(h1)
        return o

class CustomTransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self,
                 size: int = 0,
                 bb_size: int = 64,
                 ff_size: int = 0,
                 num_heads: int = 0,
                 dropout: float = 0.1):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(CustomTransformerDecoderLayer, self).__init__()
        self.size = size

        self.trg_trg_att = ContMultiHeadedAttention(num_heads, bb_size, bb_size,
                                                dropout=dropout)
        self.src_trg_att = ContMultiHeadedAttention(num_heads, size, size,
                                                dropout=dropout)

        self.feed_forward_h1 = PositionwiseFeedForward(bb_size, ff_size=ff_size)
        self.feed_forward_h2 = PositionwiseFeedForward(size, ff_size=ff_size)

        self.x_layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.spa_layer_norm = nn.LayerNorm(bb_size, eps=1e-6)

        self.dropout = nn.Dropout(dropout)

    # pylint: disable=arguments-differ
    def forward(self,
                spatial_x: Tensor = None,
                semantic_x: Tensor = None,
                memory: Tensor = None,
                src_mask: Tensor = None,
                trg_mask: Tensor = None) -> Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        # decoder/target self-attention

        spatial_x_norm = self.spa_layer_norm(spatial_x) #64
        semantic_x_norm = self.x_layer_norm(semantic_x) #256

        # layout attention

        h1 = self.trg_trg_att(spatial_x_norm, spatial_x_norm, spatial_x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + spatial_x

#         h2 = self.src_trg_att(memory, memory, semantic_x_norm, mask=src_mask)
#         h2 = self.dropout(h2)
        
        o1 = self.feed_forward_h1(h1)
        o2 = memory[:,1:,:]
        o = torch.cat((o2, o1), dim=-1)
        
        return o

class TransformerRefineLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self,
                 size: int = 0,
                 ff_size: int = 0,
                 num_heads: int = 0,
                 dropout: float = 0.1,
                 sent_length: int = 128,):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerRefineLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.box_norm = nn.LayerNorm(size//4, eps=1e-6)
        self.src_src_att = Custom_Attention(num_heads, size,
                                                dropout=dropout, sent_length=sent_length)
        self.combine_layer = nn.Linear(size+size//4, size)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    # pylint: disable=arguments-differ
    def forward(self, context: Tensor, box: Tensor, mask: Tensor, xy_pdf_score: Tensor) -> Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        context: torch.Size([64, 128, 256])
        box: torch.Size([64, 128, 64])
        context_norm: torch.Size([64, 128, 256])
        box_norm: torch.Size([64, 128, 64])
        xy_pdf_score: torch.Size([64, 128])
        """
        
#         print('context:',context.shape)
#         print('box:',box.shape)
        context_norm = self.layer_norm(context)
#         print('context_norm:',context_norm.shape)
        box_norm = self.box_norm(box)
#         print('box_norm:',box_norm.shape)
        # Calculate xy_pdf_score
#         combine_contex = torch.cat((context_norm, box_norm), dim=-1)
#         combine_contex = self.combine_layer(combine_contex)
        # k, v, q
        h = self.src_src_att(box_norm, context_norm, box_norm, mask, xy_pdf_score)
        h = self.dropout(h) + context_norm
        o = self.feed_forward(h)
        return o

3b542aaa247bfb8065c48790c0be18ccedf2d3a9
<mask token>


class PositionwiseFeedForward(nn.Module):
    <mask token>
    <mask token>
    <mask token>


class PositionalEncoding(nn.Module):
    """
    Pre-compute position encodings (PE).
    In forward pass, this adds the position-encodings to the
    input for as many time steps as necessary.
    Implementation based on OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """

    def __init__(self, size: int=0, max_len: int=5000):
        """
        Positional Encoding with maximum length max_len
        :param size:
        :param max_len:
        :param dropout:
        """
        if size % 2 != 0:
            raise ValueError(
                'Cannot use sin/cos positional encoding with odd dim (got dim={:d})'
                .format(size))
        pe = torch.zeros(max_len, size)
        pe.require_grad = False
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, size, 2, dtype=torch.float) * 
            -(math.log(10000.0) / size))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)
        super(PositionalEncoding, self).__init__()
        self.register_buffer('pe', pe)
        self.dim = size

    def forward(self, emb):
        """Embed inputs.
        Args:
            emb (FloatTensor): Sequence of word vectors
                ``(seq_len, batch_size, self.dim)``
        """
        return emb + self.pe[:, :emb.size(1)]


class TransformerEncoderLayer(nn.Module):
    """
    One Transformer encoder layer has a Multi-head attention layer plus
    a position-wise feed-forward layer.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1):
        """
        A single Transformer layer.
        :param size:
        :param ff_size:
        :param num_heads:
        :param dropout:
        """
        super(TransformerEncoderLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.src_src_att = MultiHeadedAttention(num_heads, size, dropout=
            dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    def forward(self, x: Tensor, mask: Tensor) ->Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        """
        x_norm = self.layer_norm(x)
        h = self.src_src_att(x_norm, x_norm, x_norm, mask)
        h = self.dropout(h) + x
        o = self.feed_forward(h)
        return o


class TransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1, src_trg_att: bool=True):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerDecoderLayer, self).__init__()
        self.size = size
        self.src_trg_att = src_trg_att
        self.trg_trg_att = MultiHeadedAttention(num_heads, size, dropout=
            dropout)
        if src_trg_att:
            self.src_trg_att = MultiHeadedAttention(num_heads, size,
                dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.x_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.dec_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: Tensor=None, memory: Tensor=None, src_mask: Tensor
        =None, trg_mask: Tensor=None) ->Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        x_norm = self.x_layer_norm(x)
        h1 = self.trg_trg_att(x_norm, x_norm, x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + x
        if self.src_trg_att:
            h1_norm = self.dec_layer_norm(h1)
            h2 = self.src_trg_att(memory, memory, h1_norm, mask=src_mask)
            o = self.feed_forward(self.dropout(h2) + h1)
        else:
            o = self.feed_forward(h1)
        return o


class CustomTransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, bb_size: int=64, ff_size: int=0,
        num_heads: int=0, dropout: float=0.1):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(CustomTransformerDecoderLayer, self).__init__()
        self.size = size
        self.trg_trg_att = ContMultiHeadedAttention(num_heads, bb_size,
            bb_size, dropout=dropout)
        self.src_trg_att = ContMultiHeadedAttention(num_heads, size, size,
            dropout=dropout)
        self.feed_forward_h1 = PositionwiseFeedForward(bb_size, ff_size=ff_size
            )
        self.feed_forward_h2 = PositionwiseFeedForward(size, ff_size=ff_size)
        self.x_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.spa_layer_norm = nn.LayerNorm(bb_size, eps=1e-06)
        self.dropout = nn.Dropout(dropout)

    def forward(self, spatial_x: Tensor=None, semantic_x: Tensor=None,
        memory: Tensor=None, src_mask: Tensor=None, trg_mask: Tensor=None
        ) ->Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        spatial_x_norm = self.spa_layer_norm(spatial_x)
        semantic_x_norm = self.x_layer_norm(semantic_x)
        h1 = self.trg_trg_att(spatial_x_norm, spatial_x_norm,
            spatial_x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + spatial_x
        o1 = self.feed_forward_h1(h1)
        o2 = memory[:, 1:, :]
        o = torch.cat((o2, o1), dim=-1)
        return o


class TransformerRefineLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1, sent_length: int=128):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerRefineLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.box_norm = nn.LayerNorm(size // 4, eps=1e-06)
        self.src_src_att = Custom_Attention(num_heads, size, dropout=
            dropout, sent_length=sent_length)
        self.combine_layer = nn.Linear(size + size // 4, size)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    def forward(self, context: Tensor, box: Tensor, mask: Tensor,
        xy_pdf_score: Tensor) ->Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        context: torch.Size([64, 128, 256])
        box: torch.Size([64, 128, 64])
        context_norm: torch.Size([64, 128, 256])
        box_norm: torch.Size([64, 128, 64])
        xy_pdf_score: torch.Size([64, 128])
        """
        context_norm = self.layer_norm(context)
        box_norm = self.box_norm(box)
        h = self.src_src_att(box_norm, context_norm, box_norm, mask,
            xy_pdf_score)
        h = self.dropout(h) + context_norm
        o = self.feed_forward(h)
        return o

<mask token>


class PositionwiseFeedForward(nn.Module):
    <mask token>
    <mask token>

    def forward(self, x):
        x_norm = self.layer_norm(x)
        return self.pwff_layer(x_norm) + x


class PositionalEncoding(nn.Module):
    """
    Pre-compute position encodings (PE).
    In forward pass, this adds the position-encodings to the
    input for as many time steps as necessary.
    Implementation based on OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """

    def __init__(self, size: int=0, max_len: int=5000):
        """
        Positional Encoding with maximum length max_len
        :param size:
        :param max_len:
        :param dropout:
        """
        if size % 2 != 0:
            raise ValueError(
                'Cannot use sin/cos positional encoding with odd dim (got dim={:d})'
                .format(size))
        pe = torch.zeros(max_len, size)
        pe.require_grad = False
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, size, 2, dtype=torch.float) * 
            -(math.log(10000.0) / size))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)
        super(PositionalEncoding, self).__init__()
        self.register_buffer('pe', pe)
        self.dim = size

    def forward(self, emb):
        """Embed inputs.
        Args:
            emb (FloatTensor): Sequence of word vectors
                ``(seq_len, batch_size, self.dim)``
        """
        return emb + self.pe[:, :emb.size(1)]


class TransformerEncoderLayer(nn.Module):
    """
    One Transformer encoder layer has a Multi-head attention layer plus
    a position-wise feed-forward layer.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1):
        """
        A single Transformer layer.
        :param size:
        :param ff_size:
        :param num_heads:
        :param dropout:
        """
        super(TransformerEncoderLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.src_src_att = MultiHeadedAttention(num_heads, size, dropout=
            dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    def forward(self, x: Tensor, mask: Tensor) ->Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        """
        x_norm = self.layer_norm(x)
        h = self.src_src_att(x_norm, x_norm, x_norm, mask)
        h = self.dropout(h) + x
        o = self.feed_forward(h)
        return o


class TransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1, src_trg_att: bool=True):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerDecoderLayer, self).__init__()
        self.size = size
        self.src_trg_att = src_trg_att
        self.trg_trg_att = MultiHeadedAttention(num_heads, size, dropout=
            dropout)
        if src_trg_att:
            self.src_trg_att = MultiHeadedAttention(num_heads, size,
                dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.x_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.dec_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: Tensor=None, memory: Tensor=None, src_mask: Tensor
        =None, trg_mask: Tensor=None) ->Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        x_norm = self.x_layer_norm(x)
        h1 = self.trg_trg_att(x_norm, x_norm, x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + x
        if self.src_trg_att:
            h1_norm = self.dec_layer_norm(h1)
            h2 = self.src_trg_att(memory, memory, h1_norm, mask=src_mask)
            o = self.feed_forward(self.dropout(h2) + h1)
        else:
            o = self.feed_forward(h1)
        return o


class CustomTransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, bb_size: int=64, ff_size: int=0,
        num_heads: int=0, dropout: float=0.1):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(CustomTransformerDecoderLayer, self).__init__()
        self.size = size
        self.trg_trg_att = ContMultiHeadedAttention(num_heads, bb_size,
            bb_size, dropout=dropout)
        self.src_trg_att = ContMultiHeadedAttention(num_heads, size, size,
            dropout=dropout)
        self.feed_forward_h1 = PositionwiseFeedForward(bb_size, ff_size=ff_size
            )
        self.feed_forward_h2 = PositionwiseFeedForward(size, ff_size=ff_size)
        self.x_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.spa_layer_norm = nn.LayerNorm(bb_size, eps=1e-06)
        self.dropout = nn.Dropout(dropout)

    def forward(self, spatial_x: Tensor=None, semantic_x: Tensor=None,
        memory: Tensor=None, src_mask: Tensor=None, trg_mask: Tensor=None
        ) ->Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        spatial_x_norm = self.spa_layer_norm(spatial_x)
        semantic_x_norm = self.x_layer_norm(semantic_x)
        h1 = self.trg_trg_att(spatial_x_norm, spatial_x_norm,
            spatial_x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + spatial_x
        o1 = self.feed_forward_h1(h1)
        o2 = memory[:, 1:, :]
        o = torch.cat((o2, o1), dim=-1)
        return o


class TransformerRefineLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1, sent_length: int=128):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerRefineLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.box_norm = nn.LayerNorm(size // 4, eps=1e-06)
        self.src_src_att = Custom_Attention(num_heads, size, dropout=
            dropout, sent_length=sent_length)
        self.combine_layer = nn.Linear(size + size // 4, size)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    def forward(self, context: Tensor, box: Tensor, mask: Tensor,
        xy_pdf_score: Tensor) ->Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        context: torch.Size([64, 128, 256])
        box: torch.Size([64, 128, 64])
        context_norm: torch.Size([64, 128, 256])
        box_norm: torch.Size([64, 128, 64])
        xy_pdf_score: torch.Size([64, 128])
        """
        context_norm = self.layer_norm(context)
        box_norm = self.box_norm(box)
        h = self.src_src_att(box_norm, context_norm, box_norm, mask,
            xy_pdf_score)
        h = self.dropout(h) + context_norm
        o = self.feed_forward(h)
        return o

<mask token>


class Custom_Attention(nn.Module):

    def __init__(self, num_heads: int, size: int, dropout: float=0.1,
        sent_length: int=128):
        super(Custom_Attention, self).__init__()
        assert size % num_heads == 0
        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads
        self.k_layer = nn.Linear(size // 4, num_heads * head_size // 4)
        self.v_layer = nn.Linear(size, num_heads * head_size)
        self.q_layer = nn.Linear(size // 4, num_heads * head_size // 4)
        self.confident_layer = nn.Sequential(nn.Linear(sent_length,
            sent_length), nn.ReLU())
        self.output_layer = nn.Linear(size, size)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k: Tensor, v: Tensor, q: Tensor, mask: Tensor=None,
        xy_pdf_score: Tensor=None):
        batch_size = k.size(0)
        num_heads = self.num_heads
        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)
        k = k.view(batch_size, -1, num_heads, self.head_size // 4).transpose(
            1, 2)
        v = v.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, num_heads, self.head_size // 4).transpose(
            1, 2)
        q = q / math.sqrt(self.head_size)
        scores = torch.matmul(q, k.transpose(2, 3))
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))
        if xy_pdf_score is None:
            attention = self.softmax(scores)
        else:
            xy_pdf_score = self.confident_layer(xy_pdf_score)
            xy_pdf_score = xy_pdf_score.view(batch_size, 1, 1, -1)
            scores_exp = scores.exp()
            new_scores = scores_exp * xy_pdf_score
            new_scores_sum = new_scores.sum(-1)
            attention = new_scores / new_scores_sum.unsqueeze(-1)
        attention = self.dropout(attention)
        context = torch.matmul(attention, v)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1,
            num_heads * self.head_size)
        output = self.output_layer(context)
        return output


class GELU(nn.Module):
    """
    Paper Section 3.4, last paragraph notice that BERT used the GELU instead of RELU
    """

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 
            0.044715 * torch.pow(x, 3))))


class PositionwiseFeedForward(nn.Module):
    """
    Position-wise Feed-forward layer
    Projects to ff_size and then back down to input_size.
    """

    def __init__(self, input_size, ff_size, dropout=0.1):
        """
        Initializes position-wise feed-forward layer.
        :param input_size: dimensionality of the input.
        :param ff_size: dimensionality of intermediate representation
        :param dropout:
        """
        super(PositionwiseFeedForward, self).__init__()
        self.layer_norm = nn.LayerNorm(input_size, eps=1e-06)
        self.pwff_layer = nn.Sequential(nn.Linear(input_size, ff_size),
            GELU(), nn.Dropout(dropout), nn.Linear(ff_size, input_size), nn
            .Dropout(dropout))

    def forward(self, x):
        x_norm = self.layer_norm(x)
        return self.pwff_layer(x_norm) + x


class PositionalEncoding(nn.Module):
    """
    Pre-compute position encodings (PE).
    In forward pass, this adds the position-encodings to the
    input for as many time steps as necessary.
    Implementation based on OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """

    def __init__(self, size: int=0, max_len: int=5000):
        """
        Positional Encoding with maximum length max_len
        :param size:
        :param max_len:
        :param dropout:
        """
        if size % 2 != 0:
            raise ValueError(
                'Cannot use sin/cos positional encoding with odd dim (got dim={:d})'
                .format(size))
        pe = torch.zeros(max_len, size)
        pe.require_grad = False
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, size, 2, dtype=torch.float) * 
            -(math.log(10000.0) / size))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)
        super(PositionalEncoding, self).__init__()
        self.register_buffer('pe', pe)
        self.dim = size

    def forward(self, emb):
        """Embed inputs.
        Args:
            emb (FloatTensor): Sequence of word vectors
                ``(seq_len, batch_size, self.dim)``
        """
        return emb + self.pe[:, :emb.size(1)]


class TransformerEncoderLayer(nn.Module):
    """
    One Transformer encoder layer has a Multi-head attention layer plus
    a position-wise feed-forward layer.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1):
        """
        A single Transformer layer.
        :param size:
        :param ff_size:
        :param num_heads:
        :param dropout:
        """
        super(TransformerEncoderLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.src_src_att = MultiHeadedAttention(num_heads, size, dropout=
            dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    def forward(self, x: Tensor, mask: Tensor) ->Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        """
        x_norm = self.layer_norm(x)
        h = self.src_src_att(x_norm, x_norm, x_norm, mask)
        h = self.dropout(h) + x
        o = self.feed_forward(h)
        return o


class TransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1, src_trg_att: bool=True):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerDecoderLayer, self).__init__()
        self.size = size
        self.src_trg_att = src_trg_att
        self.trg_trg_att = MultiHeadedAttention(num_heads, size, dropout=
            dropout)
        if src_trg_att:
            self.src_trg_att = MultiHeadedAttention(num_heads, size,
                dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.x_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.dec_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: Tensor=None, memory: Tensor=None, src_mask: Tensor
        =None, trg_mask: Tensor=None) ->Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        x_norm = self.x_layer_norm(x)
        h1 = self.trg_trg_att(x_norm, x_norm, x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + x
        if self.src_trg_att:
            h1_norm = self.dec_layer_norm(h1)
            h2 = self.src_trg_att(memory, memory, h1_norm, mask=src_mask)
            o = self.feed_forward(self.dropout(h2) + h1)
        else:
            o = self.feed_forward(h1)
        return o


class CustomTransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, bb_size: int=64, ff_size: int=0,
        num_heads: int=0, dropout: float=0.1):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(CustomTransformerDecoderLayer, self).__init__()
        self.size = size
        self.trg_trg_att = ContMultiHeadedAttention(num_heads, bb_size,
            bb_size, dropout=dropout)
        self.src_trg_att = ContMultiHeadedAttention(num_heads, size, size,
            dropout=dropout)
        self.feed_forward_h1 = PositionwiseFeedForward(bb_size, ff_size=ff_size
            )
        self.feed_forward_h2 = PositionwiseFeedForward(size, ff_size=ff_size)
        self.x_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.spa_layer_norm = nn.LayerNorm(bb_size, eps=1e-06)
        self.dropout = nn.Dropout(dropout)

    def forward(self, spatial_x: Tensor=None, semantic_x: Tensor=None,
        memory: Tensor=None, src_mask: Tensor=None, trg_mask: Tensor=None
        ) ->Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        spatial_x_norm = self.spa_layer_norm(spatial_x)
        semantic_x_norm = self.x_layer_norm(semantic_x)
        h1 = self.trg_trg_att(spatial_x_norm, spatial_x_norm,
            spatial_x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + spatial_x
        o1 = self.feed_forward_h1(h1)
        o2 = memory[:, 1:, :]
        o = torch.cat((o2, o1), dim=-1)
        return o


class TransformerRefineLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1, sent_length: int=128):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerRefineLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.box_norm = nn.LayerNorm(size // 4, eps=1e-06)
        self.src_src_att = Custom_Attention(num_heads, size, dropout=
            dropout, sent_length=sent_length)
        self.combine_layer = nn.Linear(size + size // 4, size)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    def forward(self, context: Tensor, box: Tensor, mask: Tensor,
        xy_pdf_score: Tensor) ->Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        context: torch.Size([64, 128, 256])
        box: torch.Size([64, 128, 64])
        context_norm: torch.Size([64, 128, 256])
        box_norm: torch.Size([64, 128, 64])
        xy_pdf_score: torch.Size([64, 128])
        """
        context_norm = self.layer_norm(context)
        box_norm = self.box_norm(box)
        h = self.src_src_att(box_norm, context_norm, box_norm, mask,
            xy_pdf_score)
        h = self.dropout(h) + context_norm
        o = self.feed_forward(h)
        return o

<mask token>


class ContMultiHeadedAttention(nn.Module):
    <mask token>

    def __init__(self, num_heads: int, size: int, size_v: int, dropout:
        float=0.1):
        """
        Create a multi-headed attention layer.
        :param num_heads: the number of heads
        :param size: model size (must be divisible by num_heads)
        :param dropout: probability of dropping a unit
        """
        super(ContMultiHeadedAttention, self).__init__()
        assert size % num_heads == 0
        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads
        self.k_layer = nn.Linear(size, num_heads * head_size)
        self.v_layer = nn.Linear(size_v, num_heads * head_size)
        self.q_layer = nn.Linear(size, num_heads * head_size)
        self.output_layer = nn.Linear(size, size_v)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k: Tensor, v: Tensor, q: Tensor, mask: Tensor=None):
        """
        Computes multi-headed attention.
        :param k: keys   [B, M, D] with M being the sentence length.
        :param v: values [B, M, D]
        :param q: query  [B, M, D]
        :param mask: optional mask [B, 1, M]
        :return:
        """
        batch_size = k.size(0)
        num_heads = self.num_heads
        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)
        k = k.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        v = v.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q / math.sqrt(self.head_size)
        scores = torch.matmul(q, k.transpose(2, 3))
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))
        attention = self.softmax(scores)
        attention = self.dropout(attention)
        try:
            context = torch.matmul(attention, v)
        except:
            print(attention.shape)
            print(v.shape)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1,
            num_heads * self.head_size)
        output = self.output_layer(context)
        return output


class Custom_Attention(nn.Module):

    def __init__(self, num_heads: int, size: int, dropout: float=0.1,
        sent_length: int=128):
        super(Custom_Attention, self).__init__()
        assert size % num_heads == 0
        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads
        self.k_layer = nn.Linear(size // 4, num_heads * head_size // 4)
        self.v_layer = nn.Linear(size, num_heads * head_size)
        self.q_layer = nn.Linear(size // 4, num_heads * head_size // 4)
        self.confident_layer = nn.Sequential(nn.Linear(sent_length,
            sent_length), nn.ReLU())
        self.output_layer = nn.Linear(size, size)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k: Tensor, v: Tensor, q: Tensor, mask: Tensor=None,
        xy_pdf_score: Tensor=None):
        batch_size = k.size(0)
        num_heads = self.num_heads
        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)
        k = k.view(batch_size, -1, num_heads, self.head_size // 4).transpose(
            1, 2)
        v = v.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, num_heads, self.head_size // 4).transpose(
            1, 2)
        q = q / math.sqrt(self.head_size)
        scores = torch.matmul(q, k.transpose(2, 3))
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))
        if xy_pdf_score is None:
            attention = self.softmax(scores)
        else:
            xy_pdf_score = self.confident_layer(xy_pdf_score)
            xy_pdf_score = xy_pdf_score.view(batch_size, 1, 1, -1)
            scores_exp = scores.exp()
            new_scores = scores_exp * xy_pdf_score
            new_scores_sum = new_scores.sum(-1)
            attention = new_scores / new_scores_sum.unsqueeze(-1)
        attention = self.dropout(attention)
        context = torch.matmul(attention, v)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1,
            num_heads * self.head_size)
        output = self.output_layer(context)
        return output


class GELU(nn.Module):
    """
    Paper Section 3.4, last paragraph notice that BERT used the GELU instead of RELU
    """

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 
            0.044715 * torch.pow(x, 3))))


class PositionwiseFeedForward(nn.Module):
    """
    Position-wise Feed-forward layer
    Projects to ff_size and then back down to input_size.
    """

    def __init__(self, input_size, ff_size, dropout=0.1):
        """
        Initializes position-wise feed-forward layer.
        :param input_size: dimensionality of the input.
        :param ff_size: dimensionality of intermediate representation
        :param dropout:
        """
        super(PositionwiseFeedForward, self).__init__()
        self.layer_norm = nn.LayerNorm(input_size, eps=1e-06)
        self.pwff_layer = nn.Sequential(nn.Linear(input_size, ff_size),
            GELU(), nn.Dropout(dropout), nn.Linear(ff_size, input_size), nn
            .Dropout(dropout))

    def forward(self, x):
        x_norm = self.layer_norm(x)
        return self.pwff_layer(x_norm) + x


class PositionalEncoding(nn.Module):
    """
    Pre-compute position encodings (PE).
    In forward pass, this adds the position-encodings to the
    input for as many time steps as necessary.
    Implementation based on OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """

    def __init__(self, size: int=0, max_len: int=5000):
        """
        Positional Encoding with maximum length max_len
        :param size:
        :param max_len:
        :param dropout:
        """
        if size % 2 != 0:
            raise ValueError(
                'Cannot use sin/cos positional encoding with odd dim (got dim={:d})'
                .format(size))
        pe = torch.zeros(max_len, size)
        pe.require_grad = False
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, size, 2, dtype=torch.float) * 
            -(math.log(10000.0) / size))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)
        super(PositionalEncoding, self).__init__()
        self.register_buffer('pe', pe)
        self.dim = size

    def forward(self, emb):
        """Embed inputs.
        Args:
            emb (FloatTensor): Sequence of word vectors
                ``(seq_len, batch_size, self.dim)``
        """
        return emb + self.pe[:, :emb.size(1)]


class TransformerEncoderLayer(nn.Module):
    """
    One Transformer encoder layer has a Multi-head attention layer plus
    a position-wise feed-forward layer.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1):
        """
        A single Transformer layer.
        :param size:
        :param ff_size:
        :param num_heads:
        :param dropout:
        """
        super(TransformerEncoderLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.src_src_att = MultiHeadedAttention(num_heads, size, dropout=
            dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    def forward(self, x: Tensor, mask: Tensor) ->Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        """
        x_norm = self.layer_norm(x)
        h = self.src_src_att(x_norm, x_norm, x_norm, mask)
        h = self.dropout(h) + x
        o = self.feed_forward(h)
        return o


class TransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1, src_trg_att: bool=True):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerDecoderLayer, self).__init__()
        self.size = size
        self.src_trg_att = src_trg_att
        self.trg_trg_att = MultiHeadedAttention(num_heads, size, dropout=
            dropout)
        if src_trg_att:
            self.src_trg_att = MultiHeadedAttention(num_heads, size,
                dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.x_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.dec_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: Tensor=None, memory: Tensor=None, src_mask: Tensor
        =None, trg_mask: Tensor=None) ->Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        x_norm = self.x_layer_norm(x)
        h1 = self.trg_trg_att(x_norm, x_norm, x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + x
        if self.src_trg_att:
            h1_norm = self.dec_layer_norm(h1)
            h2 = self.src_trg_att(memory, memory, h1_norm, mask=src_mask)
            o = self.feed_forward(self.dropout(h2) + h1)
        else:
            o = self.feed_forward(h1)
        return o


class CustomTransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, bb_size: int=64, ff_size: int=0,
        num_heads: int=0, dropout: float=0.1):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(CustomTransformerDecoderLayer, self).__init__()
        self.size = size
        self.trg_trg_att = ContMultiHeadedAttention(num_heads, bb_size,
            bb_size, dropout=dropout)
        self.src_trg_att = ContMultiHeadedAttention(num_heads, size, size,
            dropout=dropout)
        self.feed_forward_h1 = PositionwiseFeedForward(bb_size, ff_size=ff_size
            )
        self.feed_forward_h2 = PositionwiseFeedForward(size, ff_size=ff_size)
        self.x_layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.spa_layer_norm = nn.LayerNorm(bb_size, eps=1e-06)
        self.dropout = nn.Dropout(dropout)

    def forward(self, spatial_x: Tensor=None, semantic_x: Tensor=None,
        memory: Tensor=None, src_mask: Tensor=None, trg_mask: Tensor=None
        ) ->Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        spatial_x_norm = self.spa_layer_norm(spatial_x)
        semantic_x_norm = self.x_layer_norm(semantic_x)
        h1 = self.trg_trg_att(spatial_x_norm, spatial_x_norm,
            spatial_x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + spatial_x
        o1 = self.feed_forward_h1(h1)
        o2 = memory[:, 1:, :]
        o = torch.cat((o2, o1), dim=-1)
        return o


class TransformerRefineLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self, size: int=0, ff_size: int=0, num_heads: int=0,
        dropout: float=0.1, sent_length: int=128):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerRefineLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-06)
        self.box_norm = nn.LayerNorm(size // 4, eps=1e-06)
        self.src_src_att = Custom_Attention(num_heads, size, dropout=
            dropout, sent_length=sent_length)
        self.combine_layer = nn.Linear(size + size // 4, size)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    def forward(self, context: Tensor, box: Tensor, mask: Tensor,
        xy_pdf_score: Tensor) ->Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        context: torch.Size([64, 128, 256])
        box: torch.Size([64, 128, 64])
        context_norm: torch.Size([64, 128, 256])
        box_norm: torch.Size([64, 128, 64])
        xy_pdf_score: torch.Size([64, 128])
        """
        context_norm = self.layer_norm(context)
        box_norm = self.box_norm(box)
        h = self.src_src_att(box_norm, context_norm, box_norm, mask,
            xy_pdf_score)
        h = self.dropout(h) + context_norm
        o = self.feed_forward(h)
        return o

# -*- coding: utf-8 -*-

import math
import torch
import torch.nn as nn
from torch import Tensor


# pylint: disable=arguments-differ
class MultiHeadedAttention(nn.Module):
    """
    Multi-Head Attention module from "Attention is All You Need"
    Implementation modified from OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """

    def __init__(self, num_heads: int, size: int, dropout: float = 0.1):
        """
        Create a multi-headed attention layer.
        :param num_heads: the number of heads
        :param size: model size (must be divisible by num_heads)
        :param dropout: probability of dropping a unit
        """
        super(MultiHeadedAttention, self).__init__()

        assert size % num_heads == 0

        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads

        self.k_layer = nn.Linear(size, num_heads * head_size)
        self.v_layer = nn.Linear(size, num_heads * head_size)
        self.q_layer = nn.Linear(size, num_heads * head_size)

        self.output_layer = nn.Linear(size, size)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k: Tensor, v: Tensor, q: Tensor, mask: Tensor = None):
        """
        Computes multi-headed attention.
        :param k: keys   [B, M, D] with M being the sentence length.
        :param v: values [B, M, D]
        :param q: query  [B, M, D]
        :param mask: optional mask [B, 1, M]
        :return:
        """
        batch_size = k.size(0)
        num_heads = self.num_heads

        # project the queries (q), keys (k), and values (v)
        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)

        # reshape q, k, v for our computation to [batch_size, num_heads, ..]
        k = k.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        v = v.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)

        # compute scores
        q = q / math.sqrt(self.head_size)

        # batch x num_heads x query_len x key_len
        scores = torch.matmul(q, k.transpose(2, 3))
        # apply the mask (if we have one)
        # we add a dimension for the heads to it beslow: [B, 1, 1, M]
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))

        # apply attention dropout and compute context vectors.
        attention = self.softmax(scores)
        attention = self.dropout(attention)

        # get context vector (select values with attention) and reshape
        # back to [B, M, D]
        context = torch.matmul(attention, v)
        context = context.transpose(1, 2).contiguous().view(
            batch_size, -1, num_heads * self.head_size)

        output = self.output_layer(context)

        return output

class ContMultiHeadedAttention(nn.Module):
    """
    Multi-Head Attention module from "Attention is All You Need"
    Implementation modified from OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """

    def __init__(self, num_heads: int, size: int, size_v: int, dropout: float = 0.1):
        """
        Create a multi-headed attention layer.
        :param num_heads: the number of heads
        :param size: model size (must be divisible by num_heads)
        :param dropout: probability of dropping a unit
        """
        super(ContMultiHeadedAttention, self).__init__()

        assert size % num_heads == 0

        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads

        self.k_layer = nn.Linear(size, num_heads * head_size)
        self.v_layer = nn.Linear(size_v, num_heads * head_size)
        self.q_layer = nn.Linear(size, num_heads * head_size)

        self.output_layer = nn.Linear(size, size_v)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k: Tensor, v: Tensor, q: Tensor, mask: Tensor = None):
        """
        Computes multi-headed attention.
        :param k: keys   [B, M, D] with M being the sentence length.
        :param v: values [B, M, D]
        :param q: query  [B, M, D]
        :param mask: optional mask [B, 1, M]
        :return:
        """
        batch_size = k.size(0)
        num_heads = self.num_heads

        # project the queries (q), keys (k), and values (v)
        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)

        # reshape q, k, v for our computation to [batch_size, num_heads, ..]
        k = k.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        v = v.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)

        # compute scores
        q = q / math.sqrt(self.head_size)

        # batch x num_heads x query_len x key_len
        scores = torch.matmul(q, k.transpose(2, 3))

        # apply the mask (if we have one)
        # we add a dimension for the heads to it beslow: [B, 1, 1, M]
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))

        # apply attention dropout and compute context vectors.
        attention = self.softmax(scores)
        attention = self.dropout(attention)

        # get context vector (select values with attention) and reshape
        # back to [B, M, D]
        try:
            context = torch.matmul(attention, v)
        except:
            print(attention.shape)
            print(v.shape)
        context = context.transpose(1, 2).contiguous().view(
            batch_size, -1, num_heads * self.head_size)

        output = self.output_layer(context)

        return output
    
class Custom_Attention(nn.Module):
    def __init__(self, num_heads: int, size: int, dropout: float = 0.1, sent_length: int = 128):
        super(Custom_Attention, self).__init__()

        assert size % num_heads == 0

        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads

        self.k_layer = nn.Linear(size//4, num_heads * head_size//4)
        self.v_layer = nn.Linear(size, num_heads * head_size)
        self.q_layer = nn.Linear(size//4, num_heads * head_size//4)
    
        self.confident_layer = nn.Sequential(nn.Linear(sent_length, sent_length),nn.ReLU())
        
        self.output_layer = nn.Linear(size, size)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k: Tensor, v: Tensor, q: Tensor, mask: Tensor = None, xy_pdf_score: Tensor = None):
        batch_size = k.size(0)
        num_heads = self.num_heads

        # project the queries (q), keys (k), and values (v)
        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)

        # reshape q, k, v for our computation to [batch_size, num_heads, ..]
        k = k.view(batch_size, -1, num_heads, self.head_size//4).transpose(1, 2)
        v = v.view(batch_size, -1, num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, num_heads, self.head_size//4).transpose(1, 2)

        # compute scores
        q = q / math.sqrt(self.head_size)

        # batch x num_heads x query_len x key_len
        scores = torch.matmul(q, k.transpose(2, 3))
        # apply the mask (if we have one)
        # we add a dimension for the heads to it below: [B, 1, 1, M]
        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))
#         print(scores.shape) #xy_pdf_score
        # apply attention dropout and compute context vectors.
        if xy_pdf_score is None:
            attention = self.softmax(scores)
        else:
            xy_pdf_score = self.confident_layer(xy_pdf_score)
#             xy_pdf_score = (xy_pdf_score - xy_pdf_score.mean(1).unsqueeze(-1)) + 1.
            xy_pdf_score = xy_pdf_score.view(batch_size,1,1,-1)
    
            scores_exp = scores.exp()
            new_scores = scores_exp * xy_pdf_score
            new_scores_sum = new_scores.sum(-1)
            attention = new_scores/new_scores_sum.unsqueeze(-1)
        
        attention = self.dropout(attention)

        # get context vector (select values with attention) and reshape
        # back to [B, M, D]
        context = torch.matmul(attention, v)
        context = context.transpose(1, 2).contiguous().view(
            batch_size, -1, num_heads * self.head_size)

        output = self.output_layer(context)

        return output

class GELU(nn.Module):
    """
    Paper Section 3.4, last paragraph notice that BERT used the GELU instead of RELU
    """

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))


# pylint: disable=arguments-differ
class PositionwiseFeedForward(nn.Module):
    """
    Position-wise Feed-forward layer
    Projects to ff_size and then back down to input_size.
    """

    def __init__(self, input_size, ff_size, dropout=0.1):
        """
        Initializes position-wise feed-forward layer.
        :param input_size: dimensionality of the input.
        :param ff_size: dimensionality of intermediate representation
        :param dropout:
        """
        super(PositionwiseFeedForward, self).__init__()
        self.layer_norm = nn.LayerNorm(input_size, eps=1e-6)
        self.pwff_layer = nn.Sequential(
            nn.Linear(input_size, ff_size),
            GELU(),
            nn.Dropout(dropout),
            nn.Linear(ff_size, input_size),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        x_norm = self.layer_norm(x)
        return self.pwff_layer(x_norm) + x


# pylint: disable=arguments-differ
class PositionalEncoding(nn.Module):
    """
    Pre-compute position encodings (PE).
    In forward pass, this adds the position-encodings to the
    input for as many time steps as necessary.
    Implementation based on OpenNMT-py.
    https://github.com/OpenNMT/OpenNMT-py
    """
    def __init__(self,
                 size: int = 0,
                 max_len: int = 5000):
        """
        Positional Encoding with maximum length max_len
        :param size:
        :param max_len:
        :param dropout:
        """
        if size % 2 != 0:
            raise ValueError("Cannot use sin/cos positional encoding with "
                             "odd dim (got dim={:d})".format(size))
        pe = torch.zeros(max_len, size)
        pe.require_grad = False

        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp((torch.arange(0, size, 2, dtype=torch.float) *
                              -(math.log(10000.0) / size)))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)  # shape: [1, size, max_len]
        super(PositionalEncoding, self).__init__()
        self.register_buffer('pe', pe)
        self.dim = size

    def forward(self, emb):
        """Embed inputs.
        Args:
            emb (FloatTensor): Sequence of word vectors
                ``(seq_len, batch_size, self.dim)``
        """
        # Add position encodings
        return emb + self.pe[:, :emb.size(1)]


class TransformerEncoderLayer(nn.Module):
    """
    One Transformer encoder layer has a Multi-head attention layer plus
    a position-wise feed-forward layer.
    """

    def __init__(self,
                 size: int = 0,
                 ff_size: int = 0,
                 num_heads: int = 0,
                 dropout: float = 0.1):
        """
        A single Transformer layer.
        :param size:
        :param ff_size:
        :param num_heads:
        :param dropout:
        """
        super(TransformerEncoderLayer, self).__init__()

        self.layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.src_src_att = MultiHeadedAttention(num_heads, size,
                                                dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    # pylint: disable=arguments-differ
    def forward(self, x: Tensor, mask: Tensor) -> Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        """
        x_norm = self.layer_norm(x)
        h = self.src_src_att(x_norm, x_norm, x_norm, mask)
        h = self.dropout(h) + x
        o = self.feed_forward(h)
        return o


class TransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self,
                 size: int = 0,
                 ff_size: int = 0,
                 num_heads: int = 0,
                 dropout: float = 0.1,
                 src_trg_att: bool = True):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerDecoderLayer, self).__init__()
        self.size = size
        self.src_trg_att = src_trg_att
        self.trg_trg_att = MultiHeadedAttention(num_heads, size,
                                                dropout=dropout)
        if src_trg_att:
            self.src_trg_att = MultiHeadedAttention(num_heads, size,
                                                dropout=dropout)

        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)

        self.x_layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.dec_layer_norm = nn.LayerNorm(size, eps=1e-6)

        self.dropout = nn.Dropout(dropout)

    # pylint: disable=arguments-differ
    def forward(self,
                x: Tensor = None,
                memory: Tensor = None,
                src_mask: Tensor = None,
                trg_mask: Tensor = None) -> Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        # decoder/target self-attention

        x_norm = self.x_layer_norm(x)
#         print(x_norm)
#         trg_mask()
#         print(trg_mask)

        h1 = self.trg_trg_att(x_norm, x_norm, x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + x


        if self.src_trg_att:
            # source-target attention
            h1_norm = self.dec_layer_norm(h1)
            h2 = self.src_trg_att(memory, memory, h1_norm, mask=src_mask)
            # final position-wise feed-forward layer
            o = self.feed_forward(self.dropout(h2) + h1)
        else:
            o = self.feed_forward(h1)
        return o

class CustomTransformerDecoderLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self,
                 size: int = 0,
                 bb_size: int = 64,
                 ff_size: int = 0,
                 num_heads: int = 0,
                 dropout: float = 0.1):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(CustomTransformerDecoderLayer, self).__init__()
        self.size = size

        self.trg_trg_att = ContMultiHeadedAttention(num_heads, bb_size, bb_size,
                                                dropout=dropout)
        self.src_trg_att = ContMultiHeadedAttention(num_heads, size, size,
                                                dropout=dropout)

        self.feed_forward_h1 = PositionwiseFeedForward(bb_size, ff_size=ff_size)
        self.feed_forward_h2 = PositionwiseFeedForward(size, ff_size=ff_size)

        self.x_layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.spa_layer_norm = nn.LayerNorm(bb_size, eps=1e-6)

        self.dropout = nn.Dropout(dropout)

    # pylint: disable=arguments-differ
    def forward(self,
                spatial_x: Tensor = None,
                semantic_x: Tensor = None,
                memory: Tensor = None,
                src_mask: Tensor = None,
                trg_mask: Tensor = None) -> Tensor:
        """
        Forward pass of a single Transformer decoder layer.
        :param x: inputs
        :param memory: source representations
        :param src_mask: source mask
        :param trg_mask: target mask (so as to not condition on future steps)
        :return: output tensor
        """
        # decoder/target self-attention

        spatial_x_norm = self.spa_layer_norm(spatial_x) #64
        semantic_x_norm = self.x_layer_norm(semantic_x) #256

        # layout attention

        h1 = self.trg_trg_att(spatial_x_norm, spatial_x_norm, spatial_x_norm, mask=trg_mask)
        h1 = self.dropout(h1) + spatial_x

#         h2 = self.src_trg_att(memory, memory, semantic_x_norm, mask=src_mask)
#         h2 = self.dropout(h2)
        
        o1 = self.feed_forward_h1(h1)
        o2 = memory[:,1:,:]
        o = torch.cat((o2, o1), dim=-1)
        
        return o

class TransformerRefineLayer(nn.Module):
    """
    Transformer decoder layer.
    Consists of self-attention, source-attention, and feed-forward.
    """

    def __init__(self,
                 size: int = 0,
                 ff_size: int = 0,
                 num_heads: int = 0,
                 dropout: float = 0.1,
                 sent_length: int = 128,):
        """
        Represents a single Transformer decoder layer.
        It attends to the source representation and the previous decoder states.
        :param size: model dimensionality
        :param ff_size: size of the feed-forward intermediate layer
        :param num_heads: number of heads
        :param dropout: dropout to apply to input
        """
        super(TransformerRefineLayer, self).__init__()
        self.layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.box_norm = nn.LayerNorm(size//4, eps=1e-6)
        self.src_src_att = Custom_Attention(num_heads, size,
                                                dropout=dropout, sent_length=sent_length)
        self.combine_layer = nn.Linear(size+size//4, size)
        self.feed_forward = PositionwiseFeedForward(size, ff_size=ff_size)
        self.dropout = nn.Dropout(dropout)
        self.size = size

    # pylint: disable=arguments-differ
    def forward(self, context: Tensor, box: Tensor, mask: Tensor, xy_pdf_score: Tensor) -> Tensor:
        """
        Forward pass for a single transformer encoder layer.
        First applies layer norm, then self attention,
        then dropout with residual connection (adding the input to the result),
        and then a position-wise feed-forward layer.
        :param x: layer input
        :param mask: input mask
        :return: output tensor
        context: torch.Size([64, 128, 256])
        box: torch.Size([64, 128, 64])
        context_norm: torch.Size([64, 128, 256])
        box_norm: torch.Size([64, 128, 64])
        xy_pdf_score: torch.Size([64, 128])
        """
        
#         print('context:',context.shape)
#         print('box:',box.shape)
        context_norm = self.layer_norm(context)
#         print('context_norm:',context_norm.shape)
        box_norm = self.box_norm(box)
#         print('box_norm:',box_norm.shape)
        # Calculate xy_pdf_score
#         combine_contex = torch.cat((context_norm, box_norm), dim=-1)
#         combine_contex = self.combine_layer(combine_contex)
        # k, v, q
        h = self.src_src_att(box_norm, context_norm, box_norm, mask, xy_pdf_score)
        h = self.dropout(h) + context_norm
        o = self.feed_forward(h)
        return o
