add Einsum

Author: RandySheriffH

tests/test_backend.py

@@ -3138,20 +3138,22 @@ def func(x):
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
     @check_opset_min_version(12)
-    def test_less_or_equal(self):
+    def test_squared_distance(self):
         x_val = np.random.random([4, 5]).astype(np.float32)
         y_val = np.random.random([4, 5]).astype(np.float32)
         def func(x, y):
-            return tf.math.less_equal(x, y, name=_TFOUTPUT)
+            return tf.math.squared_difference(x, y, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
 
     @check_opset_min_version(12)
-    def test_squared_distance(self):
-        x_val = np.random.random([4, 5]).astype(np.float32)
-        y_val = np.random.random([4, 5]).astype(np.float32)
+    @check_tf_min_version("2.1")
+    def test_einsum(self):
+        x_val = np.random.random([10]).astype(np.float32)
+        y_val = np.random.random([10]).astype(np.float32)
         def func(x, y):
-            return tf.math.squared_difference(x, y, name=_TFOUTPUT)
-        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
+            ret = tf.einsum("i,j->ij", x, y)
+            return tf.identity(ret, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT],  {_INPUT: x_val, _INPUT1: y_val})
 
 
 if __name__ == '__main__':

tf2onnx/onnx_opset/math.py

@@ -561,3 +561,12 @@ class SquaredDistance:
     def version_12(cls, ctx, node, **kwargs):
         node.attr["reduction"] = "none"
 
+@tf_op("Einsum")
+class Einsum:
+    @classmethod
+    def version_12(cls, ctx, node, **kwargs):
+        del node.attr["N"]
+
+
+
+