Hetu/python/hetu/gpu_ops/Where.py

from __future__ import absolute_import
import numpy as np
from .Node import Op
from ..gpu_links import where


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

    def compute(self, input_vals, output_val, stream_handle=None):
        if self.on_cpu:
            output_val[:] = np.where(input_vals[0].asnumpy(
            ), input_vals[1].asnumpy(), input_vals[2].asnumpy())
        else:

            where(input_vals[0], input_vals[1],
                  input_vals[2], output_val, stream_handle)

    def gradient(self, output_grad):
        from .ZerosLike import zeroslike_op
        zeros = zeroslike_op(self.inputs[0], ctx=self.raw_ctx)
        grad_A = where_op(self.inputs[0], output_grad, zeros, ctx=self.raw_ctx)
        grad_B = where_op(self.inputs[0], zeros, output_grad, ctx=self.raw_ctx)
        return [None, grad_A, grad_B]

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


def where_op(cond, node_A, node_B, ctx=None):
    """Creates a node that represents np.where.

    Parameters:
    ----
    cond : Node of a condition array
    node_A : Node, output if cond
    node_B : Node, output if not cond

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

    """
    return WhereOp(cond, node_A, node_B, ctx=ctx)