Add a contexts config variable to specify the context map.

theano/__init__.py

@@ -112,7 +112,8 @@ if config.device.startswith('gpu') or config.init_gpu_device.startswith('gpu'):
 if (config.device.startswith('cuda') or
         config.device.startswith('opencl') or
         config.init_gpu_device.startswith('cuda') or
-        config.init_gpu_device.startswith('opencl')):
+        config.init_gpu_device.startswith('opencl') or
+        config.contexts != ''):
     import theano.sandbox.gpuarray
 
 # Use config.numpy to call numpy.seterr

theano/configdefaults.py

@@ -111,6 +111,29 @@ AddConfigVar(
     BoolParam(False, allow_override=False),
     in_c_key=False)
 
+
+class ContextsParam(ConfigParam):
+    def __init__(self):
+        def filter(val):
+            if val == '':
+                return val
+            for v in val.split(';'):
+                s = v.split('->')
+                if len(s) != 2:
+                    raise ValueError("Malformed context map: %s" % (v,))
+            return val
+        ConfigParam.__init__(self, '', filter, False)
+
+AddConfigVar(
+    'contexts',
+    """
+    Context map for multi-gpu operation. Format is a
+    semicolon-separated list of names and device names in the
+    'name->dev_name' format. An example that would map name 'test' to
+    device 'cuda0' and name 'test2' to device 'opencl0:0' follows:
+    "test->cuda0;test2->opencl0:0".
+    """, ContextsParam(), in_c_key=False)
+
 AddConfigVar(
     'print_active_device',
     "Print active device at when the GPU device is initialized.",

theano/sandbox/gpuarray/__init__.py

@@ -51,6 +51,12 @@ if pygpu:
         elif (config.init_gpu_device.startswith('cuda') or
               config.init_gpu_device.startswith('opencl')):
             init_dev(config.init_gpu_device)
+        if config.contexts != '':
+            for n, d in (c.split('->') for c in config.contexts.split(';')):
+                init_dev(d, n)
+            import theano.compile
+            theano.compile.shared_constructor(gpuarray_shared_constructor)
+            optdb.add_tags('gpuarray_opt', 'fast_run', 'fast_compile')
 
         from .basic_ops import (GpuAlloc, GpuContiguous, GpuEye, GpuFromHost,
                                 GpuJoin, GpuReshape, GpuSplit, HostFromGpu)
@@ -65,5 +71,6 @@ else:
     if (config.init_gpu_device.startswith('cuda') or
             config.init_gpu_device.startswith('opencl') or
             config.device.startswith('opencl') or
-            config.device.startswith('cuda')):
+            config.device.startswith('cuda') or
+            config.contexts != ''):
         error("pygpu was configured but could not be imported", exc_info=True)

theano/sandbox/gpuarray/opt.py

@@ -170,8 +170,8 @@ class InputToGpuOptimizer(Optimizer):
                 continue
 
             try:
-                ctx = getattr(input.tag, 'context_name', None)
-                new_input = host_from_gpu(GpuFromHost(ctx)(input))
+                ctx_name = getattr(input.tag, 'context_name', None)
+                new_input = host_from_gpu(GpuFromHost(ctx_name)(input))
                 fgraph.replace_validate(input, new_input,
                                         "InputToGpuOptimizer")
             except TypeError:
@@ -180,7 +180,8 @@ class InputToGpuOptimizer(Optimizer):
             except ValueError:
                 # If there is no context tag and no default context
                 # then it stays on the CPU
-                assert ctx is None
+                if ctx is not None:
+                    raise
                 pass
 
 

theano/sandbox/gpuarray/type.py

@@ -469,6 +469,15 @@ def gpuarray_shared_constructor(value, name=None, strict=False,
     if not isinstance(value, (numpy.ndarray, pygpu.gpuarray.GpuArray)):
         raise TypeError('ndarray or GpuArray required')
 
+    try:
+        get_context(context_name)
+    except ValueError:
+        # Don't make this a hard error if we attempt to make a shared
+        # variable while there is no default context.
+        if context_name is None:
+            raise TypeError('No default context and no context specified')
+        raise
+
     if broadcastable is None:
         broadcastable = (False,) * value.ndim
     type = GpuArrayType(value.dtype, broadcastable)