L2_NORMALIZATION support for tflite

Signed-off-by: Shesung <[email protected]>

Author: shesung

tests/test_backend.py

@@ -5873,6 +5873,14 @@ def func(x):
         x_val = np.array([0.5, 1.0, -0.5, -1.0], dtype=np.float32).reshape((2, 2))
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @skip_tfjs("not supported in tfjs")
+    def test_l2normalization(self):
+        def func(x):
+            op_ = tf.math.l2_normalize(x)
+            return tf.identity(op_, name=_TFOUTPUT)
+
+        x_val = make_xval([3, 4])
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/onnx_opset/math.py

@@ -826,3 +826,15 @@ class HardSwish:
     @classmethod
     def version_14(cls, ctx, node, **kwargs):
         pass
+
+
+@tf_op(["L2Normalization"], onnx_op="LpNormalization")
+class L2Normalization:
+    @classmethod
+    def version_1(cls, ctx, node, **kwargs):
+        axis = node.get_attr_value("axis")
+        if axis is None:
+            # by default use the last dim
+            axis = -1
+        node.set_attr("axis", axis)
+        node.set_attr("p", 2)

tf2onnx/tflite_handlers/tfl_direct.py

@@ -88,6 +88,7 @@
 @tfl_op("TFL_RFFT2D", tf_op="RFFT2D")
 @tfl_op("TFL_COMPLEX_ABS", tf_op="ComplexAbs")
 @tfl_op("TFL_HARD_SWISH", tf_op="HardSwish")
+@tfl_op("TFL_L2_NORMALIZATION", tf_op="L2Normalization")
 class TflDirectOp:
     @classmethod
     def to_tf(cls, ctx, node, **kwargs):