
.. _extending_theano:

****************
Extending Theano
****************

Theano graphs
-------------

- Theano works with symbolic graphs
- Those graphs are bi-partite graphs (graph with 2 types of nodes)
- The 2 types of nodes are Apply and Variable nodes
- Each Apply node has a link to the Op that it executes

Inputs and Outputs are lists of Theano variables

.. image:: ../hpcs2011_tutorial/pics/apply_node.png
    :width: 500 px

Op contract
-----------


.. code-block:: python

    import theano

    class MyOp(theano.Op):
        def make_node(self, *inputs):
            pass

        def __eq__(self, other):
            pass

        def __hash__(self):
            pass

        def __str__(self):
            pass

        # Python implementation:
        def perform(self, node, inputs_storage, output_storage):
            pass

        # C implementation: [see theano web site for other functions]
	def c_code(...):
	# ...
            pass

        # others implementation (pycuda, ...):
        def make_thunk(self, node, storage_map, _, _2):
            pass

        # optional:
        def __init__(self, ...):
            pass

        def grad(self, inputs, g):
            pass

	def R_op(self, inputs, eval_points):
            pass

        def infer_shape(node, (i0_shapes, ...))
            pass

.. ../extending/op.txt

There are 2 mandatory methods. The first is :func:`make_node`. The
second is the one that expresses what computation should be done at run
time. Currently you have 4 possibilities: implement the :func:`perform`
and/or :func:`c_code <Op.c_code>` (and other related :ref:`C functions
<cop>`), or the :func:`make_thunk` method. The ``perform`` method allows you
to easily wrap an existing Python function in Theano. The ``c_code``
and related methods allow you to have your op generate C code and
have Theano compile and link to it. The ``make_thunk`` method will
be called during compilation and should generate a ``thunk``: a
method that when called will do the desired computation. This is
useful if you want to generate code and compile it yourself. For
example, this allow you to use PyCUDA to compile GPU code.

There are 2 mandatory/highly recommended methods. They are needed for a basic
optimization that merges duplicate computations in a Theano function. Thus,
if you don't want Theano to perform your computations multiple times for no
good reason, implement these! Those methods are :func:`__eq__` and
:func:`__hash__`.

The :func:`infer_shape` method allows for some very interesting
optimizations, such as not performing your op's computations simply to
determine the shape your Op's output.

The :func:`grad` method is needed if you want symbolic differentiation to
work with your Op.

The :func:`__str__` is useful in order to provide a more meaningful string
representation of your Op.

The :func:`R_op` is needed if you want `theano.tensor.Rop` to work with your op.

Op example
----------

.. code-block:: python

    import theano

    class DoubleOp(theano.Op):
        def __eq__(self, other):
            return type(self) == type(other)

        def __hash__(self):
            return hash(type(self))

        def __str__(self):
            return self.__class__.__name__

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

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

Test it!

>>> x = theano.tensor.matrix()
>>> f = theano.function([x],DoubleOp()(x))
>>> import numpy
>>> inp = numpy.random.rand(5,5)
>>> out = f(inp)
>>> assert numpy.allclose(inp*2, out)
>>> print inp
>>> print out


Exercises 8
-----------

- Run the code in the file double_op.py.
- Modify and execute to compute: x * y
- Modify and execute the example to return 2 outputs: x + y and x - y

  - Our current element-wise fusion generates computation with only 1 output.



