Hetu/python/hetu/gpu_ops/Broadcast.py

from __future__ import absolute_import
import numpy as np
from .Node import Op
from .ReduceSum import reduce_sum_op
from ..gpu_links import broadcast_shape_simple
from .. import ndarray


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

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            input_shape = list(input_vals[1].shape)
            output_val[:] = np.broadcast_to(
                input_vals[0].asnumpy(), input_shape)
        else:
            if self.inplace:
                input_vals[0].broadcast_to(input_vals[1].shape, output_val)
            else:
                # broadcast_shape(input_vals[0], output_val, None, stream_handle)
                broadcast_shape_simple(
                    input_vals[0], output_val, self.out_strides, self.in_dims, stream_handle)

    def gradient(self, output_grad):
        self.grad_node = reduce_sum_op(
            output_grad, None, None, ctx=self.raw_ctx)
        return [self.grad_node, None]

    def infer_shape(self, input_shapes):
        assert len(input_shapes) == 2
        input_shape = list(input_shapes[0])
        output_shape = list(input_shapes[1])
        output_ndim = len(output_shape)
        assert len(input_shape) <= output_ndim
        diff = output_ndim - len(input_shape)
        axes = list(range(diff))
        keepdims = [False] * diff
        input_shape = [1] * diff + input_shape
        for i in range(output_ndim):
            assert output_shape[i] > 0 and isinstance(output_shape[i], int)
            assert input_shape[i] == 1 or input_shape[i] == output_shape[i]
            if i >= diff and input_shape[i] == 1 and output_shape[i] > 1:
                axes.append(i)
                keepdims.append(True)
        if hasattr(self, 'grad_node'):
            self.grad_node.axes = axes
            self.grad_node.keepdims = keepdims

        # here we save the output strides and input dimensions for GPU computation
        if self.on_gpu and not self.inplace:
            input_shape = list(input_shapes[0])
            out_strides = [0 for _ in range(output_ndim)]
            temp_size = 1
            for i in range(output_ndim - 1, -1, -1):
                out_strides[i] = temp_size
                temp_size *= output_shape[i]
            in_dims = [1 for _ in range(diff)] + input_shape

            self.out_strides = ndarray.array(
                out_strides, self.ctx, data_type=np.int32)
            self.in_dims = ndarray.array(in_dims, self.ctx, data_type=np.int32)
        return input_shapes[1]

    def backward_hook(self, config):
        self.inplace = config.enable_lazy and self not in config.eval_node_list


def 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 broadcast.
    node_b : Node
        Another Node with the target shape.

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

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