Repository: intel/ipex-llm
Pull request: fix null pointer exception in bn when affine is false
## What changes were proposed in this pull request?
fix null pointer exception in bn when affine is false

## How was this patch tested?
Unit tests

## Related links or issues (optional)
fixed https://github.com/intel-analytics/BigDL/issues/2012



Problem statement:
SpatialBatchNormalization does not work while affine is false
In BatchNormalization, weight initialization is based on `nOutput` and `affine`
```scala
 val weight: Tensor[T] =
    if (initWeight != null) initWeight else if (affine) Tensor[T](nOutput) else null
```
Weight will be null if affine is false, but it's not handled properly in forward and backward.

Hints from issue discussion:

Commit message: fix null pointer exception in bn when affine is false

Patch to review:
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
index 376e0b5dbd8..d981aab782f 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
@@ -704,14 +704,13 @@ object SpatialBatchNormalization {
     require(gradOutput.nDimension() == 4, "BN require a 4D gradient")
     require(gradOutput.isContiguous(), "gradient is not contiguous")
     val nChannel = gradOutput.size(4)
-    require(scale.size(1) == nChannel, "scale length is not consist with channel number")
     require(saveMean.size(1) == nChannel, "saveMean length is not consist with channel number")
     require(saveStd.size(1) == nChannel, "saveStd length is not consist with channel number")
 
     gradInput.resizeAs(gradOutput)
     if (gMean.isEmpty) {
-      gMean.resize(scale.size(1))
-      gxMean.resize(scale.size(1))
+      gMean.resize(nChannel)
+      gxMean.resize(nChannel)
     }
 
     val inputData = input.storage().array()
@@ -749,6 +748,8 @@ object SpatialBatchNormalization {
     }
 
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
+
       val scaleData = scale.storage().array()
       val scaleOffset = scale.storageOffset() - 1
       i = 0
@@ -796,14 +797,13 @@ object SpatialBatchNormalization {
     require(gradOutput.nDimension() == 4, "BN require a 4D gradient")
     require(gradOutput.isContiguous(), "gradient is not contiguous")
     val nChannel = gradOutput.size(4)
-    require(scale.size(1) == nChannel, "scale length is not consist with channel number")
     require(saveMean.size(1) == nChannel, "saveMean length is not consist with channel number")
     require(saveStd.size(1) == nChannel, "saveStd length is not consist with channel number")
 
     gradInput.resizeAs(gradOutput)
     if (gMean.isEmpty) {
-      gMean.resize(scale.size(1))
-      gxMean.resize(scale.size(1))
+      gMean.resize(nChannel)
+      gxMean.resize(nChannel)
     }
 
     val inputData = input.storage().array()
@@ -841,6 +841,7 @@ object SpatialBatchNormalization {
     }
 
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
       val scaleData = scale.storage().array()
       val scaleOffset = scale.storageOffset() - 1
       i = 0
@@ -882,7 +883,6 @@ object SpatialBatchNormalization {
     require(gradOutput.nDimension() == 4, "BN require a 4D gradient")
     require(gradOutput.isContiguous(), "gradient is not contiguous")
     val nChannel = gradOutput.size(4)
-    require(scale.size(1) == nChannel, "scale length is not consist with channel number")
     require(saveStd.size(1) == nChannel, "saveStd length is not consist with channel number")
 
     gradInput.resizeAs(gradOutput)
@@ -896,6 +896,7 @@ object SpatialBatchNormalization {
     val n = gradOutput.nElement()
 
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
       val scaleData = scale.storage().array()
       val scaleOffset = scale.storageOffset() - 1
       var i = 0
@@ -933,7 +934,6 @@ object SpatialBatchNormalization {
     require(gradOutput.nDimension() == 4, "BN require a 4D gradient")
     require(gradOutput.isContiguous(), "gradient is not contiguous")
     val nChannel = gradOutput.size(4)
-    require(scale.size(1) == nChannel, "scale length is not consist with channel number")
     require(saveStd.size(1) == nChannel, "saveStd length is not consist with channel number")
 
     gradInput.resizeAs(gradOutput)
@@ -948,6 +948,7 @@ object SpatialBatchNormalization {
     var i = 0
 
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
       val scaleData = scale.storage().array()
       val scaleOffset = scale.storageOffset() - 1
       while (i < n) {
@@ -989,14 +990,13 @@ object SpatialBatchNormalization {
     require(gradOutput.nDimension() == 4, "BN require a 4D gradient")
     require(gradOutput.isContiguous(), "gradient is not contiguous")
     val nChannel = gradOutput.size(2)
-    require(scale.size(1) == nChannel, "scale length is not consist with channel number")
     require(saveMean.size(1) == nChannel, "saveMean length is not consist with channel number")
     require(saveStd.size(1) == nChannel, "saveStd length is not consist with channel number")
 
     gradInput.resizeAs(gradOutput)
     if (gMean.isEmpty) {
-      gMean.resize(scale.size(1))
-      gxMean.resize(scale.size(1))
+      gMean.resize(nChannel)
+      gxMean.resize(nChannel)
     }
 
     val inputData = input.storage().array()
@@ -1044,6 +1044,7 @@ object SpatialBatchNormalization {
     i = 0
     b = 0
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
       val scaleData = scale.storage().array()
       val scaleOffset = scale.storageOffset() - 1
       while (b < nBatch) {
@@ -1504,7 +1505,6 @@ object SpatialBatchNormalization {
     require(gradOutput.nDimension() == 4, "BN require a 4D gradient")
     require(gradOutput.isContiguous(), "gradient is not contiguous")
     val nChannel = gradOutput.size(2)
-    require(scale.size(1) == nChannel, "scale length is not consist with channel number")
     require(saveStd.size(1) == nChannel, "saveStd length is not consist with channel number")
 
     gradInput.resizeAs(gradOutput)
@@ -1520,6 +1520,7 @@ object SpatialBatchNormalization {
     var b = 0
     var i = 0
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
       val scaleData = scale.storage().array()
       val scaleOffset = scale.storageOffset() - 1
       while (b < nBatch) {
@@ -1565,7 +1566,6 @@ object SpatialBatchNormalization {
     require(gradOutput.nDimension() == 4, "BN require a 4D gradient")
     require(gradOutput.isContiguous(), "gradient is not contiguous")
     val nChannel = gradOutput.size(2)
-    require(scale.size(1) == nChannel, "scale length is not consist with channel number")
     require(saveStd.size(1) == nChannel, "saveStd length is not consist with channel number")
 
     gradInput.resizeAs(gradOutput)
@@ -1581,6 +1581,7 @@ object SpatialBatchNormalization {
     var b = 0
     var i = 0
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
       val scaleData = scale.storage().array()
       val scaleOffset = scale.storageOffset() - 1
       while (b < nBatch) {
diff --git a/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/BatchNormalizationSpec.scala b/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/BatchNormalizationSpec.scala
index 3b73e7599f7..4e908752024 100644
--- a/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/BatchNormalizationSpec.scala
+++ b/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/BatchNormalizationSpec.scala
@@ -236,6 +236,15 @@ class BatchNormalizationSpec extends FlatSpec with Matchers {
 
     bn2.gradWeight should be(bn1.gradWeight.mul(0.5))
     bn2.gradBias should be(bn1.gradBias.mul(2))
+  }
 
+  "BatchNormalization backward" should "be good when affine is false" in {
+    val layer = BatchNormalization[Float](3, affine = false)
+    val input = Tensor[Float](4, 3).fill(1)
+    val gradOutput = Tensor[Float](4, 3).fill(1)
+    val output = layer.forward(input)
+    output should be(Tensor[Float](4, 3).fill(0))
+    val gradInput = layer.backward(input, gradOutput)
+    gradInput should be(Tensor[Float](4, 3).fill(0))
   }
 }
diff --git a/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalizationSpec.scala b/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalizationSpec.scala
index 10d7acf2b44..527e7527ec3 100644
--- a/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalizationSpec.scala
+++ b/spark/dl/src/test/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalizationSpec.scala
@@ -34,6 +34,16 @@ class SpatialBatchNormalizationSpec extends FlatSpec with Matchers {
     checker.checkLayer[Double](sbn, input, 1e-3) should be(true)
   }
 
+  "SpatialBatchNormalization backward" should "be good when affine is false" in {
+    val layer = SpatialBatchNormalization[Float](3, affine = false)
+    val input = Tensor[Float](4, 3, 24, 24).fill(1)
+    val gradOutput = Tensor[Float](4, 3, 24, 24).fill(1)
+    val output = layer.forward(input)
+    output should be(Tensor[Float](4, 3, 24, 24).fill(0))
+    val gradInput = layer.backward(input, gradOutput)
+    gradInput should be(Tensor[Float](4, 3, 24, 24).fill(0))
+  }
+
   "SpatialBatchNormalization module in batch mode" should "be good in gradient check " +
     "for weight" in {
     val seed = 100


Code review comments:
diff_hunk: @@ -896,6 +896,7 @@ object SpatialBatchNormalization {
     val n = gradOutput.nElement()
 
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
text: @user1:
consist => consistent

diff_hunk: @@ -1520,6 +1520,7 @@ object SpatialBatchNormalization {
     var b = 0
     var i = 0
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
text: @user1:
consist => consistent

diff_hunk: @@ -1581,6 +1581,7 @@ object SpatialBatchNormalization {
     var b = 0
     var i = 0
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
text: @user1:
consist => consistent

diff_hunk: @@ -749,6 +748,8 @@ object SpatialBatchNormalization {
     }
 
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
text: @user1:
consist => consistent

diff_hunk: @@ -1044,6 +1044,7 @@ object SpatialBatchNormalization {
     i = 0
     b = 0
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
text: @user1:
consist => consistent

diff_hunk: @@ -841,6 +841,7 @@ object SpatialBatchNormalization {
     }
 
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
text: @user1:
consist => consistent

diff_hunk: @@ -948,6 +948,7 @@ object SpatialBatchNormalization {
     var i = 0
 
     if (scale != null) {
+      require(scale.size(1) == nChannel, "scale length is not consist with channel number")
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/SpatialBatchNormalization.scala
text: @user1:
consist => consistent