Prod axis one

include/other.h

@@ -13,4 +13,7 @@ expr *new_prod(expr *child);
 /* product of entries along axis=0 (columnwise products) */
 expr *new_prod_axis_zero(expr *child);
 
+/* product of entries along axis=1 (rowwise products) */
+expr *new_prod_axis_one(expr *child);
+
 #endif /* OTHER_H */

include/subexpr.h

@@ -56,14 +56,15 @@ typedef struct prod_expr
     double prod_nonzero; /* product of non-zero elements */
 } prod_expr;
 
-/* Product of entries along axis=0 (columnwise products) */
-typedef struct prod_axis_zero_expr
+/* Product of entries along axis=0 (columnwise products) or axis = 1 (rowwise
+ * products) */
+typedef struct prod_axis
 {
     expr base;
-    int *num_of_zeros;    /* num of zeros for each column */
-    int *zero_index;      /* stores idx of zero element per column */
-    double *prod_nonzero; /* product of non-zero elements per column */
-} prod_axis_zero_expr;
+    int *num_of_zeros; /* num of zeros for each column / row depending on the axis*/
+    int *zero_index;   /* stores idx of zero element per column / row */
+    double *prod_nonzero; /* product of non-zero elements per column / row */
+} prod_axis;
 
 /* Horizontal stack (concatenate) */
 typedef struct hstack_expr

python/atoms/prod_axis_one.h

@@ -0,0 +1,32 @@
+#ifndef ATOM_PROD_AXIS_ONE_H
+#define ATOM_PROD_AXIS_ONE_H
+
+#include "common.h"
+#include "other.h"
+
+static PyObject *py_make_prod_axis_one(PyObject *self, PyObject *args)
+{
+    (void) self;
+    PyObject *child_capsule;
+    if (!PyArg_ParseTuple(args, "O", &child_capsule))
+    {
+        return NULL;
+    }
+    expr *child = (expr *) PyCapsule_GetPointer(child_capsule, EXPR_CAPSULE_NAME);
+    if (!child)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_prod_axis_one(child);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create prod_axis_one node");
+        return NULL;
+    }
+    expr_retain(node); /* Capsule owns a reference */
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_PROD_AXIS_ONE_H */

python/bindings.c

@@ -23,6 +23,7 @@
 #include "atoms/neg.h"
 #include "atoms/power.h"
 #include "atoms/prod.h"
+#include "atoms/prod_axis_one.h"
 #include "atoms/prod_axis_zero.h"
 #include "atoms/promote.h"
 #include "atoms/quad_form.h"
@@ -80,6 +81,8 @@ static PyMethodDef DNLPMethods[] = {
     {"make_prod", py_make_prod, METH_VARARGS, "Create prod node"},
     {"make_prod_axis_zero", py_make_prod_axis_zero, METH_VARARGS,
      "Create prod_axis_zero node"},
+    {"make_prod_axis_one", py_make_prod_axis_one, METH_VARARGS,
+     "Create prod_axis_one node"},
     {"make_sin", py_make_sin, METH_VARARGS, "Create sin node"},
     {"make_cos", py_make_cos, METH_VARARGS, "Create cos node"},
     {"make_tan", py_make_tan, METH_VARARGS, "Create tan node"},