use is_odd input to irfft op

theano/gpuarray/fft.py

@@ -38,7 +38,7 @@ class CuRFFTOp(Op):
                             broadcastable=[False] * (inp.type.ndim + 1),
                             context_name=inp.type.context_name)
 
-    def make_node(self, inp, s):
+    def make_node(self, inp):
         if not scikits_cuda_available:
             raise RuntimeError("scikits.cuda is needed for CuFFTOp")
 
@@ -51,13 +51,10 @@ class CuRFFTOp(Op):
         inp = basic_ops.gpu_contiguous(
             basic_ops.as_gpuarray_variable(inp,
                                            basic_ops.infer_context_name(inp)))
-        s = T.as_tensor_variable(s)
 
         assert inp.dtype == "float32"
-        assert s.ndim == 1
-        assert 'int' in s.dtype
 
-        return theano.Apply(self, [inp, s], [self.output_type(inp)()])
+        return theano.Apply(self, [inp], [self.output_type(inp)()])
 
     def make_thunk(self, node, storage_map, _, _2):
 
@@ -73,12 +70,9 @@ class CuRFFTOp(Op):
 
         def thunk():
             input_shape = inputs[0][0].shape
-            s = inputs[1][0]
-
-            assert (input_shape[1:] == s).all()
 
             # construct output shape
-            output_shape = [input_shape[0]] + list(s)
+            output_shape = list(input_shape)
             # DFT of real input is symmetric, no need to store
             # redundant coefficients
             output_shape[-1] = output_shape[-1] // 2 + 1
@@ -105,7 +99,7 @@ class CuRFFTOp(Op):
                 # only initialise plan if necessary
                 if plan[0] is None or plan_input_shape[0] != input_shape:
                     plan_input_shape[0] = input_shape
-                    plan[0] = fft.Plan(s, np.float32, np.complex64,
+                    plan[0] = fft.Plan(input_shape[1:], np.float32, np.complex64,
                                        batch=input_shape[0])
                 # Sync GPU variables before computation
                 input_pycuda.sync()
@@ -123,17 +117,15 @@ class CuRFFTOp(Op):
 
     def grad(self, inputs, output_grads):
         gout, = output_grads
-        s = inputs[1]
+        s = inputs[0].shape[1:]
+        is_odd = s[-1] % 2
         # Divide the last dimension of the output gradients by 2, they are 
         # double-counted by the real-IFFT due to symmetry, except the first
         # and last elements (for even transforms) which are unique.
         idx = [slice(None)] * (gout.ndim - 2) \
-                + [slice(1, (s[-1] // 2) + (s[-1] % 2))] + [slice(None)]
+                + [slice(1, (s[-1] // 2) + is_odd)] + [slice(None)]
         gout = T.set_subtensor(gout[idx], gout[idx]*0.5) 
-        return [cuirfft_op(gout, s), DisconnectedType()()]
-    
-    def connection_pattern(self, node):
-        return [[True],[False]]
+        return [cuirfft_op(gout, is_odd)]
 
 curfft_op = CuRFFTOp()
 
@@ -148,7 +140,7 @@ class CuIRFFTOp(Op):
                             broadcastable=[False] * (inp.type.ndim - 1),
                             context_name=inp.type.context_name)
 
-    def make_node(self, inp, s):
+    def make_node(self, inp, is_odd):
         if not scikits_cuda_available:
             raise RuntimeError("scikits.cuda is needed for CuIFFTOp")
 
@@ -161,12 +153,12 @@ class CuIRFFTOp(Op):
         inp = basic_ops.gpu_contiguous(
             basic_ops.as_gpuarray_variable(inp,
                                            basic_ops.infer_context_name(inp)))
-        s = T.as_tensor_variable(s)
+        is_odd = T.as_tensor_variable(is_odd)
 
         assert inp.dtype == "float32"
-        assert s.ndim == 1
+        assert 'int' in is_odd.dtype
 
-        return theano.Apply(self, [inp, s], [self.output_type(inp)()])
+        return theano.Apply(self, [inp, is_odd], [self.output_type(inp)()])
 
     def make_thunk(self, node, storage_map, _, _2):
 
@@ -182,13 +174,16 @@ class CuIRFFTOp(Op):
 
         def thunk():
             input_shape = inputs[0][0].shape
-            s = inputs[1][0]
+            is_odd = inputs[1][0]
+            assert is_odd in (0, 1)
             
             # construct output shape
             # chop off the extra length-2 dimension for real/imag
-            output_shape = [input_shape[0]] + list(s)
+            output_shape = list(input_shape[:-1])
+            # restore full signal length
+            output_shape[-1] = (output_shape[-1] - 1) * 2 + is_odd
             output_shape = tuple(output_shape)
-
+            
             z = outputs[0]
 
             # only allocate if there is no previous allocation of the
@@ -207,7 +202,7 @@ class CuIRFFTOp(Op):
                 # only initialise plan if necessary
                 if plan[0] is None or plan_input_shape[0] != input_shape:
                     plan_input_shape[0] = input_shape
-                    plan[0] = fft.Plan(s,np.complex64, np.float32,
+                    plan[0] = fft.Plan(output_shape[1:], np.complex64, np.float32,
                                        batch=output_shape[0])
                 # Sync GPU variables before computation
                 input_pycuda.sync()
@@ -229,8 +224,8 @@ class CuIRFFTOp(Op):
         
     def grad(self, inputs, output_grads):
         gout, = output_grads
-        s = inputs[1]
-        gf = curfft_op(gout, s)
+        s = gout.shape
+        gf = curfft_op(gout)
         # Multiply the last dimension of the gradient by 2, they represent
         # both positive and negative frequencies, except the first
         # and last elements (for even transforms) which are unique.
@@ -273,16 +268,15 @@ def curfft(inp, norm=None):
     """
 
     s = inp.shape[1:]
-        
     cond_norm = _unitary(norm)
     if cond_norm is None or cond_norm == "no_norm":
         scaling = 1
     elif cond_norm == "ortho":
         scaling = T.sqrt(s.prod().astype('float32'))
     
-    return curfft_op(inp, s) / scaling                                  
+    return curfft_op(inp) / scaling                                  
 
-def cuirfft(inp, norm=None, is_odd=False):
+def cuirfft(inp, norm=None, is_odd=0):
     """
     Performs the real-valued output inverse Fourier Transform using the
     gpuarray backend.
@@ -310,12 +304,12 @@ def cuirfft(inp, norm=None, is_odd=False):
         
     """
 
+    if is_odd != 0:
+        is_odd = 1
+    
     s = inp.shape[1:-1]
-    if is_odd:
-        s = T.set_subtensor(s[-1], (s[-1] - 1) * 2 + 1)
-    else:
-        s = T.set_subtensor(s[-1], (s[-1] - 1) * 2)
-            
+    s = T.set_subtensor(s[-1], (s[-1] - 1) * 2 + is_odd)
+
     cond_norm = _unitary(norm)
     if cond_norm is None:
         scaling = s.prod().astype('float32')
@@ -324,7 +318,7 @@ def cuirfft(inp, norm=None, is_odd=False):
     if cond_norm == "no_norm":
         scaling = 1
 
-    return cuirfft_op(inp, s) / scaling
+    return cuirfft_op(inp, is_odd) / scaling
 
 def _unitary(norm):
     if norm not in (None, "ortho", "no_norm"):

theano/gpuarray/tests/test_fft.py

@@ -30,6 +30,42 @@ N = 16
 
 class TestFFT(unittest.TestCase):
 
+    def test_1Dfft(self):
+        inputs_val = np.random.random((1, N)).astype('float32')
+        inputs = theano.shared(inputs_val)
+
+        x = T.matrix('x', dtype='float32')
+        rfft = theano.gpuarray.fft.curfft(x)
+        f_rfft = theano.function([x], rfft, mode=mode_with_gpu)
+        res_rfft = f_rfft(inputs_val)
+        res_rfft_comp = (np.asarray(res_rfft[:, :, 0]) +
+                         1j * np.asarray(res_rfft[:, :, 1]))
+    
+        rfft_ref = numpy.fft.rfft(inputs_val, axis=1)
+    
+        utt.assert_allclose(rfft_ref, res_rfft_comp)        
+        
+        m = rfft.type()
+        irfft = theano.gpuarray.fft.cuirfft(m)
+        f_irfft = theano.function([m], irfft, mode=mode_with_gpu)
+        res_irfft = f_irfft(res_rfft)
+        
+        utt.assert_allclose(inputs_val, np.asarray(res_irfft))
+        
+        # The numerical gradient of the FFT is sensitive, must set large
+        # enough epsilon to get good accuracy.
+        eps = 1e-1
+        
+        def f_rfft(inp):
+            return theano.gpuarray.fft.curfft(inp)
+        inputs_val = np.random.random((1, N)).astype('float32')
+        utt.verify_grad(f_rfft, [inputs_val], eps=eps)
+        
+        def f_irfft(inp):
+            return theano.gpuarray.fft.cuirfft(inp)
+        inputs_val = np.random.random((1, N//2+1, 2)).astype('float32')
+        utt.verify_grad(f_irfft, [inputs_val], eps=eps)
+    
     def test_rfft(self):
         inputs_val = np.random.random((1, N, N)).astype('float32')
         inputs = theano.shared(inputs_val)
@@ -116,7 +152,7 @@ class TestFFT(unittest.TestCase):
     
         utt.assert_allclose(irfft_ref_ortho * np.sqrt(N*N),
             res_irfft, atol=1e-4, rtol=1e-4)
-
+    
     def test_grad(self):
         # The numerical gradient of the FFT is sensitive, must set large
         # enough epsilon to get good accuracy.
@@ -131,7 +167,7 @@ class TestFFT(unittest.TestCase):
             return theano.gpuarray.fft.cuirfft(inp)
         inputs_val = np.random.random((1, N, N // 2 + 1, 2)).astype('float32')
         utt.verify_grad(f_irfft, [inputs_val], eps=eps)
-
+    
         def f_rfft(inp):
             return theano.gpuarray.fft.curfft(inp, norm='ortho')
         inputs_val = np.random.random((1, N, N)).astype('float32')