## @package dot_product
# Module caffe2.python.layers.dot_product
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from caffe2.python import schema
from caffe2.python.layers.layers import (
    ModelLayer,
)


class PairwiseDotProduct(ModelLayer):

    def __init__(self, model, input_record, output_dim,
                 name='pairwise_dot_product', **kwargs):
        super(PairwiseDotProduct, self).__init__(model, name, input_record, **kwargs)
        assert isinstance(input_record, schema.Struct), (
            "Incorrect input type. Excpected Struct, but received: {0}".
            format(input_record))
        assert 'all_embeddings' in input_record, "all_embeddings is not given."
        all_embeddings = input_record['all_embeddings']
        assert isinstance(all_embeddings, schema.Scalar), (
            "Incorrect input type. Excpected Scalar, but received: {0}".
            format(all_embeddings))
        if 'indices_to_gather' in input_record:
            indices_to_gather = input_record['indices_to_gather']
            assert isinstance(indices_to_gather, schema.Scalar), (
                "Incorrect type of indices_to_gather. "
                "Expected Scalar, but received: {0}".format(indices_to_gather)
            )

        self.all_embeddings = all_embeddings
        self.indices_to_gather = indices_to_gather

        dtype = all_embeddings.field_types()[0].base
        self.output_schema = schema.Scalar(
            (dtype, (output_dim)),
            self.get_next_blob_reference('output')
        )

    def add_ops(self, net):
        Y = net.BatchMatMul(
            [self.all_embeddings(), self.all_embeddings()],
            [self.all_embeddings() + '_matmul'],
            trans_b=1,
        )
        if self.indices_to_gather:
            flattened = net.Flatten(
                Y, Y + '_flatten',
            )
            net.BatchGather(
                [flattened, self.indices_to_gather()],
                self.output_schema(),
            )
        else:
            net.Flatten(Y, self.output_schema())

@@ -722,16 +722,31 @@
  assert 
+(%0A            (
 'all_emb
@@ -773,52 +773,392 @@
 cord
-, %22all_embeddings is not given.%22%0A        all
+) %5E%0A            ('x_embeddings' in input_record and 'y_embeddings' in input_record)%0A        ), (%0A            %22either (all_embeddings) xor (x_embeddings and y_embeddings) %22 +%0A            %22should be given.%22%0A        )%0A        x_embeddings = (%0A            input_record%5B'all_embeddings'%5D if 'all_embeddings' in input_record%0A            else input_record%5B'x_embeddings'%5D%0A        )%0A        y
 _emb
@@ -1167,16 +1167,17 @@
 dings = 
+(
 input_re
@@ -1191,32 +1191,151 @@
 ll_embeddings'%5D%0A
+                        if 'all_embeddings' in input_record%0A                        else input_record%5B'y_embeddings'%5D)%0A
         assert i
@@ -1344,19 +1344,17 @@
 nstance(
-all
+x
 _embeddi
@@ -1401,37 +1401,42 @@
 rrect input type
+ for x
 . Ex
-c
 pected Scalar, b
@@ -1426,32 +1426,49 @@
 xpected Scalar, 
+%22 +%0A            %22
 but received: %7B0
@@ -1474,31 +1474,245 @@
 0%7D%22.
-%0A            
+format(x_embeddings))%0A        if 'y_embeddings' in input_record:%0A            assert isinstance(y_embeddings, schema.Scalar), (%0A                %22Incorrect input type for y. Expected Scalar, %22 +%0A                %22but received: %7B0%7D%22.
 format(
-all
+y
 _emb
@@ -1711,33 +1711,47 @@
 at(y_embeddings)
-)
+%0A            )%0A
 %0A        if 'ind
@@ -2057,33 +2057,36 @@
 )%0A            )%0A
-%0A
+    
         self.all
@@ -2086,90 +2086,184 @@
 elf.
-all_embeddings = all_embeddings%0A        self.indices_to_gather = indices_to_gather
+indices_to_gather = indices_to_gather%0A        else:%0A            self.indices_to_gather = None%0A%0A        self.x_embeddings = x_embeddings%0A        self.y_embeddings = y_embeddings
 %0A%0A  
@@ -2276,19 +2276,17 @@
 dtype = 
-all
+x
 _embeddi
@@ -2519,35 +2519,33 @@
           %5Bself.
-all
+x
 _embeddings(), s
@@ -2544,27 +2544,25 @@
 ngs(), self.
-all
+y
 _embeddings(
@@ -2583,19 +2583,17 @@
   %5Bself.
-all
+x
 _embeddi

extend pairwise dot product for non-equal x & y dimension size