Forward of Decomposable Attention layer

def hybrid_forward(self, F, a, b):
        """
        Forward of Decomposable Attention layer
        """
        # a.shape = [B, L1, H]
        # b.shape = [B, L2, H]
        # extract features
        tilde_a = self.f(a)  # shape = [B, L1, H]
        tilde_b = self.f(b)  # shape = [B, L2, H]
        # attention
        # e.shape = [B, L1, L2]
        e = F.batch_dot(tilde_a, tilde_b, transpose_b=True)
        # beta: b align to a, [B, L1, H]
        beta = F.batch_dot(e.softmax(), tilde_b)
        # alpha: a align to b, [B, L2, H]
        alpha = F.batch_dot(e.transpose([0, 2, 1]).softmax(), tilde_a)
        # compare
        feature1 = self.g(F.concat(tilde_a, beta, dim=2))
        feature2 = self.g(F.concat(tilde_b, alpha, dim=2))
        feature1 = feature1.sum(axis=1)
        feature2 = feature2.sum(axis=1)
        yhat = self.h(F.concat(feature1, feature2, dim=1))
        return yhat