diff --git a/src/ops/cholesky.ts b/src/ops/cholesky.ts new file mode 100644 index 0000000000..b515420a42 --- /dev/null +++ b/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_ }) diff --git a/src/ops/cholesky_test.ts b/src/ops/cholesky_test.ts new file mode 100644 index 0000000000..a22dcd263a --- /dev/null +++ b/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 + ); + }); + + +});