Add cholesky and gradients

src/ops/cholesky.ts

@@ -0,0 +1,131 @@
+import { Tensor } from '../tensor';
+import { op } from './operation';
+import * as ops from './ops';
+import { Tensor1D } from '../tensor';
+
+// function choleskey(a: tf.Tensor) : tf.Tensor {
+//   // Based on https://rosettacode.org/wiki/Cholesky_decomposition
+//   const n = a.shape[0]
+//   const L = tf.buffer([n, n], a.dtype);
+
+//   let Ldata = L.values as TypedArray;
+//   let aData = a.dataSync();
+
+//   for (let i = 0; i < n; i++) {
+//     for (let k = 0; k < (i + 1); k++) {
+//       let sum = 0.0;
+//       for (let j = 0; j < k; j++) {
+//         sum = sum + Ldata[i * n + j] * Ldata[k * n + j];
+//       }
+//       Ldata[i * n + k] = (i === k) ? Math.sqrt(aData[i * n  + i] - sum)
+//         : (1.0 / Ldata[k * n + k] * (aData[i * n + k] - sum))
+//     }
+//   }
+
+//   return L.toTensor();
+// }
+
+function level2partition(A: Tensor, j: number): [Tensor, Tensor, Tensor, Tensor] {
+  // +-----+
+  // | r d |
+  // | B c |
+  // +-----+
+  const n = A.shape[0];
+  const rr = A.slice([j, 0], [1, j]).as1D();
+  const dd = A.slice([j, j], [1, 1]).asScalar();
+  const B = A.slice([j + 1, 0], [n - (j + 1), j]);
+  const cc = A.slice([j + 1, j], [n - (j + 1), 1]).as1D();
+  return [rr, dd, B, cc];
+}
+
+function cholesky_unblocked_(A: Tensor): Tensor {
+  let n = A.shape[0]
+
+  const Adata = A.dataSync()
+  const res = ops.zerosLike(A);
+  const resData = res.buffer();
+  for (let i = 0; i < n; i++) {
+    for (let j = 0; j < n; j++) {
+      resData.values[i * n + j] = Adata[i * n + j];
+    }
+  }
+
+
+  for (let j = 0; j < n; j++) {
+    let [rr, dd, B, cc] = level2partition(res, j);
+    const ddnew = dd.sub(rr.dot(rr)).sqrt();
+    const ccnew = cc.sub(B.dot(rr)).div(ddnew);
+
+    const ddnewVals = ddnew.dataSync();
+    const ccnewVals = ccnew.dataSync();
+    // update ddnew
+    resData.values[j * n + j] = ddnewVals[0];
+    // update ccnew
+    for (let k = (j + 1); k < n; k++) {
+      resData.values[k * n + j] = ccnewVals[k - (j + 1)];
+    }
+  }
+
+  for (let i = 0; i < n; i++)
+    for (let j = i + 1; j < n; j++)
+      resData.values[i * n + j] = 0;
+  return resData.toTensor();
+}
+
+
+function cholesky_unblocked_grad_(L: Tensor, Abar: Tensor) {
+  let n = L.shape[0];
+
+  const Adata = Abar.dataSync()
+  const res = ops.zerosLike(Abar);
+  const resData = res.buffer();
+  for (let i = 0; i < n; i++) {
+    for (let j = 0; j < n; j++) {
+      resData.values[i * n + j] = Adata[i * n + j];
+    }
+  }
+
+  for (let j = n - 1; j > -1; j--) {
+
+    let [rr, dd, B, cc] = level2partition(L, j);
+    let [rbar, dbar, Bbar, cbar] = level2partition(res, j);
+
+    dbar = dbar.sub(cc.dot(cbar).div(dd));
+    dbar = dbar.div(dd)
+    cbar = cbar.div(dd)
+
+    rbar = rbar.sub(dbar.mul(rr));
+    rbar = rbar.sub(B.transpose().dot(cbar));
+    Bbar = Bbar.sub(ops.outerProduct(cbar as Tensor1D, rr as Tensor1D));
+    dbar = dbar.div(2);
+
+    // Copy into result
+    // update dbar
+    resData.values[j * n + j] = dbar.dataSync()[0];
+    // update cbar
+    let ccnewVals = cbar.dataSync();
+    for (let k = (j + 1); k < n; k++) {
+      resData.values[k * n + j] = ccnewVals[k - (j + 1)];
+    }
+    // update rbar
+    let rbarVals = rbar.dataSync();
+    for (let k = 0; k < j; k++) {
+      resData.values[j * n + k] = rbarVals[k];
+    }
+    // // update Bbar
+    let BbarVals = Bbar.dataSync();
+    for (let r = j + 1; r < n; r++) {
+      for (let c = 0; c < j; c++) {
+        resData.values[r * n + c] = BbarVals[(r - (j + 1)) * n + c];
+      }
+    }
+  }
+
+  for (let i = 0; i < n; i++)
+    for (let j = i + 1; j < n; j++)
+      resData.values[i * n + j] = 0;
+  return resData.toTensor();
+}
+
+export const cholesky = op({ cholesky_unblocked_ })
+export const cholesky_grad = op({ cholesky_unblocked_grad_ })

src/ops/cholesky_test.ts

@@ -0,0 +1,36 @@
+import * as tf from '../index';
+import { describeWithFlags } from '../jasmine_util';
+import { cholesky, cholesky_grad } from './cholesky';
+import { expectArraysClose, ALL_ENVS } from '../test_util';
+
+describeWithFlags('cholesky-small', ALL_ENVS, () => {
+  it('Compute cholesky', () => {
+    const a = tf.tensor2d([25, 15, -5, 15, 18, 0, -5, 0, 11],
+      [3, 3], "float32");
+
+    let L = cholesky(a);
+    let res = tf.tensor2d([5, 0, 0, 3, 3, 0, -1, 1, 3], [3, 3], "float32");
+    expectArraysClose(L, res);
+    let rec = L.matMul(L.transpose());
+
+    expectArraysClose(a, rec);
+  })
+
+  it('Compute gradients', () => {
+    const a = tf.tensor2d([25, 15, -5, 15, 18, 0, -5, 0, 11],
+      [3, 3], "float32");
+
+    let L = cholesky(a);
+    let dL = cholesky_grad(L, tf.ones([3, 3]));
+
+    let expected = tf.tensor2d(
+      [0.0867, 0.0000, 0.0000, 0.0889, 0.1296, 0.0000, 0.1333, 0.2222, 0.1667],
+      [3, 3], 'float32')
+
+    expectArraysClose(
+      dL, expected
+    );
+  });
+
+
+});