Merge pull request #1449 from nouiz/mpi

doc/library/tensor/index.txt

@@ -24,3 +24,4 @@ They are grouped into the following sections:
     signal/index
     utils
     extra_ops
+    io

doc/library/tensor/io.txt

+===================================================================
+:mod:`tensor.io` --  Tensor IO Ops
+===================================================================
+
+.. module:: tensor.io
+   :platform: Unix, Windows
+   :synopsis: Tensor IO Ops
+.. moduleauthor:: LISA
+
+File operation
+==============
+
+- Load from disk with the function :func:`load <theano.tensor.io.load>` and its associated op :class:`LoadFromDisk <theano.tensor.io.LoadFromDisk>`
+
+MPI operation
+=============
+- Non-blocking transfer: :func:`isend <theano.tensor.io.isend>` and :func:`irecv <theano.tensor.io.irecv>`.
+- Blocking transfer: :func:`send <theano.tensor.io.send>` and :func:`recv <theano.tensor.io.recv>`
+
+Details
+=======
+
+.. automodule:: theano.tensor.io
+    :members:
+

theano/tensor/io.py

@@ -52,6 +52,7 @@ class LoadFromDisk(Op):
     def __str__(self):
         return "Load{dtype: %s, broadcastable: %s, mmep: %s}" % self._info
 
+
 def load(path, dtype, broadcastable, mmap_mode=None):
     """
     Load an array from an .npy file.
@@ -91,6 +92,7 @@ else:
     comm = MPI.COMM_WORLD
     mpi_enabled = True
 
+
 class MPIRecv(Op):
     """
     An operation to asynchronously receive an array to a remote host using MPI
@@ -120,6 +122,7 @@ class MPIRecv(Op):
         return gof.Apply(self, [], [theano.Variable(Generic()),
                                     tensor(self.dtype,
                                            broadcastable=self.broadcastable)])
+
     def perform(self, node, inp, out):
 
         data = numpy.zeros(self.shape, dtype=self.dtype)
@@ -137,6 +140,7 @@ class MPIRecv(Op):
     def do_constant_folding(self, node):
         return False
 
+
 class MPIRecvWait(Op):
     """
     An operation to wait on a previously received array using MPI
@@ -160,6 +164,7 @@ class MPIRecvWait(Op):
         return gof.Apply(self, [request, data],
                                [tensor(data.dtype,
                                        broadcastable=data.broadcastable)])
+
     def perform(self, node, inp, out):
 
         request = inp[0]
@@ -177,6 +182,7 @@ class MPIRecvWait(Op):
 
     view_map = {0: [1]}
 
+
 class MPISend(Op):
     """
     An operation to asynchronously Send an array to a remote host using MPI
@@ -216,6 +222,7 @@ class MPISend(Op):
     def __str__(self):
         return "MPISend{dest: %d, tag: %d}" % self._info
 
+
 class MPISendWait(Op):
     """
     An operation to wait on a previously sent array using MPI
@@ -247,18 +254,35 @@ class MPISendWait(Op):
     def __str__(self):
         return "MPISendWait"
 
+
 def isend(var, dest, tag):
+    """
+    Non blocking send
+    """
     return MPISend(dest, tag)(var)
 
+
 def send(var, dest, tag):
+    """
+    blocking send
+    """
     return MPISendWait(tag)(*isend(var, dest, tag))
 
+
 def irecv(shape, dtype, source, tag):
+    """
+    non-blocking receive
+    """
     return MPIRecv(source, tag, shape, dtype)()
 
+
 def recv(shape, dtype, source, tag):
+    """
+    blocking receive
+    """
     return MPIRecvWait(tag)(*irecv(shape, dtype, source, tag))
 
+
 # Ordering keys for scheduling
 def mpi_send_wait_key(a):
     """ Wait as long as possible on Waits, Start Send/Recvs early """
@@ -268,6 +292,7 @@ def mpi_send_wait_key(a):
         return -1
     return 0
 
+
 def mpi_tag_key(a):
     """ Break MPI ties by using the variable tag - prefer lower tags first """
     if isinstance(a.op, (MPISend, MPIRecv, MPIRecvWait, MPISendWait)):