Merge pull request #1785 from abergeron/op_decorator

doc/library/compile/index.txt

@@ -16,6 +16,7 @@
     shared
     function
     io
+    ops
     mode
     module
     debugmode

doc/library/compile/ops.txt

+==================================================
+:mod:`ops` --  Some Common Ops and extra Ops stuff
+==================================================
+
+.. automodule:: theano.compile.ops
+    :members:

theano/compile/__init__.py

@@ -2,7 +2,8 @@ from theano.compile.ops import (
         DeepCopyOp, deep_copy_op, register_deep_copy_op_c_code,
         Shape, shape, register_shape_c_code,
         Shape_i, register_shape_i_c_code,
-        ViewOp, view_op, register_view_op_c_code)
+        ViewOp, view_op, register_view_op_c_code, FromFunctionOp, 
+        as_op)
 
 from theano.compile.function_module import *
 

theano/compile/ops.py

-"""This file contain auxiliary Ops, used during the compilation phase."""
+"""This file contain auxiliary Ops, used during the compilation phase
+and Ops building class (:class:`FromFunctionOp`) and decorator
+(:func:`as_op`) that help make new Ops more rapidly.
+
+"""
 import copy
 import warnings
 
@@ -364,3 +368,96 @@ def register_shape_i_c_code(typ, code, version=()):
 # List of Theano Types that one can add an extra dimension and for which
 # Scan can deal with.
 expandable_types = ()
+
+class FromFunctionOp(gof.Op):
+    """
+    Build a basic Theano Op around a function.
+
+    Since the resulting Op is very basic and is missing most of the
+    optional functionality, some optimization may not apply.  If you
+    want to help, you can supply an infer_shape function that computes
+    the shapes of the output given the shapes of the inputs.
+
+    Also the gradient is undefined in the resulting op and Theano will
+    raise an error if you attempt to get the gradient of a graph
+    containing this op.
+    """
+    def __init__(self, fn, itypes, otypes, infer_shape):
+        self.__fn = fn
+        self.itypes = itypes
+        self.otypes = otypes
+        self.__infer_shape = infer_shape
+        if self.__infer_shape is not None:
+            self.infer_shape = self._infer_shape
+
+    def __eq__(self, other):
+        return (type(self) == type(other) and
+                self.__fn == other.__fn)
+
+    def __hash__(self):
+        return hash(type(self)) ^ hash(self.__fn)
+
+    def __str__(self):
+        return 'FromFunctionOp{%s}' % self.__fn.__name__
+
+    def make_node(self, *inputs):
+        assert len(inputs) == len(self.itypes)
+        assert all(inp.type == it for inp, it in zip(inputs, self.itypes))
+        return theano.Apply(self, inputs, [o() for o in self.otypes])
+
+    def perform(self, node, inputs, outputs):
+        outs = self.__fn(*inputs)
+        if not isinstance(outs, (list, tuple)):
+            outs = (outs,)
+        assert len(outs) == len(outputs)
+        for i in range(len(outs)):
+            outputs[i][0] = outs[i]
+
+    def _infer_shape(self, node, input_shapes):
+        return self.__infer_shape(node, input_shapes)
+
+def as_op(itypes, otypes, infer_shape=None):
+    """
+    Decorator that converts a function into a basic Theano op that
+    will call the supplied function as its implementation.
+
+    It takes an optional infer_shape parameter that should be a
+    callable with this signature:
+
+        def infer_shape(node, input_shapes):
+            ...
+            return output_shapes
+
+    Here `input_shapes` and `output_shapes` are lists of tuples that
+    represent the shape of the corresponding inputs/outputs.
+
+    This should not be used when performance is a concern since the
+    very basic nature of the resulting Op may interfere with certain
+    graph optimizations.
+
+    Example usage:
+
+       @as_op(itypes=[theano.tensor.fmatrix, theano.tensor.fmatrix],
+              otypes=[theano.tensor.fmatrix])
+       def numpy_dot(a, b):
+           return numpy.dot(a, b)
+    """
+    if not isinstance(itypes, (list, tuple)):
+        itypes = [itypes]
+    if any(not isinstance(t, theano.Type) for t in itypes):
+        raise TypeError("itypes has to be a list of Theano types")
+    if not isinstance(otypes, (list, tuple)):
+        otypes = [otypes]
+    if any(not isinstance(t, theano.Type) for t in otypes):
+        raise TypeError("otypes has to be a list of Theano types")
+
+    # make sure they are lists and not tuples
+    itypes = list(itypes)
+    otypes = list(otypes)
+
+    if infer_shape is not None and not callable(infer_shape):
+        raise TypeError("infer_shape needs to be a callable")
+
+    def make_op(fn):
+        return FromFunctionOp(fn, itypes, otypes, infer_shape)
+    return make_op

theano/compile/tests/test_ops.py

+"""
+ Tests for the Op decorator
+"""
+import numpy as np
+
+from theano.tests import unittest_tools as utt
+from theano import function
+import theano
+from theano import tensor
+from theano.tensor import dmatrix, dvector
+from numpy import allclose
+from theano.compile import as_op
+
+
+class OpDecoratorTests(utt.InferShapeTester):
+    def test_1arg(self):
+        x = dmatrix('x')
+
+        @as_op(dmatrix, dvector)
+        def diag(x):
+            return np.diag(x)
+
+        fn = function([x], diag(x))
+        r = fn([[1.5, 5], [2, 2]])
+        r0 = np.array([1.5, 2])
+
+        assert allclose(r, r0), (r, r0)
+
+    def test_2arg(self):
+        x = dmatrix('x')
+        x.tag.test_value = np.zeros((2, 2))
+        y = dvector('y')
+        y.tag.test_value = [0, 0]
+
+        @as_op([dmatrix, dvector], dvector)
+        def diag_mult(x, y):
+            return np.diag(x) * y
+
+        fn = function([x, y], diag_mult(x, y))
+        r = fn([[1.5, 5], [2, 2]], [1, 100])
+        r0 = np.array([1.5, 200])
+
+        assert allclose(r, r0), (r, r0)
+
+    def test_infer_shape(self):
+        x = dmatrix('x')
+        x.tag.test_value = np.zeros((2, 2))
+        y = dvector('y')
+        y.tag.test_value = [0, 0]
+
+        def infer_shape(node, shapes):
+            x, y = shapes
+            return [y]
+
+        @as_op([dmatrix, dvector], dvector, infer_shape)
+        def diag_mult(x, y):
+            return np.diag(x) * y
+
+        self._compile_and_check([x, y], [diag_mult(x, y)],
+                                [[[1.5, 5], [2, 2]], [1, 100]],
+                                diag_mult.__class__, warn=False)