C code for Alloc in all cases.

theano/tensor/basic.py

@@ -2849,40 +2849,47 @@ class Alloc(gof.Op):
             out[0][...] = v  # broadcast v to fill us up
 
     def c_code(self, node, name, inp, out, sub):
-        # TODO: use the elemwise code generator here
-        if python_all(node.inputs[0].broadcastable):
-            # filling with a scalar is a common use of alloc
-            # that we can implement relatively easily
-            vv = inp[0]
-            zz, = out
-            fail = sub['fail']
-            if node.outputs[0].ndim == 1:
-                N0 = inp[1]
-                return """
-                npy_intp N0 = ((dtype_%(N0)s*)%(N0)s->data)[0];
-                dtype_%(vv)s vv;
-                dtype_%(zz)s* zz;
-                if ((NULL == %(zz)s) || (%(zz)s->dimensions[0] != N0))
-                {
-                    if (%(zz)s) Py_XDECREF(%(zz)s);
-                    %(zz)s = (PyArrayObject*)PyArray_SimpleNew(1,
-                        &N0, type_num_%(vv)s);
-                    if(!%(zz)s) {
-                        PyErr_SetString(PyExc_MemoryError, "alloc failed");
-                        %(fail)s
-                    }
-                }
-                vv = ((dtype_%(vv)s*)%(vv)s->data)[0];
-                zz = ((dtype_%(zz)s*)%(zz)s->data);
-                assert (%(zz)s->strides[0] == sizeof(dtype_%(zz)s));
-                for (int i = 0; i < N0; ++i)
+        vv = inp[0]
+        ndim = len(inp[1:])
+        zz, = out
+        fail = sub['fail']
+
+        code = """
+            npy_intp shape[%(ndim)s];
+            """ % dict(ndim=ndim)
+
+        # Initialize shape
+        for i, shp_i in enumerate(inp[1:]):
+            code += """
+                shape[%(i)s] = ((dtype_%(shp_i)s*) %(shp_i)s->data)[0];
+                """ % dict(i=i, shp_i=shp_i)
+
+        code += """
+            int need_new_out = (NULL == %(zz)s);
+            for (int i = 0; i < %(ndim)s; i++)
+                need_new_out = (need_new_out
+                                || (%(zz)s->dimensions[i] != shape[i]));
+
+            if (need_new_out)
+            {
+                Py_XDECREF(%(zz)s);
+                %(zz)s = (PyArrayObject*) PyArray_SimpleNew(%(ndim)s,
+                    shape, type_num_%(vv)s);
+                if (!%(zz)s)
                 {
-                    zz[i] = vv;
+                    PyErr_SetString(PyExc_MemoryError, "alloc failed");
+                    %(fail)s
                 }
-                """ % locals()
+            }
+
+            // This function takes care of broadcasting
+            PyArray_CopyInto(%(zz)s, %(vv)s);
+            """ % dict(vv=vv, ndim=ndim, zz=zz, fail=fail)
 
-        # else pretend this never happened
-        return super(Alloc, self).c_code(node, name, inp, out, sub)
+        return code
+
+    def c_code_cache_version(self):
+        return (1,)
 
     def infer_shape(self, node, input_shapes):
         return [node.inputs[1:]]