From b97e2a1b89e1600138ea7e426a308bc9407ae0f2 Mon Sep 17 00:00:00 2001
From: Leon Chen <leon@md.ai>
Date: Sat, 8 Oct 2016 23:21:11 -0400
Subject: [PATCH] fix BatchNormalization layer axis normalization

---
 src/layers/normalization/BatchNormalization.js | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/layers/normalization/BatchNormalization.js b/src/layers/normalization/BatchNormalization.js
index 60dafd4..af6ac23 100644
--- a/src/layers/normalization/BatchNormalization.js
+++ b/src/layers/normalization/BatchNormalization.js
@@ -23,7 +23,12 @@ export default class BatchNormalization extends Layer {
 
     this.epsilon = epsilon
     this.mode = mode
+
+    // no batch axis, so axis is less 1 compared to representation in keras
+    // will be set in call(), as input tensor shape is needed to calculate axis
+    // if axis < 0
     this.axis = axis
+    this.axisNormalized = false
 
     // Layer weights specification
     // running mean and std are non_trainable_weights in mode 0
@@ -38,8 +43,10 @@ export default class BatchNormalization extends Layer {
    * @returns {Tensor} x
    */
   call (x) {
-    // no batch axis, so axis is less 1 compared to representation in keras
-    this.axis = this.axis < 0 ? x.tensor.shape.length + this.axis : this.axis - 1
+    if (!this.axisNormalized) {
+      this.axis = this.axis < 0 ? x.tensor.shape.length + this.axis : this.axis - 1
+      this.axisNormalized = true
+    }
 
     let broadcast = []
     for (let d = 0; d < x.tensor.shape.length; d++) {