import torch


def calc_gated_attention(x, c, gate, rnn, proj=None):
    """
    Processes attention output by a gate and an RNN.

    Args:
        x: input vectors
        c: attention output
        gate: an element-wise multiplicative gate
        rnn: an RNN (or variant thereof) encoder
        proj: a linear layer or None

    Returns:
        attention output passed through a gate and processed by an RNN
    """
    rnn_in = torch.cat([x, c], dim=2)       # (batch_size, seq_len, 2 * input_size)
    rnn_in = rnn_in * gate(rnn_in)      # (batch_size, seq_len, 2 * input_size)
    h, _ = rnn(rnn_in)  # (batch_size, seq_len, 2 * hidden_size)
    return h if proj is None else proj(h)