Hetu/python/hetu/gpu_ops/BroadcastShape.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 BroadcastShapeOp(Op):
    def __init__(self, node_A, shape, add_axes=(), ctx=None):
        super().__init__(BroadcastShapeOp, [node_A], ctx)
        self.target_shape = shape
        self.add_axes = add_axes

    def compute(self, input_vals, output_val, stream_handle=None):
        assert self.target_shape is not None and self.add_axes is not None
        if self.on_cpu:
            input_shape = list(input_vals[0].shape)
            for i in range(len(input_shape)):
                if self.add_axes and i in self.add_axes:
                    input_shape[i] = 1
            output_val[:] = np.broadcast_to(
                input_vals[0].asnumpy().reshape(input_shape), self.target_shape)
        else:
            if self.inplace:
                input_vals[0].broadcast_to(
                    self.target_shape, output_val, self.add_axes)
            else:
                # broadcast_shape(input_vals[0], output_val, self.add_axes, 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]

    def infer_shape(self, input_shapes):
        assert self.target_shape is not None and self.add_axes is not None
        assert len(input_shapes) == 1
        input_shape = list(input_shapes[0])
        output_shape = list(self.target_shape)
        output_ndim = len(output_shape)
        assert len(input_shape) <= output_ndim
        diff = output_ndim - len(input_shape)
        if self.add_axes:
            assert diff == len(self.add_axes) or input_shape == [1]
            assert all([axis < output_ndim for axis in self.add_axes])
            in_ind = 0
            for i in range(output_ndim):
                if i not in self.add_axes:
                    assert input_shape[in_ind] == output_shape[i]
                    in_ind += 1
            if hasattr(self, 'grad_node'):
                self.grad_node.axes = tuple(self.add_axes)
                self.grad_node.axes.keepdims = [False] * len(self.add_axes)
        else:
            axes = list(range(diff))
            keepdims = [False] * diff
            input_shape = [1] * diff + input_shape
            for i in range(output_ndim):
                if output_shape[i] == -1:
                    output_shape[i] = input_shape[i]
                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]
            if self.add_axes:
                in_dims = [0 for _ in range(output_ndim)]
                for i in range(diff):
                    in_dims[self.add_axes[i]] = 1
                temp_ind = 0
                for dim in input_shape:
                    while in_dims[temp_ind] == 1:
                        temp_ind += 1
                    in_dims[temp_ind] = dim
            else:
                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 tuple(output_shape)

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


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

    Parameters:
    ----
    node_a : Node
        The Node to be broadcast.
    shape : tuple
        Target shape.

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

    """
    return BroadcastShapeOp(node_A, shape, add_axes=add_axes, ctx=ctx)