Conceptual example of how to use template ops

template_op

class TemplateOp(theano.Op):
    @property
    def __impls__(self):
        return {}

    def __call__(self, *inputs, **kwargs):
        tag = kwargs.pop('tag', None)
        if tag is None:
            return super(TemplateOp, self).__call__(*inputs, **kwargs)

        try:
            return self.__impls__[tag](*inputs)
        except KeyError:
            raise KeyError('Tag "%s" has not been registered.' % tag)

###########################################################################

class Dropout(TemplateOp):
    def __init__(self, drop=0.5, rng=theano_rng):
        drop = float(drop)
        
        if drop <= 0 or drop >= 1:
            raise ValueError('`drop` must be in the range (0,1).')

        self._drop = drop
        self._rng = rng
    
    @property
    def __impls__(self):
        return {'train': self.__fit, 'evaluate': self.__evaluate}

    def __train(self, x):
        p = 1.0 - self._drop
        mask = self._rng.binomial(x.shape, p=p, dtype=x.dtype, nstreams=1)
        return mask * x

    def __evaluate(self, x):
        p = 1.0 - self._drop
        return p * x


X = T.matrix('X')

drop = Dropout()

# or alternatively we could define a function
#
# def dropout(x, p=0.5, rng=theano_rng, tag=None):
#     return Dropout(p=p, rng=rng)(x, tag=tag)

Xd = drop(X)

# or if we want to force a specific implementation
#
# Xd = drop(X, tag='evaluate')
#
# or using the above defined function
#
# Xd = dropout(x, p=0.5, tag='evaluate')

f_train = theano.function([X], Xd, tags=['train'])
f_evaluate = theano.function([X], Xd, tags=['evaluate'])