Add functions for the new backend.

theano/compile/nanguardmode.py

@@ -11,6 +11,14 @@ import theano.tensor as T
 import theano.sandbox.cuda as cuda
 from theano.compile import Mode
 
+try:
+    from theano.gpuarray.type import GpuArrayType, _name_for_ctx
+    from pygpu.gpuarray import GpuArray
+    pygpu_available = True
+except ImportError:
+    pygpu_available = False
+
+
 logger = logging.getLogger("theano.compile.nanguardmode")
 
 
@@ -86,6 +94,8 @@ def contains_nan(arr, node=None, var=None):
         else:
             compile_gpu_func(True, False, False)
             return np.isnan(f_gpumin(arr.reshape(arr.size)))
+    elif pygpu_available and isinstance(arr, GpuArray):
+        return np.isnan(f_gpua_min(arr.reshape(arr.size)))
 
     return np.isnan(np.min(arr))
 
@@ -136,6 +146,9 @@ def contains_inf(arr, node=None, var=None):
             compile_gpu_func(False, True, False)
             return (np.isinf(f_gpumin(arr.reshape(arr.size))) or
                     np.isinf(f_gpumax(arr.reshape(arr.size))))
+    elif pygpu_available and isinstance(arr, GpuArray):
+        return (np.isinf(f_gpua_min(arr.reshape(arr.size))) or
+                np.isinf(f_gpua_max(arr.reshape(arr.size))))
 
     return np.isinf(np.nanmax(arr)) or np.isinf(np.nanmin(arr))
 
@@ -187,6 +200,47 @@ def compile_gpu_func(nan_is_error, inf_is_error, big_is_error):
             cuda_compile_failed = True
 
 
+def f_gpua_min(inp):
+    dt = inp.dtype
+    ctx_name = _name_for_ctx(inp.context)
+    key = (dt, ctx_name)
+    f = f_gpua_min.cache.get(key, None)
+    if f is None:
+        guard_in = GpuArrayType(str(dt), (False,), context_name=ctx_name)()
+        f = theano.function([guard_in], T.min(guard_in),
+                            mode='FAST_RUN', profile=False)
+        f_gpua_min.cache[key] = f
+    return f(inp)
+f_gpua_min.cache = dict()
+
+def f_gpua_max(inp):
+    dt = inp.dtype
+    ctx_name = _name_for_ctx(inp.context)
+    key = (dt, ctx_name)
+    f = f_gpua_min.cache.get(key, None)
+    if f is None:
+        guard_in = GpuArrayType(str(dt), (False,), context_name=ctx_name)()
+        f = theano.function([guard_in], T.max(guard_in),
+                            mode='FAST_RUN', profile=False)
+        f_gpua_min.cache[key] = f
+    return f(inp)
+f_gpua_max.cache = dict()
+
+
+def f_gpua_absmax(inp):
+    dt = inp.dtype
+    ctx_name = _name_for_ctx(inp.context)
+    key = (dt, ctx_name)
+    f = f_gpua_min.cache.get(key, None)
+    if f is None:
+        guard_in = GpuArrayType(str(dt), (False,), context_name=ctx_name)()
+        f = theano.function([guard_in], T.max(T.abs_(guard_in)),
+                            mode='FAST_RUN', profile=False)
+        f_gpua_min.cache[key] = f
+    return f(inp)
+f_gpua_absmax.cache = dict()
+
+
 class NanGuardMode(Mode):
     """
     A Theano compilation Mode that makes the compiled function automatically
@@ -262,7 +316,10 @@ class NanGuardMode(Mode):
                 elif value.size == 0:
                     err = False
                 elif cuda.cuda_available and isinstance(value, cuda.CudaNdarray):
+                    compile_gpu_func(False, False, True)
                     err = (f_gpuabsmax(value.reshape(value.size)) > 1e10)
+                elif pygpu_available and isinstance(value, GpuArray):
+                    err = (f_gpua_absmax(value.reshape(value.size)) > 1e10)
                 else:
                     err = (np.abs(value).max() > 1e10)
                 if err: