diff --git a/src/ops/lrn.ts b/src/ops/lrn.ts index b7b11ce2a1..56bae55108 100644 --- a/src/ops/lrn.ts +++ b/src/ops/lrn.ts @@ -29,9 +29,8 @@ export class LRNOps { * @param x The input tensor. The 4-D input tensor is treated as a 3-D array * of 1D vectors (along the last dimension), and each vector is * normalized independently. - * @param radius The number of adjacent channels or spatial locations of the - * 1D normalization window. In Tensorflow this param is called - * 'depth_radius' because only 'acrossChannels' mode is supported. + * @param depthRadius The number of adjacent channels in the 1D normalization + * window. * @param bias A constant bias term for the basis. * @param alpha A scale factor, usually positive. * @param beta An exponent. @@ -39,16 +38,16 @@ export class LRNOps { @doc({heading: 'Operations', subheading: 'Normalization'}) @operation static localResponseNormalization( - x: T, radius = 5, bias = 1, alpha = 1, beta = 0.5): T { + x: T, depthRadius = 5, bias = 1, alpha = 1, beta = 0.5): T { util.assertArgumentsAreTensors({x}, 'localResponseNormalization'); util.assert( x.rank === 4 || x.rank === 3, `Error in localResponseNormalization: x must be rank 3 or 4 but got rank ${x.rank}.`); util.assert( - util.isInt(radius), - `Error in localResponseNormalization3D: radius must be an integer - but got radius ${radius}.`); + util.isInt(depthRadius), + `Error in localResponseNormalization: depthRadius must be an integer + but got depthRadius ${depthRadius}.`); let x4D = x as Tensor4D; let reshapedTo4D = false; @@ -58,7 +57,7 @@ export class LRNOps { } const res = ENV.engine.runKernel( backend => backend.localResponseNormalization4D( - x4D, radius, bias, alpha, beta), + x4D, depthRadius, bias, alpha, beta), {x4D}); if (reshapedTo4D) { return res.as3D(res.shape[1], res.shape[2], res.shape[3]) as T;