code: def hybrid_forward(self, F, a, b):
          
          # 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
docs: Forward of Decomposable Attention layer