Fix L2Normalization.

PiperOrigin-RevId: 197052728

diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
index 3b59f24..6e5ceec 100644 (file)
--- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
@@ -2288,7 +2288,7 @@ void L2Normalization(const float* input_data, const Dims<4>& input_dims,
   for (int i = 0; i < outer_size; ++i) {
     float squared_l2_norm = 0;
     for (int c = 0; c < depth; ++c) {
-      const float val = input_data[depth * i + c];
+      const float val = input_data[c];
       squared_l2_norm += val * val;
     }
     const float l2_norm = std::sqrt(squared_l2_norm);

diff --git a/tensorflow/contrib/lite/kernels/l2norm_test.cc b/tensorflow/contrib/lite/kernels/l2norm_test.cc
index 042314c..11cc666 100644 (file)
--- a/tensorflow/contrib/lite/kernels/l2norm_test.cc
+++ b/tensorflow/contrib/lite/kernels/l2norm_test.cc
@@ -76,6 +76,23 @@ TEST(L2NormOpTest, SimpleTest) {
               ElementsAreArray({-0.55, 0.3, 0.35, 0.6, -0.35, 0.05}));
 }
 
+TEST(L2NormOpTest, MultipleBatchesTest) {
+  L2NormOpModel m({3, 1, 1, 6}, TensorType_FLOAT32,
+                  ActivationFunctionType_NONE);
+  m.SetInput({
+      -1.1, 0.6, 0.7, 1.2, -0.7, 0.1,  // batch 1
+      -1.1, 0.6, 0.7, 1.2, -0.7, 0.1,  // batch 2
+      -1.1, 0.6, 0.7, 1.2, -0.7, 0.1,  // batch 3
+  });
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput<float>(),
+              ElementsAreArray({
+                  -0.55, 0.3, 0.35, 0.6, -0.35, 0.05,  // batch 1
+                  -0.55, 0.3, 0.35, 0.6, -0.35, 0.05,  // batch 2
+                  -0.55, 0.3, 0.35, 0.6, -0.35, 0.05,  // batch 3
+              }));
+}
+
 TEST(L2NormOpTest, SimpleUint8Test) {
   L2NormOpModel m({1, 1, 1, 6}, TensorType_UINT8, ActivationFunctionType_NONE);