This repository has been archived by the owner on Aug 15, 2019. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 949
/
linalg_ops.ts
267 lines (251 loc) · 8.93 KB
/
linalg_ops.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
/**
* @license
* Copyright 2018 Google LLC. All Rights Reserved.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =============================================================================
*/
/**
* Linear algebra ops.
*/
import {ENV} from '../environment';
import {dispose} from '../globals';
import {Tensor, Tensor1D, Tensor2D} from '../tensor';
import {assert} from '../util';
import {eye, squeeze, stack, unstack} from './array_ops';
import {split} from './concat_split';
import {norm} from './norm';
import {op} from './operation';
import {sum} from './reduction_ops';
import {tensor2d} from './tensor_ops';
/**
* Gram-Schmidt orthogonalization.
*
* ```js
* const x = tf.tensor2d([[1, 2], [3, 4]]);
* let y = tf.linalg.gramSchmidt(x);
* y.print();
* console.log('Othogonalized:');
* y.dot(y.transpose()).print(); // should be nearly the identity matrix.
* console.log('First row direction maintained:');
* console.log(y.get(0, 1) / y.get(0, 0)); // should be nearly 2.
* ```
*
* @param xs The vectors to be orthogonalized, in one of the two following
* formats:
* - An Array of `tf.Tensor1D`.
* - A `tf.Tensor2D`, i.e., a matrix, in which case the vectors are the rows
* of `xs`.
* In each case, all the vectors must have the same length and the length
* must be greater than or equal to the number of vectors.
* @returns The orthogonalized and normalized vectors or matrix.
* Orthogonalization means that the vectors or the rows of the matrix
* are orthogonal (zero inner products). Normalization means that each
* vector or each row of the matrix has an L2 norm that equals `1`.
*/
/**
* @doc {heading:'Operations',
* subheading:'Linear Algebra',
* namespace:'linalg'}
*/
function gramSchmidt_(xs: Tensor1D[]|Tensor2D): Tensor1D[]|Tensor2D {
let inputIsTensor2D: boolean;
if (Array.isArray(xs)) {
inputIsTensor2D = false;
assert(
xs != null && xs.length > 0,
() => 'Gram-Schmidt process: input must not be null, undefined, or ' +
'empty');
const dim = xs[0].shape[0];
for (let i = 1; i < xs.length; ++i) {
assert(
xs[i].shape[0] === dim,
() =>
'Gram-Schmidt: Non-unique lengths found in the input vectors: ' +
`(${(xs as Tensor1D[])[i].shape[0]} vs. ${dim})`);
}
} else {
inputIsTensor2D = true;
xs = split(xs, xs.shape[0], 0).map(x => squeeze(x, [0]));
}
assert(
xs.length <= xs[0].shape[0],
() => `Gram-Schmidt: Number of vectors (${
(xs as Tensor1D[]).length}) exceeds ` +
`number of dimensions (${(xs as Tensor1D[])[0].shape[0]}).`);
const ys: Tensor1D[] = [];
const xs1d = xs as Tensor1D[];
for (let i = 0; i < xs.length; ++i) {
ys.push(ENV.engine.tidy(() => {
let x = xs1d[i];
if (i > 0) {
for (let j = 0; j < i; ++j) {
const proj = sum(ys[j].mulStrict(x)).mul(ys[j]);
x = x.sub(proj);
}
}
return x.div(norm(x, 'euclidean'));
}));
}
if (inputIsTensor2D) {
return stack(ys, 0) as Tensor2D;
} else {
return ys;
}
}
/**
* Compute QR decomposition of m-by-n matrix using Householder transformation.
*
* Implementation based on
* [http://www.cs.cornell.edu/~bindel/class/cs6210-f09/lec18.pdf]
* (http://www.cs.cornell.edu/~bindel/class/cs6210-f09/lec18.pdf)
*
* ```js
* const a = tf.tensor2d([[1, 2], [3, 4]]);
* let [q, r] = tf.linalg.qr(a);
* console.log('Q');
* q.print();
* console.log('R');
* r.print();
* console.log('Orthogonalized');
* q.dot(q.transpose()).print() // should be nearly the identity matrix.
* console.log('Reconstructed');
* q.dot(r).print(); // should be nearly [[1, 2], [3, 4]];
* ```
*
* @param x The `tf.Tensor` to be QR-decomposed. Must have rank >= 2. Suppose
* it has the shape `[..., M, N]`.
* @param fullMatrices An optional boolean parameter. Defaults to `false`.
* If `true`, compute full-sized `Q`. If `false` (the default),
* compute only the leading N columns of `Q` and `R`.
* @returns An `Array` of two `tf.Tensor`s: `[Q, R]`. `Q` is a unitary matrix,
* i.e., its columns all have unit norm and are mutually orthogonal.
* If `M >= N`,
* If `fullMatrices` is `false` (default),
* - `Q` has a shape of `[..., M, N]`,
* - `R` has a shape of `[..., N, N]`.
* If `fullMatrices` is `true` (default),
* - `Q` has a shape of `[..., M, M]`,
* - `R` has a shape of `[..., M, N]`.
* If `M < N`,
* - `Q` has a shape of `[..., M, M]`,
* - `R` has a shape of `[..., M, N]`.
* @throws If the rank of `x` is less than 2.
*/
/**
* @doc {heading:'Operations',
* subheading:'Linear Algebra',
* namespace:'linalg'}
*/
function qr_(x: Tensor, fullMatrices = false): [Tensor, Tensor] {
if (x.rank < 2) {
throw new Error(
`qr() requires input tensor to have a rank >= 2, but got rank ${
x.rank}`);
} else if (x.rank === 2) {
return qr2d(x as Tensor2D, fullMatrices);
} else {
// Rank > 2.
// TODO(cais): Below we split the input into individual 2D tensors,
// perform QR decomposition on them and then stack the results back
// together. We should explore whether this can be parallelized.
const outerDimsProd = x.shape.slice(0, x.shape.length - 2)
.reduce((value, prev) => value * prev);
const x2ds = unstack(
x.reshape([
outerDimsProd, x.shape[x.shape.length - 2],
x.shape[x.shape.length - 1]
]),
0);
const q2ds: Tensor2D[] = [];
const r2ds: Tensor2D[] = [];
x2ds.forEach(x2d => {
const [q2d, r2d] = qr2d(x2d as Tensor2D, fullMatrices);
q2ds.push(q2d);
r2ds.push(r2d);
});
const q = stack(q2ds, 0).reshape(x.shape);
const r = stack(r2ds, 0).reshape(x.shape);
return [q, r];
}
}
function qr2d(x: Tensor2D, fullMatrices = false): [Tensor2D, Tensor2D] {
return ENV.engine.tidy(() => {
if (x.shape.length !== 2) {
throw new Error(
`qr2d() requires a 2D Tensor, but got a ${x.shape.length}D Tensor.`);
}
const m = x.shape[0];
const n = x.shape[1];
let q = eye(m) as Tensor2D; // Orthogonal transform so far.
let r = x.clone(); // Transformed matrix so far.
const one2D = tensor2d([[1]], [1, 1]);
let w: Tensor2D = one2D.clone();
const iters = m >= n ? n : m;
for (let j = 0; j < iters; ++j) {
// This tidy within the for-loop ensures we clean up temporary
// tensors as soon as they are no longer needed.
const rTemp = r;
const wTemp = w;
const qTemp = q;
[w, r, q] = ENV.engine.tidy((): [Tensor2D, Tensor2D, Tensor2D] => {
// Find H = I - tau * w * w', to put zeros below R(j, j).
const rjEnd1 = r.slice([j, j], [m - j, 1]);
const normX = rjEnd1.norm();
const rjj = r.slice([j, j], [1, 1]);
const s = rjj.sign().neg() as Tensor2D;
const u1 = rjj.sub(s.mul(normX)) as Tensor2D;
const wPre = rjEnd1.div(u1);
if (wPre.shape[0] === 1) {
w = one2D.clone();
} else {
w = one2D.concat(
wPre.slice([1, 0], [wPre.shape[0] - 1, wPre.shape[1]]) as
Tensor2D,
0);
}
const tau = s.matMul(u1).div(normX).neg() as Tensor2D;
// -- R := HR, Q := QH.
const rjEndAll = r.slice([j, 0], [m - j, n]);
const tauTimesW = tau.mul(w) as Tensor2D;
if (j === 0) {
r = rjEndAll.sub(tauTimesW.matMul(w.transpose().matMul(rjEndAll)));
} else {
r = r.slice([0, 0], [j, n])
.concat(
rjEndAll.sub(tauTimesW.matMul(
w.transpose().matMul(rjEndAll))) as Tensor2D,
0) as Tensor2D;
}
const qAllJEnd = q.slice([0, j], [m, q.shape[1] - j]);
if (j === 0) {
q = qAllJEnd.sub(qAllJEnd.matMul(w).matMul(tauTimesW.transpose()));
} else {
q = q.slice([0, 0], [m, j])
.concat(
qAllJEnd.sub(qAllJEnd.matMul(w).matMul(
tauTimesW.transpose())) as Tensor2D,
1) as Tensor2D;
}
return [w, r, q];
});
dispose([rTemp, wTemp, qTemp]);
}
if (!fullMatrices && m > n) {
q = q.slice([0, 0], [m, n]);
r = r.slice([0, 0], [n, n]);
}
return [q, r];
}) as [Tensor2D, Tensor2D];
}
export const gramSchmidt = op({gramSchmidt_});
export const qr = op({qr_});