jacobian of hstack

Author: dance858

include/affine.h

@@ -8,6 +8,7 @@ expr *new_linear(expr *u, const CSR_Matrix *A);
 
 expr *new_add(expr *left, expr *right);
 expr *new_sum(expr *child, int axis);
+expr *new_hstack(expr **args, int n_args, int n_vars);
 
 expr *new_constant(int d1, int d2, int n_vars, const double *values);
 expr *new_variable(int d1, int d2, int var_id, int n_vars);

include/expr.h

@@ -30,6 +30,8 @@ typedef struct expr
     int refcount;
     struct expr *left;
     struct expr *right;
+    struct expr **args; /* hstack can have multiple arguments */
+    int n_args;
     double *dwork;
     int *iwork;
     struct int_double_pair *int_double_pairs; /* for sorting jacobian entries */

src/affine/add.c

@@ -1,6 +1,6 @@
 #include "affine.h"
 
-static void add_forward(expr *node, const double *u)
+static void forward(expr *node, const double *u)
 {
     /* children's forward passes */
     node->left->forward(node->left, u);
@@ -52,7 +52,7 @@ expr *new_add(expr *left, expr *right)
     node->right = right;
     expr_retain(left);
     expr_retain(right);
-    node->forward = add_forward;
+    node->forward = forward;
     node->is_affine = is_affine;
     node->jacobian_init = jacobian_init;
     node->eval_jacobian = eval_jacobian;

src/affine/hstack.c

@@ -0,0 +1,103 @@
+#include "affine.h"
+#include <assert.h>
+#include <string.h>
+
+static void forward(expr *node, const double *u)
+{
+
+    /* children's forward passes */
+    for (int i = 0; i < node->n_args; i++)
+    {
+        node->args[i]->forward(node->args[i], u);
+    }
+
+    /* concatenate values horizontally */
+    int offset = 0;
+    for (int i = 0; i < node->n_args; i++)
+    {
+        expr *child = node->args[i];
+        memcpy(node->value + offset, child->value, child->size * sizeof(double));
+        offset += child->size;
+    }
+}
+
+static void jacobian_init(expr *node)
+{
+    /* initialize children's jacobians */
+    int nnz = 0;
+    for (int i = 0; i < node->n_args; i++)
+    {
+        node->args[i]->jacobian_init(node->args[i]);
+        nnz += node->args[i]->jacobian->nnz;
+    }
+
+    node->jacobian = new_csr_matrix(node->size, node->n_vars, nnz);
+}
+
+static void eval_jacobian(expr *node)
+{
+    /* evaluate children's jacobians */
+    int row_offset = 0;
+    CSR_Matrix *A = node->jacobian;
+    A->nnz = 0;
+
+    for (int i = 0; i < node->n_args; i++)
+    {
+        expr *child = node->args[i];
+        child->eval_jacobian(child);
+        CSR_Matrix *B = child->jacobian;
+
+        /* copy columns and values */
+        memcpy(A->x + A->nnz, B->x, B->nnz * sizeof(double));
+        memcpy(A->i + A->nnz, B->i, B->nnz * sizeof(int));
+
+        /* set row pointers */
+        for (int r = 0; r < child->size; r++)
+        {
+            A->p[row_offset + r] = A->nnz + B->p[r];
+        }
+
+        A->nnz += B->nnz;
+        row_offset += child->size;
+    }
+    A->p[node->size] = A->nnz;
+}
+
+static bool is_affine(expr *node)
+{
+    for (int i = 0; i < node->n_args; i++)
+    {
+        if (!node->args[i]->is_affine(node->args[i]))
+        {
+            return false;
+        }
+    }
+    return true;
+}
+
+expr *new_hstack(expr **args, int n_args, int n_vars)
+{
+    /* compute second dimension */
+    int d2 = 0;
+    for (int i = 0; i < n_args; i++)
+    {
+        d2 += args[i]->d2;
+    }
+
+    expr *node = new_expr(args[0]->d1, d2, n_vars);
+    if (!node) return NULL;
+    node->args = args;
+    node->n_args = n_args;
+
+    for (int i = 0; i < n_args; i++)
+    {
+        expr_retain(args[i]);
+    }
+
+    node->forward = forward;
+    node->is_affine = is_affine;
+    node->jacobian_init = jacobian_init;
+    node->eval_jacobian = eval_jacobian;
+
+    return node;
+}

src/elementwise_univariate/entr.c

@@ -18,7 +18,7 @@ static void forward(expr *node, const double *u)
 static void eval_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         vals[j] = -log(child->value[j]) - 1.0;
     }

src/elementwise_univariate/exp.c

@@ -16,7 +16,7 @@ static void forward(expr *node, const double *u)
 
 static void eval_local_jacobian(expr *node, double *vals)
 {
-    memcpy(vals, node->value, node->d1 * sizeof(double));
+    memcpy(vals, node->value, node->size * sizeof(double));
 }
 
 expr *new_exp(expr *child)

src/elementwise_univariate/hyperbolic.c

@@ -14,7 +14,7 @@ static void sinh_forward(expr *node, const double *u)
 static void sinh_eval_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         vals[j] = cosh(child->value[j]);
     }
@@ -46,7 +46,7 @@ static void tanh_forward(expr *node, const double *u)
 static void tanh_eval_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         double c = cosh(child->value[j]);
         vals[j] = 1.0 / (c * c);
@@ -79,7 +79,7 @@ static void asinh_forward(expr *node, const double *u)
 static void asinh_eval_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         vals[j] = 1.0 / sqrt(1.0 + child->value[j] * child->value[j]);
     }
@@ -111,7 +111,7 @@ static void atanh_forward(expr *node, const double *u)
 static void atanh_eval_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         vals[j] = 1.0 / (1.0 - child->value[j] * child->value[j]);
     }

src/elementwise_univariate/logistic.c

@@ -26,7 +26,7 @@ static void forward(expr *node, const double *u)
 static void eval_local_jacobian(expr *node, double *vals)
 {
     double *x = node->left->value;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         if (x[j] >= 0)
         {

src/elementwise_univariate/power.c

@@ -19,7 +19,7 @@ static void forward(expr *node, const double *u)
 static void eval_local_jacobian(expr *node, double *vals)
 {
     double *x = node->left->value;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         vals[j] = node->p * pow(x[j], node->p - 1);
     }

src/elementwise_univariate/trig.c

@@ -14,7 +14,7 @@ static void sin_forward(expr *node, const double *u)
 static void sin_eval_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         vals[j] = cos(child->value[j]);
     }
@@ -46,7 +46,7 @@ static void cos_forward(expr *node, const double *u)
 static void cos_eval_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         vals[j] = -sin(child->value[j]);
     }
@@ -78,7 +78,7 @@ static void tan_forward(expr *node, const double *u)
 static void tan_eval_local_jacobian(expr *node, double *vals)
 {
     expr *child = node->left;
-    for (int j = 0; j < node->d1; j++)
+    for (int j = 0; j < node->size; j++)
     {
         double c = cos(child->value[j]);
         vals[j] = 1.0 / (c * c);