Merge pull request #18 from null-a/cholesky

Add backwards pass for cholesky.

ad/README.md

@@ -62,6 +62,7 @@ adnn comes with a large number of built-in AD primitives:
   - **`ad.tensor.inverse(x)`**: Returns the inverse of the matrix `x`.
   - **`ad.tensor.determinant(x)`**: Returns the determinant of the matrix `x`.
   - **`ad.tensor.dot(x, y)`**: Returns the inner product of the matrices `x` and `y`.
+  - **`ad.tensor.cholesky(x)`**: Returns the [Cholesky decomposition](https://en.wikipedia.org/wiki/Cholesky_decomposition) of the matrix `x`.
 - **Miscellaneous**
   - **`ad.tensor.softmax(x)`**: Compute the [Softmax](https://en.wikipedia.org/wiki/Softmax_function) function for a tensor `x`.
 

ad/functions.js

@@ -274,6 +274,36 @@ fns.tensor.dot = func.newBinaryFunction({
   }
 });
 
+// Smith, Stephen P. "Differentiation of the Cholesky algorithm."
+// Journal of Computational and Graphical Statistics 4.2 (1995):
+// 134-147.
+fns.tensor.cholesky = func.newUnaryFunction({
+	OutputType: Tensor,
+	name: 'cholesky',
+	forward: function(x) {
+		return x.cholesky();
+	},
+	backward: function(x) {
+		var i, j, k;
+		var L = this;
+		var dL = this.dx.tril();
+		var n = L.x.dims[0];
+		for (k = n - 1; k >= 0; k--) {
+			for (j = k + 1; j < n; j++) {
+				for (i = j; i < n; i++) {
+					dL.data[i * n + k] -= dL.data[i * n + j] * L.x.data[j * n + k];
+					dL.data[j * n + k] -= dL.data[i * n + j] * L.x.data[i * n + k];
+				}
+			}
+			for (j = k + 1; j < n; j++) {
+				dL.data[j * n + k] /= L.x.data[k * n + k];
+				dL.data[k * n + k] -= dL.data[j * n + k] * L.x.data[j * n + k];
+			}
+			dL.data[k * n + k] /= 2 * L.x.data[k * n + k];
+		}
+		x.dx.addeq(dL);
+	}
+});
 
 // Tensor reductions  -----------------------------------------------------
 

tensor.js

@@ -510,6 +510,22 @@ Tensor.prototype.cholesky = function() {
   return L;
 };
 
+// Return a copy of a matrix with the elements above the main diagonal
+// zeroed.
+Tensor.prototype.tril = function() {
+	var a = this;
+	assert.ok((a.rank === 2) && (a.dims[0] === a.dims[1]),
+		  'tril is only defined for square matrices.');
+
+	var n = a.dims[0];
+	var ret = new Tensor([n, n]);
+	for (var i = 0; i < n; i++) {
+		for (var j = 0; j <= i; j++) {
+			ret.data[i * n + j] = a.data[i * n + j];
+		}
+	}
+	return ret;
+};
 
 module.exports = Tensor;