diff --git a/src/layer_normalization.js b/src/layer_normalization.js
index 4e5e2b5..c159a07 100644
--- a/src/layer_normalization.js
+++ b/src/layer_normalization.js
@@ -24,8 +24,7 @@ export function getIndexOfSortedValue(axes) {
  * Normalize the tensor values of input features using
  * [layer-Normalization](https://arxiv.org/abs/1607.06450)
  * @param {Tensor} input
- * @param {Array} axes
- * @param {MLBatchNormalizationOptions} [options]
+ * @param {MLLayerNormalizationOptions} [options]
  * @return {Tensor}
  */
 export function layerNormalization(input, {scale, bias, axes, epsilon=1e-5}) {