@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
544 lines • 25.2 kB
JavaScript
"use strict";
/**
* @license
* Copyright 2019 Google Inc. 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.
* =============================================================================
*/
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : new P(function (resolve) { resolve(result.value); }).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
var __generator = (this && this.__generator) || function (thisArg, body) {
var _ = { label: 0, sent: function() { if (t[0] & 1) throw t[1]; return t[1]; }, trys: [], ops: [] }, f, y, t, g;
return g = { next: verb(0), "throw": verb(1), "return": verb(2) }, typeof Symbol === "function" && (g[Symbol.iterator] = function() { return this; }), g;
function verb(n) { return function (v) { return step([n, v]); }; }
function step(op) {
if (f) throw new TypeError("Generator is already executing.");
while (_) try {
if (f = 1, y && (t = op[0] & 2 ? y["return"] : op[0] ? y["throw"] || ((t = y["return"]) && t.call(y), 0) : y.next) && !(t = t.call(y, op[1])).done) return t;
if (y = 0, t) op = [op[0] & 2, t.value];
switch (op[0]) {
case 0: case 1: t = op; break;
case 4: _.label++; return { value: op[1], done: false };
case 5: _.label++; y = op[1]; op = [0]; continue;
case 7: op = _.ops.pop(); _.trys.pop(); continue;
default:
if (!(t = _.trys, t = t.length > 0 && t[t.length - 1]) && (op[0] === 6 || op[0] === 2)) { _ = 0; continue; }
if (op[0] === 3 && (!t || (op[1] > t[0] && op[1] < t[3]))) { _.label = op[1]; break; }
if (op[0] === 6 && _.label < t[1]) { _.label = t[1]; t = op; break; }
if (t && _.label < t[2]) { _.label = t[2]; _.ops.push(op); break; }
if (t[2]) _.ops.pop();
_.trys.pop(); continue;
}
op = body.call(thisArg, _);
} catch (e) { op = [6, e]; y = 0; } finally { f = t = 0; }
if (op[0] & 5) throw op[1]; return { value: op[0] ? op[1] : void 0, done: true };
}
};
var _this = this;
Object.defineProperty(exports, "__esModule", { value: true });
var engine_1 = require("./engine");
var tf = require("./index");
var jasmine_util_1 = require("./jasmine_util");
var test_util_1 = require("./test_util");
jasmine_util_1.describeWithFlags('gradients', jasmine_util_1.ALL_ENVS, function () {
it('matmul + relu', function () { return __awaiter(_this, void 0, void 0, function () {
var a, b, _a, da, db, dedm, transposeA, transposeB, _b, _c, _d, _e;
return __generator(this, function (_f) {
switch (_f.label) {
case 0:
a = tf.tensor2d([-1, 2, -3, 10, -20, 30], [2, 3]);
b = tf.tensor2d([2, -3, 4, -1, 2, -3], [3, 2]);
_a = tf.grads(function (a, b) {
// m = dot(a, b)
// y = relu(m)
// e = sum(y)
var m = tf.matMul(a, b);
var y = tf.relu(m);
return tf.sum(y);
})([a, b]), da = _a[0], db = _a[1];
dedm = tf.step(tf.matMul(a, b));
// de/da = dot(de/dy, bT)
expect(da.shape).toEqual(a.shape);
transposeA = false;
transposeB = true;
_b = test_util_1.expectArraysClose;
return [4 /*yield*/, da.data()];
case 1:
_c = [_f.sent()];
return [4 /*yield*/, tf.matMul(dedm, b, transposeA, transposeB).data()];
case 2:
_b.apply(void 0, _c.concat([_f.sent()]));
// de/db = dot(aT, de/dy)
expect(db.shape).toEqual(b.shape);
transposeA = true;
transposeB = false;
_d = test_util_1.expectArraysClose;
return [4 /*yield*/, db.data()];
case 3:
_e = [_f.sent()];
return [4 /*yield*/, tf.matMul(a, dedm, transposeA, transposeB).data()];
case 4:
_d.apply(void 0, _e.concat([_f.sent()]));
return [2 /*return*/];
}
});
}); });
it('grad(f)', function () { return __awaiter(_this, void 0, void 0, function () {
var grad, result, _a;
return __generator(this, function (_b) {
switch (_b.label) {
case 0:
grad = tf.grad(function (x) { return x.square(); });
result = grad(tf.tensor1d([.1, .2]));
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, result.data()];
case 1:
_a.apply(void 0, [_b.sent(), [.2, .4]]);
return [2 /*return*/];
}
});
}); });
it('calling grad(f) twice works', function () { return __awaiter(_this, void 0, void 0, function () {
var grad, result, result2, _a, _b;
return __generator(this, function (_c) {
switch (_c.label) {
case 0:
grad = tf.grad(function (x) { return x.square(); });
result = grad(tf.tensor1d([.1, .2]));
result2 = grad(tf.tensor1d([.1, .4]));
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, result.data()];
case 1:
_a.apply(void 0, [_c.sent(), [.2, .4]]);
_b = test_util_1.expectArraysClose;
return [4 /*yield*/, result2.data()];
case 2:
_b.apply(void 0, [_c.sent(), [.2, .8]]);
return [2 /*return*/];
}
});
}); });
it('grad(f): throwing an error during forward pass', function () {
var grad = tf.grad(function (x) {
throw new Error('failed forward pass');
});
expect(function () { return grad(tf.zeros([])); }).toThrowError();
expect(engine_1.ENGINE.isTapeOn()).toBe(false);
});
it('grad(f): throwing an error during backwards pass', function () {
var customOp = tf.customGrad(function (x) {
return {
value: x,
gradFunc: function () {
throw new Error('failed backward pass');
}
};
});
var grad = tf.grad(function (x) { return customOp(x); });
expect(function () { return grad(tf.zeros([])); }).toThrowError();
expect(engine_1.ENGINE.isTapeOn()).toBe(false);
});
it('grads(f)', function () { return __awaiter(_this, void 0, void 0, function () {
var grads, result, _a;
return __generator(this, function (_b) {
switch (_b.label) {
case 0:
grads = tf.grads(function (x) { return x.square(); });
result = grads([tf.tensor1d([.1, .2])]);
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, result[0].data()];
case 1:
_a.apply(void 0, [_b.sent(), [.2, .4]]);
return [2 /*return*/];
}
});
}); });
it('calling grads(f) twice works', function () { return __awaiter(_this, void 0, void 0, function () {
var grads, result, result2, _a, _b;
return __generator(this, function (_c) {
switch (_c.label) {
case 0:
grads = tf.grads(function (x) { return x.square(); });
result = grads([tf.tensor1d([.1, .2])]);
result2 = grads([tf.tensor1d([.1, .4])]);
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, result[0].data()];
case 1:
_a.apply(void 0, [_c.sent(), [.2, .4]]);
_b = test_util_1.expectArraysClose;
return [4 /*yield*/, result2[0].data()];
case 2:
_b.apply(void 0, [_c.sent(), [.2, .8]]);
return [2 /*return*/];
}
});
}); });
it('works with reshape', function () { return __awaiter(_this, void 0, void 0, function () {
var a, exponent, da, _a;
return __generator(this, function (_b) {
switch (_b.label) {
case 0:
a = tf.tensor2d([1, 2, 3, 4], [2, 2]);
exponent = tf.tensor1d([2, 2, 2, 2], 'int32');
da = tf.grad(function (a) {
var b = a.flatten();
var m = tf.pow(b, exponent);
return tf.sum(m);
})(a);
expect(da.shape).toEqual([2, 2]);
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, da.data()];
case 1:
_a.apply(void 0, [_b.sent(), [2, 4, 6, 8]]);
return [2 /*return*/];
}
});
}); });
it('reshape outside tf.grads() throws error', function () {
var a = tf.tensor2d([1, 2, 3, 4], [2, 2]);
var b = a.flatten();
var exponent = tf.tensor1d([2, 2, 2, 2], 'int32');
var f = function () {
tf.grads(function (a, b) {
var m = tf.pow(b, exponent);
return tf.sum(m);
})([a, b]);
};
expect(f).toThrowError();
});
it('does not error if irrelevant (pruned) ops are missing grads', function () { return __awaiter(_this, void 0, void 0, function () {
var a, b, da, _a;
return __generator(this, function (_b) {
switch (_b.label) {
case 0:
a = tf.tensor1d([true, true], 'bool');
b = tf.tensor1d([false, true], 'bool');
da = tf.grad(function (a) {
// Logical has no gradients, but it is irrelevant.
a.logicalAnd(b);
return a.sum();
})(a);
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, da.data()];
case 1:
_a.apply(void 0, [_b.sent(), [1, 1]]);
return [2 /*return*/];
}
});
}); });
it('errors if relevant ops are missing grads', function () {
var a = tf.tensor1d([true, true], 'bool');
var b = tf.tensor1d([false, true], 'bool');
var dfda = tf.grad(function (a) {
// Logical has no gradients, but it's relevant to the output.
return a.logicalAnd(b);
});
expect(function () { return dfda(a); }).toThrowError();
});
it('works with asType', function () { return __awaiter(_this, void 0, void 0, function () {
var a, exponent, da, _a;
return __generator(this, function (_b) {
switch (_b.label) {
case 0:
a = tf.tensor2d([1, 2, 3, 4], [2, 2], 'int32');
exponent = tf.tensor2d([2, 2, 2, 2], [2, 2], 'int32');
da = tf.grad(function (a) {
var b = a.toFloat();
var m = tf.pow(b, exponent);
return tf.sum(m);
})(a);
expect(da.shape).toEqual([2, 2]);
expect(da.dtype).toEqual('float32');
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, da.data()];
case 1:
_a.apply(void 0, [_b.sent(), [2, 4, 6, 8]]);
return [2 /*return*/];
}
});
}); });
it('asType outside of tf.grads() throws error', function () {
var a = tf.tensor2d([1, 2, 3, 4], [2, 2], 'int32');
var b = a.toFloat();
var exponent = tf.tensor2d([2, 2, 2, 2], [2, 2], 'int32');
var f = function () {
tf.grad(function (a) {
var m = tf.pow(b, exponent);
return tf.sum(m);
})(a);
};
expect(f).toThrowError();
});
it('saves tensors from the forward pass as expected', function () {
var x = tf.scalar(1).variable();
var optimizer = tf.train.sgd(0.1);
optimizer.minimize(function () {
var y = x.square();
var z = y.square();
y.dispose();
return z;
});
});
it('custom ops do not leak', function () {
var before = tf.memory().numTensors;
var x = tf.softmax([1, 2, 3, 4]);
x.dispose();
var now = tf.memory().numTensors;
expect(now).toBe(before);
});
});
jasmine_util_1.describeWithFlags('valueAndGradients', jasmine_util_1.ALL_ENVS, function () {
it('matmul + relu', function () { return __awaiter(_this, void 0, void 0, function () {
var a, b, _a, value, grads, _b, dedm, da, db, transposeA, transposeB, _c, _d, _e, _f;
return __generator(this, function (_g) {
switch (_g.label) {
case 0:
a = tf.tensor2d([-1, 2, -3, 10, -20, 30], [2, 3]);
b = tf.tensor2d([2, -3, 4, -1, 2, -3], [3, 2]);
_a = tf.valueAndGrads(function (a, b) {
// m = dot(a, b)
// y = relu(m)
// e = sum(y)
var m = tf.matMul(a, b);
var y = tf.relu(m);
return tf.sum(y);
})([a, b]), value = _a.value, grads = _a.grads;
_b = test_util_1.expectArraysClose;
return [4 /*yield*/, value.data()];
case 1:
_b.apply(void 0, [_g.sent(), 10]);
dedm = tf.step(tf.matMul(a, b));
da = grads[0], db = grads[1];
transposeA = false;
transposeB = true;
_c = test_util_1.expectArraysClose;
return [4 /*yield*/, da.data()];
case 2:
_d = [_g.sent()];
return [4 /*yield*/, tf.matMul(dedm, b, transposeA, transposeB).data()];
case 3:
_c.apply(void 0, _d.concat([_g.sent()]));
// de/db = dot(aT, de/dy)
transposeA = true;
transposeB = false;
_e = test_util_1.expectArraysClose;
return [4 /*yield*/, db.data()];
case 4:
_f = [_g.sent()];
return [4 /*yield*/, tf.matMul(a, dedm, transposeA, transposeB).data()];
case 5:
_e.apply(void 0, _f.concat([_g.sent()]));
return [2 /*return*/];
}
});
}); });
it('matmul + relu + inner tidy', function () { return __awaiter(_this, void 0, void 0, function () {
var a, b, _a, value, grads, _b, dedm, da, db, transposeA, transposeB, _c, _d, _e, _f;
return __generator(this, function (_g) {
switch (_g.label) {
case 0:
a = tf.tensor2d([-1, 2, -3, 10, -20, 30], [2, 3]);
b = tf.tensor2d([2, -3, 4, -1, 2, -3], [3, 2]);
_a = tf.valueAndGrads(function (a, b) {
// m = dot(a, b)
// y = relu(m)
// e = sum(y)
var m = tf.matMul(a, b);
return tf.tidy(function () {
var y = tf.relu(m);
return tf.sum(y);
});
})([a, b]), value = _a.value, grads = _a.grads;
_b = test_util_1.expectArraysClose;
return [4 /*yield*/, value.data()];
case 1:
_b.apply(void 0, [_g.sent(), 10]);
dedm = tf.step(tf.matMul(a, b));
da = grads[0], db = grads[1];
transposeA = false;
transposeB = true;
_c = test_util_1.expectArraysClose;
return [4 /*yield*/, da.data()];
case 2:
_d = [_g.sent()];
return [4 /*yield*/, tf.matMul(dedm, b, transposeA, transposeB).data()];
case 3:
_c.apply(void 0, _d.concat([_g.sent()]));
// de/db = dot(aT, de/dy)
transposeA = true;
transposeB = false;
_e = test_util_1.expectArraysClose;
return [4 /*yield*/, db.data()];
case 4:
_f = [_g.sent()];
return [4 /*yield*/, tf.matMul(a, dedm, transposeA, transposeB).data()];
case 5:
_e.apply(void 0, _f.concat([_g.sent()]));
return [2 /*return*/];
}
});
}); });
});
jasmine_util_1.describeWithFlags('higher-order gradients', jasmine_util_1.ALL_ENVS, function () {
it('grad(grad(f))', function () { return __awaiter(_this, void 0, void 0, function () {
var x, before, gradgrad, result, _a;
return __generator(this, function (_b) {
switch (_b.label) {
case 0:
x = tf.tensor1d([.1, .2]);
before = tf.memory().numTensors;
gradgrad = tf.grad(tf.grad(function (x) { return x.mul(x).mul(x); }));
result = gradgrad(x);
expect(tf.memory().numTensors).toBe(before + 1);
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, result.data()];
case 1:
_a.apply(void 0, [_b.sent(), [.6, 1.2]]);
return [2 /*return*/];
}
});
}); });
it('grad(grad(x^2))', function () { return __awaiter(_this, void 0, void 0, function () {
var x, gradgrad, result, _a;
return __generator(this, function (_b) {
switch (_b.label) {
case 0:
x = tf.scalar(3);
gradgrad = tf.grad(tf.grad(function (x) { return x.square(); }));
result = gradgrad(x);
// grad(grad(x^2)) = grad(2x) = 2
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, result.data()];
case 1:
// grad(grad(x^2)) = grad(2x) = 2
_a.apply(void 0, [_b.sent(), [2]]);
return [2 /*return*/];
}
});
}); });
it('grads(grads(f))', function () { return __awaiter(_this, void 0, void 0, function () {
var grads, gradsgrads, result, _a;
return __generator(this, function (_b) {
switch (_b.label) {
case 0:
grads = tf.grads(function (x) { return x.mul(x).mul(x); });
gradsgrads = tf.grads(function (x) { return grads([x])[0]; });
result = gradsgrads([tf.tensor1d([.1, .2])]);
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, result[0].data()];
case 1:
_a.apply(void 0, [_b.sent(), [.6, 1.2]]);
return [2 /*return*/];
}
});
}); });
});
jasmine_util_1.describeWithFlags('customGradient', jasmine_util_1.ALL_ENVS, function () {
it('basic', function () { return __awaiter(_this, void 0, void 0, function () {
var a, b, dy, customPow, _a, value, grad, _b, _c;
return __generator(this, function (_d) {
switch (_d.label) {
case 0:
a = tf.scalar(3);
b = tf.scalar(2, 'int32');
dy = tf.scalar(4);
customPow = tf.customGrad(function (a) {
var value = tf.pow(a, b);
var gradFunc = function (dy) { return dy.mul(tf.scalar(0.1)); };
return { value: value, gradFunc: gradFunc };
});
_a = tf.valueAndGrad(function (a) { return customPow(a); })(a, dy), value = _a.value, grad = _a.grad;
expect(value.shape).toEqual(a.shape);
_b = test_util_1.expectArraysClose;
return [4 /*yield*/, value.data()];
case 1:
_b.apply(void 0, [_d.sent(), [9]]);
expect(grad.shape).toEqual(a.shape);
_c = test_util_1.expectArraysClose;
return [4 /*yield*/, grad.data()];
case 2:
_c.apply(void 0, [_d.sent(), [.4]]);
return [2 /*return*/];
}
});
}); });
it('second order derivative through customGradient', function () { return __awaiter(_this, void 0, void 0, function () {
var a, b, dy, customPow, dda, _a, _b;
return __generator(this, function (_c) {
switch (_c.label) {
case 0:
a = tf.scalar(3);
b = tf.scalar(2, 'int32');
dy = tf.scalar(5);
customPow = tf.customGrad(function (a, save) {
var value = tf.pow(a, b);
save([a]);
var gradFunc = function (dy, saved) {
var a = saved[0];
return dy.mul(a);
};
return { value: value, gradFunc: gradFunc };
});
dda = tf.grad(tf.grad(function (a) { return customPow(a); }))(a, dy);
expect(dda.shape).toEqual(a.shape);
// First order: dy * a. Second order: dy.
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, dda.data()];
case 1:
_b = [_c.sent()];
return [4 /*yield*/, dy.data()];
case 2:
// First order: dy * a. Second order: dy.
_a.apply(void 0, _b.concat([_c.sent()]));
return [2 /*return*/];
}
});
}); });
it('calling gradient of custom op twice works', function () { return __awaiter(_this, void 0, void 0, function () {
var customOp, x, grad, _a, _b;
return __generator(this, function (_c) {
switch (_c.label) {
case 0:
customOp = tf.customGrad(function (x, save) {
// Override gradient of our custom x ^ 2 op to be dy * abs(x);
save([x]);
return {
value: x.square(),
gradFunc: function (dy, saved) {
var x = saved[0];
return dy.mul(x.abs());
}
};
});
x = tf.tensor1d([-1, -2, 3]);
grad = tf.grad(function (x) { return customOp(x); });
_a = test_util_1.expectArraysClose;
return [4 /*yield*/, grad(x).data()];
case 1:
_a.apply(void 0, [_c.sent(), [1, 2, 3]]);
_b = test_util_1.expectArraysClose;
return [4 /*yield*/, grad(x).data()];
case 2:
_b.apply(void 0, [_c.sent(), [1, 2, 3]]);
return [2 /*return*/];
}
});
}); });
});
//# sourceMappingURL=gradients_test.js.map