Product with axis = 0

TODO.md

@@ -1,6 +1,6 @@
 4. in the refactor, add consts
 10. For performance reasons, is it useful to have a dense matmul with A and B as dense matrices?
-11. right matmul, add broadcasting logic as in left matmul
+11. right matmul, add broadcasting logic as in left matmul. Is this necessary?
 
 Going through all atoms to see that sparsity pattern is computed in initialization of jacobian:
 2. trace - not ok

include/other.h

@@ -10,4 +10,7 @@ expr *new_quad_form(expr *child, CSR_Matrix *Q);
 /* product of all entries, without axis argument */
 expr *new_prod(expr *child);
 
+/* product of entries along axis=0 (columnwise products) */
+expr *new_prod_axis_zero(expr *child);
+
 #endif /* OTHER_H */

include/subexpr.h

@@ -56,6 +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
+{
+    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;
+
 /* Horizontal stack (concatenate) */
 typedef struct hstack_expr
 {

python/atoms/prod_axis_zero.h

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: Apache-2.0
+#ifndef ATOM_PROD_AXIS_ZERO_H
+#define ATOM_PROD_AXIS_ZERO_H
+
+#include "common.h"
+#include "other.h"
+
+static PyObject *py_make_prod_axis_zero(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_zero(child);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create prod_axis_zero node");
+        return NULL;
+    }
+    expr_retain(node); /* Capsule owns a reference */
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_PROD_AXIS_ZERO_H */

python/bindings.c

@@ -23,6 +23,7 @@
 #include "atoms/neg.h"
 #include "atoms/power.h"
 #include "atoms/prod.h"
+#include "atoms/prod_axis_zero.h"
 #include "atoms/promote.h"
 #include "atoms/quad_form.h"
 #include "atoms/quad_over_lin.h"
@@ -77,6 +78,8 @@ static PyMethodDef DNLPMethods[] = {
      "Create constant vector multiplication node (a ∘ f(x))"},
     {"make_power", py_make_power, METH_VARARGS, "Create power node"},
     {"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_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"},