Hetu/python/hetu/gpu_ops/Conv2dBroadcast.py

from __future__ import absolute_import
import numpy as np
from .Node import Op
from .Conv2dReduceSum import conv2d_reducesum_op
from .ZerosLike import zeroslike_op
from ..gpu_links import broadcast_to


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

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            shapeW = input_vals[1].shape
            shapeW = list(shapeW)
            tmp = shapeW[1]
            shapeW[1] = shapeW[3]
            shapeW[3] = tmp
            output_val[:] = np.broadcast_to(
                input_vals[0].asnumpy(), input_vals[1].asnumpy().shape).swapaxes(1, 3)
        else:
            broadcast_to(input_vals[0], output_val, stream_handle)

    def gradient(self, output_grad):

        grad_A = conv2d_reducesum_op(output_grad, ctx=self.raw_ctx)
        return [grad_A, None]

    def infer_shape(self, input_shapes):
        assert len(input_shapes) == 2
        return input_shapes[1]


def conv2d_broadcastto_op(node_A, node_B, ctx=None):
    """Creates a node that represents np.broadcast_to(node_A, node_B.shape).

    Parameters:
    ----
    node_a : Node
        The Node to be bcast.
    node_b : Node
        Another Node with the target shape.

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

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