@@ -49,6 +49,28 @@ export class Gradients {
49
49
* `x` is computed instead. `f(x)` must take a single tensor `x` and return a
50
50
* single tensor `y`. If `f()` takes multiple inputs, use `grads` instead.
51
51
*
52
+ * ```js
53
+ * // f(x) = x ^ 2
54
+ * const f = x => x.square();
55
+ * // f'(x) = 2x
56
+ * const g = dl.grad(f);
57
+ *
58
+ * const x = dl.tensor1d([2, 3]);
59
+ * g(x).print();
60
+ * ```
61
+ *
62
+ * ```js
63
+ * // f(x) = x ^ 3
64
+ * const f = x => x.pow(dl.scalar(3, 'int32'));
65
+ * // f'(x) = 3x ^ 2
66
+ * const g = dl.grad(f);
67
+ * // f''(x) = 6x
68
+ * const gg = dl.grad(g);
69
+ *
70
+ * const x = dl.tensor1d([2, 3]);
71
+ * gg(x).print();
72
+ * ```
73
+ *
52
74
* @param f The function f(x), to compute gradient for.
53
75
*/
54
76
@doc ( { heading : 'Training' , subheading : 'Gradients' } )
@@ -85,6 +107,21 @@ export class Gradients {
85
107
* The provided `f` must take one or more tensors and return a single tensor
86
108
* `y`. If `f()` takes a single input, we recommend using `grad` instead.
87
109
*
110
+ * ```js
111
+ * // f(a, b) = a * b
112
+ * const f = (a, b) => a.mul(b);
113
+ * // df / da = b, df / db = a
114
+ * const g = dl.grads(f);
115
+ *
116
+ * const a = dl.tensor1d([2, 3]);
117
+ * const b = dl.tensor1d([-2, -3]);
118
+ * const [da, db] = g([a, b]);
119
+ * console.log('da');
120
+ * da.print();
121
+ * console.log('db');
122
+ * db.print();
123
+ * ```
124
+ *
88
125
* @param f The function `f(x1, x2,...)` to compute gradients for.
89
126
*/
90
127
@doc ( { heading : 'Training' , subheading : 'Gradients' } )
@@ -119,6 +156,21 @@ export class Gradients {
119
156
* The result is a rich object with the following properties:
120
157
* - grad: The gradient of `f(x)` w.r.t `x` (result of `grad`).
121
158
* - value: The value returned by `f(x)`.
159
+ *
160
+ * ```js
161
+ * // f(x) = x ^ 2
162
+ * const f = x => x.square();
163
+ * // f'(x) = 2x
164
+ * const g = dl.valueAndGrad(f);
165
+ *
166
+ * const x = dl.tensor1d([2, 3]);
167
+ * const {value, grad} = g(x);
168
+ *
169
+ * console.log('value');
170
+ * value.print();
171
+ * console.log('grad');
172
+ * grad.print();
173
+ * ```
122
174
*/
123
175
@doc ( { heading : 'Training' , subheading : 'Gradients' } )
124
176
static valueAndGrad < I extends Tensor , O extends Tensor > ( f : ( x : I ) => O ) :
@@ -149,6 +201,27 @@ export class Gradients {
149
201
* The result is a rich object with the following properties:
150
202
* - grads: The gradients of `f()` w.r.t each input (result of `grads`).
151
203
* - value: The value returned by `f(x)`.
204
+ *
205
+ * ```js
206
+ * // f(a, b) = a * b
207
+ * const f = (a, b) => a.mul(b);
208
+ * // df/da = b, df/db = a
209
+ * const g = dl.valueAndGrads(f);
210
+ *
211
+ * const a = dl.tensor1d([2, 3]);
212
+ * const b = dl.tensor1d([-2, -3]);
213
+ * const {value, grads} = g([a, b]);
214
+ *
215
+ * const [da, db] = grads;
216
+ *
217
+ * console.log('value');
218
+ * value.print();
219
+ *
220
+ * console.log('da');
221
+ * da.print();
222
+ * console.log('db');
223
+ * db.print();
224
+ * ```
152
225
*/
153
226
@doc ( { heading : 'Training' , subheading : 'Gradients' } )
154
227
static valueAndGrads < O extends Tensor > ( f : ( ...args : Tensor [ ] ) => O ) :
@@ -183,9 +256,20 @@ export class Gradients {
183
256
* trainable variables provided by `varList`. If no list is provided, it
184
257
* defaults to all trainable variables.
185
258
*
259
+ * ```js
260
+ * const a = dl.variable(dl.tensor1d([3, 4]));
261
+ * const b = dl.variable(dl.tensor1d([5, 6]));
262
+ * const x = dl.tensor1d([1, 2]);
263
+ *
264
+ * // f(a, b) = a * x ^ 2 + b * x
265
+ * const f = () => a.mul(x.square()).add(b.mul(x)).sum();
266
+ * // df/da = x ^ 2, df/db = x
267
+ * const {value, grads} = dl.variableGrads(f);
268
+ *
269
+ * Object.keys(grads).forEach(varName => grads[varName].print());
270
+ * ```
271
+ *
186
272
* @param f The function to execute. f() should return a scalar.
187
- * @param varList An optional list of variables to provide gradients with
188
- * respect to. Defaults to all trainable variables.
189
273
*/
190
274
@doc ( { heading : 'Training' , subheading : 'Gradients' } )
191
275
static variableGrads ( f : ( ) => Scalar , varList ?: Variable [ ] ) :
@@ -239,6 +323,21 @@ export class Gradients {
239
323
* called, `g` returns `f().value`. In backward mode, custom gradients with
240
324
* respect to each input of `f` are computed using `f().gradFunc`.
241
325
*
326
+ * ```js
327
+ * const customOp = dl.customGrad(x => {
328
+ * // Override gradient of our custom x ^ 2 op to be dy * abs(x);
329
+ * return {value: x.square(), gradFunc: dy => [dy.mul(x.abs())]};
330
+ * });
331
+ *
332
+ * const x = dl.tensor1d([-1, -2, 3]);
333
+ * const dx = dl.grad(x => customOp(x));
334
+ *
335
+ * console.log(`f(x):`);
336
+ * customOp(x).print();
337
+ * console.log(`f'(x):`);
338
+ * dx(x).print();
339
+ * ```
340
+ *
242
341
* @param f The function to evaluate in forward mode, which should return
243
342
* `{value: Tensor, gradFunc: (dy) => Tensor[]}`, where `gradFunc` returns
244
343
* the custom gradients of `f` with respect to its inputs.
0 commit comments