scalmulop.py

from theano import Op, Apply
from theano.tensor import as_tensor_variable
from theano.scalar import as_scalar

class ScalMulV1(Op):
    __props__ = ('scal',)

    def __init__(self, scal):
        if not isinstance(scal, int):
            raise TypeError('expected an int')
        self.scal = scal

    def make_node(self, x):
        x = as_tensor_variable(x)
        return Apply(self, [x], [x.type()])

    def perform(self, node, inputs, output_storage):
        x = inputs[0]
        z = output_storage[0]
        z[0] = x * self.scal

    def infer_shape(self, node, input_shapes):
        return input_shapes

    def grad(self, inputs, output_grads):
        return [output_grads[0] * self.scal]

    def R_op(self, inputs, eval_points):
        if eval_points[0] is None:
            return eval_points
        return self.grad(inputs, eval_points)


class ScalMulV2(Op):
    __props__ = ()

    def make_node(self, x, scal):
        x = as_tensor_variable(x)
        scal = as_scalar(scal)
        return Apply(self, [x, scal], [x.type()])

    def perform(self, node, inputs, output_storage):
        x = inputs[0]
        scal = inputs[1]
        z = output_storage[0]
        z[0] = x * scal

    def infer_shape(self, node, input_shapes):
        return [input_shapes[0]]

    def grad(self, inputs, output_grads):
        return [output_grads[0] * inputs[1], (inputs[0] * outputs_grads[1]).sum()]

#    def R_op(self, inputs, eval_points):
#        if eval_points[0] is None:
#            return eval_points
#        return self.grad(inputs, eval_points)