From d59768cb1026e5a4636ecc4647593d1e2fb42ffe Mon Sep 17 00:00:00 2001 From: ethanluoyc Date: Thu, 10 Jan 2019 10:55:32 +0000 Subject: [PATCH 1/2] Add cholesky and gradients --- src/ops/cholesky.ts | 111 +++++++++++++++++++++++++++++++++++++++ src/ops/cholesky_test.ts | 31 +++++++++++ 2 files changed, 142 insertions(+) create mode 100644 src/ops/cholesky.ts create mode 100644 src/ops/cholesky_test.ts diff --git a/src/ops/cholesky.ts b/src/ops/cholesky.ts new file mode 100644 index 0000000000..ffa5bb9047 --- /dev/null +++ b/src/ops/cholesky.ts @@ -0,0 +1,111 @@ +import { Tensor } from '../tensor'; +import { op } from './operation'; +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 res = A.clone(); + const resData = res.buffer(); + + 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 res = Abar.clone(); + let resData = res.buffer(); + let n = L.shape[0]; + + 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( + tf.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]; + } + } + } + return resData.toTensor(); +} + +export const cholesky = op({ cholesky_unblocked_ }) diff --git a/src/ops/cholesky_test.ts b/src/ops/cholesky_test.ts new file mode 100644 index 0000000000..f8b8a1cf93 --- /dev/null +++ b/src/ops/cholesky_test.ts @@ -0,0 +1,31 @@ +import * as tf from '../index'; +import { describeWithFlags } from '../jasmine_util'; +import { cholesky } from './cholesky'; +import { expectArraysClose, CPU_ENVS } from '../test_util'; + +// import { scalar, tensor1d, tensor2d, tensor3d, tensor4d } from './ops'; + +describeWithFlags('cholesky-tiny', CPU_ENVS, () => { + it('Compute cholesky decomposition', () => { + const a = tf.tensor2d([25, 15, -5, 15, 18, 0, -5, 0, 11], + [3, 3], "float32"); + + a.print(); + let L = cholesky(a.clone()); + a.print(); + 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 correctly', () => { + // console.log('----grad----') + // let dL = cholesky_unblocked_grad(L, tf.eye(3, 3)); + }); + + +}); From 7d06e4cc91c161c17f835c6286ba77cffdd9c9ca Mon Sep 17 00:00:00 2001 From: ethanluoyc Date: Thu, 10 Jan 2019 12:13:44 +0000 Subject: [PATCH 2/2] Add more tests and grad op --- src/ops/cholesky.ts | 32 ++++++++++++++++++++++++++------ src/ops/cholesky_test.ts | 35 ++++++++++++++++++++--------------- 2 files changed, 46 insertions(+), 21 deletions(-) diff --git a/src/ops/cholesky.ts b/src/ops/cholesky.ts index ffa5bb9047..b515420a42 100644 --- a/src/ops/cholesky.ts +++ b/src/ops/cholesky.ts @@ -1,5 +1,6 @@ import { Tensor } from '../tensor'; import { op } from './operation'; +import * as ops from './ops'; import { Tensor1D } from '../tensor'; // function choleskey(a: tf.Tensor) : tf.Tensor { @@ -39,8 +40,16 @@ function level2partition(A: Tensor, j: number): [Tensor, Tensor, Tensor, Tensor] function cholesky_unblocked_(A: Tensor): Tensor { let n = A.shape[0] - const res = A.clone(); + + 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); @@ -64,11 +73,18 @@ function cholesky_unblocked_(A: Tensor): Tensor { } -function cholesky_unblocked_grad(L: Tensor, Abar: Tensor) { - let res = Abar.clone(); - let resData = res.buffer(); +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); @@ -80,8 +96,7 @@ function cholesky_unblocked_grad(L: Tensor, Abar: Tensor) { rbar = rbar.sub(dbar.mul(rr)); rbar = rbar.sub(B.transpose().dot(cbar)); - Bbar = Bbar.sub( - tf.outerProduct(cbar as Tensor1D, rr as Tensor1D)); + Bbar = Bbar.sub(ops.outerProduct(cbar as Tensor1D, rr as Tensor1D)); dbar = dbar.div(2); // Copy into result @@ -105,7 +120,12 @@ function cholesky_unblocked_grad(L: Tensor, Abar: Tensor) { } } } + + 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_ }) diff --git a/src/ops/cholesky_test.ts b/src/ops/cholesky_test.ts index f8b8a1cf93..a22dcd263a 100644 --- a/src/ops/cholesky_test.ts +++ b/src/ops/cholesky_test.ts @@ -1,30 +1,35 @@ import * as tf from '../index'; import { describeWithFlags } from '../jasmine_util'; -import { cholesky } from './cholesky'; -import { expectArraysClose, CPU_ENVS } from '../test_util'; +import { cholesky, cholesky_grad } from './cholesky'; +import { expectArraysClose, ALL_ENVS } from '../test_util'; -// import { scalar, tensor1d, tensor2d, tensor3d, tensor4d } from './ops'; - -describeWithFlags('cholesky-tiny', CPU_ENVS, () => { - it('Compute cholesky decomposition', () => { +describeWithFlags('cholesky-small', ALL_ENVS, () => { + it('Compute cholesky', () => { const a = tf.tensor2d([25, 15, -5, 15, 18, 0, -5, 0, 11], [3, 3], "float32"); - a.print(); - let L = cholesky(a.clone()); - a.print(); + 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 - ); + expectArraysClose(a, rec); }) - it('Compute gradients correctly', () => { - // console.log('----grad----') - // let dL = cholesky_unblocked_grad(L, tf.eye(3, 3)); + 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 + ); });