Hetu/python/hetu/gpu_ops/MatrixDot.py

from __future__ import absolute_import
from .Node import Op
import numpy as np
from .MultiplyElewise import mul_op
from .ReduceSum import reduce_sum_op
from ..gpu_links import matrix_dot

# TODO: This op may have bugs and is not complete!
# Use other ops to replace it


class MatrixDotOp(Op):
    def __init__(self, node_A, node_B, axes=0, ctx=None):
        super().__init__(MatrixDotOp, [node_A, node_B], ctx)
        self.axes = axes

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            output_val[:] = np.tensordot(
                input_vals[0], input_vals[1], axes=self.axes)
        else:
            matrix_dot(input_vals[0], input_vals[1], output_val, stream_handle)

    def gradient(self, output_grad):
        return [matrix_dot_op(output_grad, self.inputs[1], axes=1, ctx=self.raw_ctx),
                reduce_sum_op(mul_op(self.inputs[0], output_grad, ctx=self.raw_ctx), axes=1, keepdims=True, ctx=self.raw_ctx)]

    def infer_shape(self, input_shapes):
        """Need to handle input_vals[0].shape != input_vals[1].shape"""
        return input_shapes[0]


def matrix_dot_op(node_A, node_B, axes=0, ctx=None):
    """Make a new instance of matrixs elementwise multiplication and call the instance.

    Parameters:
    ----
    node_a : Node
        The Node to be multiplied.
    node_b : Node
        Another Node to be multiplied.

    Returns:
    ----
    A new Node instance created by Op.

    """
    return MatrixDotOp(node_A, node_B, ctx=ctx)