/
variance.js
73 lines (68 loc) · 1.97 KB
/
variance.js
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
import Layer from './base.js'
import Matrix from '../../../util/matrix.js'
import Tensor from '../../../util/tensor.js'
/**
* Variance layer
*/
export default class VarLayer extends Layer {
/**
* @param {object} config config
* @param {number | number[] | string} [config.axis] axis
* @param {boolean} [config.keepdims] Keep dimensions or not
*/
constructor({ axis = -1, keepdims = true, ...rest }) {
super(rest)
this._axis = null
if (typeof axis === 'string') {
this._axisname = axis
} else {
this._axis = typeof axis === 'number' ? [axis] : axis
}
this._keepdims = keepdims
}
calc(x) {
if (this._axisname) {
this._axis = this.graph.getNode(this._axisname).outputValue.toArray()
}
if (this._axis.includes(-1)) {
this._size = Array(x.dimension).fill(1)
} else {
this._size = x.sizes.concat()
for (let k = 0; k < this._axis.length; k++) {
this._size[this._axis[k]] = 1
}
}
this._i = x
if (!this._keepdims && this._axis.includes(-1)) {
this._m = x.reduce((s, v) => s + v, 0) / x.length
return new Tensor([], x.reduce((s, v) => s + (v - this._m) ** 2, 0) / x.length)
}
if (!this._keepdims && x instanceof Matrix) {
x = Tensor.fromArray(x)
}
const len = this._axis.includes(-1) ? x.length : this._axis.reduce((s, v) => s * x.sizes[v], 1)
this._m = x.reduce((s, v) => s + v, 0, this._axis, true)
this._m.map(v => v / len)
const d = x.copy()
d.broadcastOperate(this._m, (a, b) => a - b)
const o = d.reduce((s, v) => s + v ** 2, 0, this._axis, this._keepdims)
o.map(v => v / len)
return o
}
grad(bo) {
bo.reshape(...this._size)
const len = this._axis.includes(-1) ? this._i.length : this._axis.reduce((s, v) => s * this._i.sizes[v], 1)
const bi = this._i.copy()
bi.broadcastOperate(this._m, (a, b) => (2 * (a - b)) / len)
bi.broadcastOperate(bo, (a, b) => a * b)
return bi
}
toObject() {
return {
type: 'variance',
axis: this._axis,
keepdims: this._keepdims,
}
}
}
VarLayer.registLayer('variance')