in the process of making tensor ops work - problems with refcounts

_test_tensor_ops.py

+
+
+import unittest
+
+from gof import ResultBase, Op, Env, modes
+import gof
+
+from tensor import *
+from tensor_ops import *
+
+import numpy
+
+import sys
+
+
+def inputs():
+    x = modes.build(tensor([[1.0, 2.0], [3.0, 4.0]], 'x'))
+    y = None
+    z = None
+    return x, y, z
+
+def env(inputs, outputs, validate = True, features = []):
+    return Env(inputs, outputs, features = features, consistency_check = validate)
+
+
+class _test_TensorOps(unittest.TestCase):
+
+    def test_0(self):
+        x, y, z = inputs()
+        e = transpose(x)
+        g = env([x], [e])
+        fn, (i, ), (o, ) = gof.cc.CLinker(g).make_thunk()
+        i.data = [[1.0, 2.0], [3.0, 4.0]]
+#        print sys.getrefcount(i.data)
+        fn()
+#        print sys.getrefcount(i.data)
+#        print sys.getrefcount(o.data)
+        print o.data
+#        assert res == numpy.asarray(arr)
+
+#     def test_1(self):
+#         x, y, z = inputs()
+#         e = mul(add(x, y), div(x, y))
+#         g = env([x, y], [e])
+#         fn = gof.cc.CLinker(g).make_function()
+#         assert fn(1.0, 2.0) == 1.5
+#         assert e.data == 1.5
+
+
+if __name__ == '__main__':
+    unittest.main()
+
+
+
+
+
+

tensor.py

 
 import numpy
 from copy import copy
+import inspect
+from gof import ResultBase, Op, utils
 
-from gof import ResultBase
-from gof import Op
 
 
 def tensor(data, name = None):
+    data = numpy.asarray(data)
     return Tensor(data.dtype, [0]*len(data.shape), data, name)
 
 def _broadcastable_pattern(pattern):
     def factory(data = None, name = None):
         if data: assert len(data.shape) == len(pattern)
         return Tensor(data.dtype, pattern, data, name)
+    return factory
 
 matrix = _broadcastable_pattern([0, 0])
 row = _broadcastable_pattern([1, 0])
@@ -23,7 +25,7 @@ class Tensor(ResultBase):
 
     def __init__(self, dtype, broadcastable, data=None, name=None):
         self.broadcastable = broadcastable
-        self.dtype = dtype
+        self.dtype = str(dtype)
         ResultBase.__init__(self, role = None, data = None, name = name)
 
     def filter(self, data):
@@ -32,31 +34,54 @@ class Tensor(ResultBase):
             assert not b or s == 1
         return arr
 
+    def dtype_specs(self):
+        return {'float64': (float, 'double')}[self.dtype]
+            
     def c_declare(self):
         return """
         PyArrayObject* %%(name)s;
         typedef %(dtype)s %%(name)s_dtype;
-        """ % dict(dtype = self.to_c_type(self.dtype))
+        """ % dict(dtype = self.dtype_specs()[1])
+
+    def c_init(self):
+        return """
+        %(name)s = NULL;
+        """
 
-    def c_data_extract(self):
+    def c_extract(self):
         return """
-        if (py_%(name)s == Py_None)
+        if (py_%(name)s == Py_None) {
             %(name)s = NULL;
-        else
+        }
+        else if (!PyArray_Check(py_%(name)s)) {
+            PyErr_SetString(PyExc_ValueError, "expected an ndarray");
+            %(fail)s
+        }
+        else {
             %(name)s = (PyArrayObject*)(py_%(name)s);
+            Py_XINCREF(%(name)s);
+        }
         """
 
-    def c_data_cleanup(self):
-        return ""
+    def c_cleanup(self):
+        return """
+        if (%(name)s) {
+            Py_XDECREF(%(name)s);
+            for (int i = 0; i < PyArray_REFCOUNT(%(name)s); i++) {
+                printf("X");
+            }
+            printf("Y\\n");
+        }
+        """
     
-    def c_data_sync(self):
+    def c_sync(self):
         return """
         if (!%(name)s) {
-            Py_XDECREF(py_%(name));
+            Py_XDECREF(py_%(name)s);
             py_%(name)s = Py_None;
         }
         else if ((void*)py_%(name)s != (void*)%(name)s) {
-            Py_XDECREF(py_%(name));
+            Py_XDECREF(py_%(name)s);
             py_%(name)s = (PyObject*)%(name)s;
         }
         """
@@ -76,48 +101,3 @@ class Tensor(ResultBase):
         return cpy
 
 
-
-def TensorOp(Op):
-
-    nin = -1
-    nout = 1
-    
-    def __init__(self, *inputs):
-
-        def wrap_as_tensor(x):
-            if isinstance(x, Tensor):
-                return x
-            else:
-                return Tensor(x)
-
-        inputs = map(wrap_as_tensor, inputs)
-        
-        if self.nin >= 0:
-            if len(inputs) != self.nin:
-                raise TypeError("Wrong number of inputs for %s (got %i, expected %i)") \
-                    % (self, len(inputs), self.nin)
-
-        i_broadcastables = [getattr(input, 'broadcastable', None) for input in inputs]
-        i_dtypes = [getattr(input, 'dtype', None) for input in inputs]
-
-        o_broadcastables = utils.from_return_values(self.propagate_broadcastable(*i_broadcastables))
-        o_dtypes = utils.from_return_values(self.propagate_dtype(*i_dtypes))
-
-        self.inputs = inputs
-        self.outputs = [Tensor(dtype, broadcastable) for broadcastable, dtype in zip(o_broadcastables, o_dtypes)]
-
-    def propagate_broadcastable(self, *inputs):
-        raise AbstractFunctionError()
-
-    def propagate_dtype(self, *inputs):
-        raise AbstractFunctionError()
-    
-    def impl(self, *inputs):
-        raise AbstractFunctionError()
-    
-    def perform(self):
-        self.outputs[0].data = self.impl(*[input.data for input in self.inputs])
-
-
-
-

tensor_ops.py

+
+from tensor import *
+from gof import Op, utils
+
+
+def upcast(dtype, *dtypes):
+    z = numpy.zeros((), dtype = dtype)
+    for dtype in dtypes:
+        z = z + numpy.zeros((), dtype = dtype)
+    return str(z.dtype)
+
+class TensorOp(Op):
+
+    nin = -1
+    nout = 1
+
+    cast_method = lambda self, *args: upcast(*args)
+    
+    def __init__(self, *inputs):
+
+        def wrap_as_tensor(x):
+            if isinstance(x, Tensor):
+                return x
+            else:
+                return Tensor(x)
+
+        inputs = map(wrap_as_tensor, inputs)
+        
+        if self.nin >= 0:
+            if len(inputs) != self.nin:
+                raise TypeError("Wrong number of inputs for %s (got %i, expected %i)") \
+                    % (self, len(inputs), self.nin)
+
+        i_broadcastables = [getattr(input, 'broadcastable', None) for input in inputs]
+        i_dtypes = [getattr(input, 'dtype', None) for input in inputs]
+
+        o_broadcastables = utils.from_return_values(self.propagate_broadcastable(*i_broadcastables))
+        o_dtypes = utils.from_return_values(self.propagate_dtype(*i_dtypes))
+
+        self.inputs = inputs
+        self.outputs = [Tensor(dtype, broadcastable) for broadcastable, dtype in zip(o_broadcastables, o_dtypes)]
+
+    def propagate_broadcastable(self, *inputs):
+        raise AbstractFunctionError()
+    
+    def propagate_dtype(self, *i_dtypes):
+        for dtype in i_dtypes:
+            if dtype is None:
+                raise TypeError("Expected a Tensor.")
+        return self.cast_method(*i_dtypes)
+    
+    def impl(self, *inputs):
+        raise AbstractFunctionError()
+    
+    def perform(self):
+        self.outputs[0].data = self.impl(*[input.data for input in self.inputs])
+    
+    def c_var_names(self):
+        (self, inames, onames), _1, _2, _3 = inspect.getargspec(self.c_impl)
+        inames = utils.from_return_values(inames)
+        onames = utils.from_return_values(onames)
+        return [inames, onames]
+    
+    def c_code(self):
+        return self.c_impl(self.inputs, self.outputs)
+
+    def c_impl(self, inputs, outputs):
+        raise AbstractFunctionError()
+
+        
+
+class UnaryTensorOp(TensorOp):
+    nin = 1
+
+class BinaryTensorOp(TensorOp):
+    nin = 2
+
+
+
+class Transpose(UnaryTensorOp):
+
+    def propagate_broadcastable(self, x):
+        x2 = copy(x)
+        x2.reverse()
+        return x2
+
+    def impl(self, x):
+        return x.T
+
+    def c_impl(self, x, z):
+        return """
+        PyArrayObject* transposed = (PyArrayObject*)PyArray_Transpose(%(x)s, NULL);
+        if (PyArray_REFCOUNT(transposed) == 1) {
+            printf("lala\\n");
+        }
+        if (%(z)s) {
+            Py_XDECREF(%(z)s);
+        }
+        %(z)s = transposed;
+        """
+
+
+from gof import modes
+modes.make_constructors(globals())
+