Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings
This repository was archived by the owner on Aug 15, 2019. It is now read-only.

Add cholesky and gradients #1492

Open
ethanluoyc wants to merge 2 commits into tensorflow:master
base: master
Choose a base branch
Loading
from ethanluoyc:cholesky
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 131 additions & 0 deletions src/ops/cholesky.ts
View file Open in desktop
Original file line number Diff line number Diff line change
@@ -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_ })
36 changes: 36 additions & 0 deletions src/ops/cholesky_test.ts
View file Open in desktop
Original file line number Diff line number Diff line change
@@ -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
);
});


});

AltStyle によって変換されたページ (->オリジナル) /