Create semantic layered directories for normalization layers

PiperOrigin-RevId: 371140922

Author: yashk2810

tensorflow_model_optimization/python/core/quantization/keras/layers/BUILD

@@ -31,7 +31,7 @@ py_strict_library(
         # python/keras:initializers tensorflow dep2,
         # python/keras/layers tensorflow dep2,
         # python/keras/layers:convolutional tensorflow dep2,
-        # python/keras/layers:normalization tensorflow dep2,
+        # python/keras/layers/normalization:batch_normalization_v1 tensorflow dep2,
         # python/keras/utils:engine_utils tensorflow dep2,
         "//tensorflow_model_optimization/python/core/keras:utils",
         "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",

tensorflow_model_optimization/python/core/quantization/keras/layers/conv_batchnorm.py

@@ -19,21 +19,20 @@
 from __future__ import print_function
 
 import tensorflow as tf
+
 from tensorflow.python.framework import dtypes
 from tensorflow.python.keras import activations
 from tensorflow.python.keras import backend as K
 from tensorflow.python.keras import initializers
 from tensorflow.python.keras.layers import convolutional
-from tensorflow.python.keras.layers import normalization
 from tensorflow.python.keras.layers import serialization
+from tensorflow.python.keras.layers.normalization import batch_normalization_v1
 from tensorflow.python.keras.utils import conv_utils
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn
 from tensorflow.python.ops import nn_ops
-
 from tensorflow_model_optimization.python.core.keras import utils
-
 from tensorflow_model_optimization.python.core.quantization.keras import quantizers
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantizers
 
@@ -232,7 +231,7 @@ def __init__(
         name=name,
         **kwargs)
 
-    self.batchnorm = normalization.BatchNormalization(
+    self.batchnorm = batch_normalization_v1.BatchNormalization(
         axis=axis,
         momentum=momentum,
         epsilon=epsilon,
@@ -417,7 +416,7 @@ def __init__(
         name=name,
         **kwargs)
 
-    self.batchnorm = normalization.BatchNormalization(
+    self.batchnorm = batch_normalization_v1.BatchNormalization(
         axis=axis,
         momentum=momentum,
         epsilon=epsilon,