(
xShape: number[], radius: number, bias: number, alpha: number,
beta: number)
| 25 | packedOutput = true; |
| 26 | |
| 27 | constructor( |
| 28 | xShape: number[], radius: number, bias: number, alpha: number, |
| 29 | beta: number) { |
| 30 | const rad = radius; |
| 31 | const maxD = xShape[3] - 1; |
| 32 | this.outputShape = xShape; |
| 33 | |
| 34 | // optimize pow(bias + alpha * sum, -beta) |
| 35 | // src: https://github.com/tensorflow/tensorflow/.. |
| 36 | // blob/26033a1644a9c4a5fbe3170ab2e864b6a4ccd4ca/.. |
| 37 | // tensorflow/core/kernels/mkl_lrn_op.cc#L320 |
| 38 | let powOperator; |
| 39 | const basis = `float(${bias}) + float(${alpha}) * sum`; |
| 40 | if (beta === 0.5) { |
| 41 | powOperator = `inversesqrt(${basis})`; |
| 42 | } else if (beta === 1.0) { |
| 43 | powOperator = `1.0/(${basis})`; |
| 44 | } else { |
| 45 | powOperator = `exp(log(${basis}) * float(-${beta}));`; |
| 46 | } |
| 47 | |
| 48 | this.userCode = ` |
| 49 | void main() { |
| 50 | ivec4 coords = getOutputCoords(); |
| 51 | int b = coords.x; |
| 52 | int r = coords.y; |
| 53 | int c = coords.z; |
| 54 | int d = coords.w; |
| 55 | |
| 56 | bool hasNextCol = d < ${this.outputShape[3]}; |
| 57 | bool hasNextRow = c < ${this.outputShape[2]}; |
| 58 | |
| 59 | vec4 sum = vec4(0.); |
| 60 | vec4 xFragAtOutputCoords = getX(b, r, c, d); |
| 61 | |
| 62 | vec4 xAtOutputCoords = vec4( |
| 63 | getChannel(xFragAtOutputCoords, vec2(c, d)), |
| 64 | hasNextCol ? |
| 65 | getChannel(xFragAtOutputCoords, vec2(c, d + 1)) : 0.0, |
| 66 | hasNextRow ? |
| 67 | getChannel(xFragAtOutputCoords , vec2(c + 1, d)) : 0.0, |
| 68 | (hasNextRow && hasNextCol) ? |
| 69 | getChannel(xFragAtOutputCoords, vec2(c + 1, d + 1)) : 0.0 |
| 70 | ); |
| 71 | |
| 72 | int firstChannel = d - ${rad}; |
| 73 | vec2 cache = vec2(0.); |
| 74 | if(firstChannel >= 0){ |
| 75 | vec4 firstChannelFrag = getX(b, r, c, firstChannel); |
| 76 | cache.x = getChannel(firstChannelFrag, vec2(c, firstChannel)); |
| 77 | if(hasNextRow){ |
| 78 | cache.y = getChannel(firstChannelFrag, vec2(c + 1, firstChannel)); |
| 79 | } |
| 80 | } |
| 81 | |
| 82 | ivec2 depth = ivec2(d, d + 1); |
| 83 | for (int j = - ${rad}; j <= ${rad}; j++) { |
| 84 | ivec2 idx = depth + j; |
nothing calls this directly
no outgoing calls
no test coverage detected