Add functions to decompose einsum summation (#1472)

* Add function to decompose einsum summation

Signed-off-by: xavier dupré <[email protected]>

* header

Signed-off-by: xavier dupré <[email protected]>

* to_onnx

Signed-off-by: xavier dupré <[email protected]>

* fix doot

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* fix opset

Signed-off-by: xavier dupré <[email protected]>

* remove useless transpose

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* use gemm if possible

Signed-off-by: xavier dupré <[email protected]>

* add main files to handle conversion

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* complete optimizer

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* add optimization

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* einsum

Signed-off-by: xavier dupré <[email protected]>

* ut

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* enable einsum optimizer

Signed-off-by: xavier dupré <[email protected]>

* fix missing output

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

* small design change

Signed-off-by: xavier dupré <[email protected]>

* lint

Signed-off-by: xavier dupré <[email protected]>

Co-authored-by: xavier dupré <[email protected]>
Co-authored-by: TomWildenhain-Microsoft <[email protected]>

Author: 3 people

tests/test_einsum_helper.py

@@ -0,0 +1,333 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""Unit Tests for einsum decomposition."""
+
+import unittest
+import itertools
+import numpy as np
+from numpy.testing import assert_almost_equal
+from onnxruntime import InferenceSession
+from tf2onnx.optimizer.einsum_optimizer import (
+    analyse_einsum_equation, decompose_einsum_equation, EinsumSubOp)
+from backend_test_base import Tf2OnnxBackendTestBase
+from common import check_opset_min_version
+
+
+class TestEinsum(Tf2OnnxBackendTestBase):
+    "unit tests for einsum optimizer"
+
+    def assert_raise(self, fct, exc_type):
+        try:
+            fct()
+        except exc_type:
+            return
+        raise AssertionError("%r was not raised." % exc_type)
+
+    def apply_einsum_sequence(self, seq, *inputs):
+        names = ["X%d" % i for i in range(len(inputs))]
+        onx = seq.to_onnx('Y', *names, opset=self.config.opset)
+        sess = InferenceSession(onx.SerializeToString())
+        inps = {n: i.astype(np.float32) for n, i in zip(names, inputs)}
+        res = sess.run(None, inps)
+        return res[0]
+
+    def test_analyse_einsum_equation(self):
+        "unit test"
+        self.assert_raise(lambda: analyse_einsum_equation("abc"), NotImplementedError)
+        self.assert_raise(lambda: analyse_einsum_equation("abc0,ch->ah"), ValueError)
+        self.assert_raise(lambda: analyse_einsum_equation("abc,ch->a0"), ValueError)
+        res = analyse_einsum_equation("abc,ch->ah")
+        self.assertEqual(len(res), 4)
+        letters, mat, lengths, duplicates = res
+        self.assertEqual(letters, "abch")
+        assert_almost_equal(lengths, np.array([3, 2, 2]))
+        assert_almost_equal(mat, np.array([[0, 1, 2, -1], [-1, -1, 0, 1], [0, -1, -1, 1]]))
+        self.assertEqual(duplicates, [None, None, None])
+
+    def test_analyse_einsum_equation_duplicates(self):
+        res = analyse_einsum_equation("aac,ca->aa")
+        self.assertEqual(len(res), 4)
+        letters, mat, lengths, duplicates = res
+        self.assertEqual(letters, "ac")
+        assert_almost_equal(lengths, np.array([3, 2, 2]))
+        self.assertEqual(duplicates, [{'a': [0, 1], 'c': [2]}, None, {'a': [0, 1]}])
+        assert_almost_equal(mat, np.array([[1, 2], [1, 0], [1, -1]]))
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_decompose_einsum_equation(self):
+        "test decompose einsum"
+        m1 = np.arange(0, 8).astype(np.float32).reshape((2, 2, 2))
+        m2 = np.arange(0, 4).astype(np.float32).reshape((2, 2))
+        exp = np.einsum("bac,ch->ah", m1, m2)
+        seq = decompose_einsum_equation("bac,ch->ah", (2, 2, 2), (2, 2))
+        dot = seq.to_dot()
+        red = dot.split('red')
+        self.assertEqual(len(red), 5)
+        res = self.apply_einsum_sequence(seq, m1, m2)
+        assert_almost_equal(exp, res)
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_decompose_einsum_equation_deep_case(self):
+        m1 = np.arange(0, 16).astype(np.float32).reshape((2, 2, 2, 2))
+        m2 = np.arange(0, 16).astype(np.float32).reshape((2, 2, 2, 2))
+        exp = np.einsum("bsnh,btnh->bnts", m1, m2)
+        seq = decompose_einsum_equation("bsnh,btnh->bnts")
+        res = self.apply_einsum_sequence(seq, m1, m2)
+        assert_almost_equal(exp, res)
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_decompose_einsum_equation_onnx(self):
+        m1 = np.arange(0, 24).astype(np.float32).reshape((2, 3, 4))
+        m2 = np.arange(0, 20).astype(np.float32).reshape((4, 5))
+        seq = decompose_einsum_equation("bac,ch->ah", (2, 3, 4), (4, 5))
+        exp = np.einsum("bac,ch->ah", m1, m2)
+        res = self.apply_einsum_sequence(seq, m1, m2)
+        assert_almost_equal(exp, res)
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_decompose_einsum_equation_noshape(self):
+        m1 = np.arange(0, 24).astype(np.float32).reshape((2, 3, 4))
+        m2 = np.arange(0, 20).astype(np.float32).reshape((4, 5))
+        seq = decompose_einsum_equation("bac,ch->ah")
+        exp = np.einsum("bac,ch->ah", m1, m2)
+        res = self.apply_einsum_sequence(seq, m1, m2)
+        assert_almost_equal(exp, res)
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_decompose_einsum_equation_onnx2(self):
+        "test bac,cd,def->ebc"
+        m1 = np.arange(0, 24).astype(np.float32).reshape((2, 3, 4))
+        m2 = np.arange(0, 20).astype(np.float32).reshape((4, 5))
+        m3 = np.arange(0, 77 * 5).astype(np.float32).reshape((5, 7, 11))
+
+        seq = decompose_einsum_equation(
+            "bac,cd,def->ebc", (2, 3, 4), (4, 5), (5, 7, 11))
+        exp = np.einsum("bac,cd,def->ebc", m1, m2, m3)
+        res = self.apply_einsum_sequence(seq, m1, m2, m3)
+        assert_almost_equal(exp, res)
+
+    def test_einsum_sub_op(self):
+        self.assert_raise(lambda: EinsumSubOp(2, "er", (2, 2)), ValueError)
+        self.assert_raise(lambda: EinsumSubOp(2, "expand_dims"), RuntimeError)
+        self.assert_raise(lambda: EinsumSubOp(2, "matmul", (2, 2)), RuntimeError)
+        self.assert_raise(lambda: EinsumSubOp(2, "id", (2, 2)), TypeError)
+
+    def common_test_case_2(self, equation):
+        m1 = np.arange(2 * 2 * 2).reshape((2, 2, 2)) + 10
+        m2 = np.arange(4).reshape((2, 2)) + 100
+        exp = np.einsum(equation, m1, m2)
+
+        seq = decompose_einsum_equation(equation, m1.shape, m2.shape)
+        res = self.apply_einsum_sequence(seq, m1, m2)
+        assert_almost_equal(exp, res)
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_case_2_a(self):
+        self.common_test_case_2('abc,cd->abc')
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_many_2(self):
+        "test many equation with 2 inputs"
+        m1 = np.arange(2 * 2 * 2).reshape((2, 2, 2)) + 10
+        m2 = np.arange(4).reshape((2, 2)) + 100
+
+        res = []
+        for p1 in itertools.permutations(list("abc")):
+            for p2 in itertools.permutations(list("cd")):
+                for i in [1, 2]:
+                    for j in [0, 1]:
+                        sp1 = "".join(p1)
+                        sp2 = "".join(p2)
+                        if len(set([sp1[0], sp1[i], sp2[j]])) != 3:
+                            continue
+                        equation = "%s,%s->%s%s%s" % (
+                            sp1, sp2, sp1[0], sp1[i], sp2[j])
+                        try:
+                            r = np.einsum(equation, m1, m2)
+                            res.append((equation, r))
+                        except ValueError:
+                            # Not viable equation.
+                            continue
+
+        for i, (eq, exp) in enumerate(res):
+            with self.subTest(equation=eq, index=i, total=len(res)):
+                seq = decompose_einsum_equation(
+                    eq, m1.shape, m2.shape)
+                res = self.apply_einsum_sequence(seq, m1, m2)
+                exp = np.einsum(eq, m1, m2)
+                assert_almost_equal(exp, res)
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_many_3(self):
+        "test many equation with 3 inputs"
+        m1 = np.arange(2 * 2 * 2).reshape((2, 2, 2)) + 10
+        m2 = np.arange(4).reshape((2, 2)) + 100
+        m3 = np.arange(8).reshape((2, 2, 2)) + 1000
+
+        res = []
+        for p1 in itertools.permutations(list("abc")):  # pylint: disable=R1702
+            for p2 in itertools.permutations(list("cd")):
+                for p3 in itertools.permutations(list("def")):
+                    for i in [1, 2]:
+                        for j in [0, 1]:
+                            sp1 = "".join(p1)
+                            sp2 = "".join(p2)
+                            sp3 = "".join(p3)
+                            equation = "%s,%s,%s->%s%s%s" % (
+                                sp1, sp2, sp3, sp1[0], sp1[i], sp3[j])
+                            try:
+                                r = np.einsum(equation, m1, m2, m3)
+                                res.append((equation, r))
+                            except ValueError:
+                                # Not viable equation.
+                                continue
+
+        for i, (eq, exp) in enumerate(res):
+            with self.subTest(equation=eq, index=i, total=len(res)):
+                seq = decompose_einsum_equation(
+                    eq, m1.shape, m2.shape, m3.shape)
+                res = self.apply_einsum_sequence(seq, m1, m2, m3)
+                exp = np.einsum(eq, m1, m2, m3)
+                assert_almost_equal(exp, res)
+
+    # Taken from https://github.com/numpy/numpy/blob/main/numpy/
+    # core/tests/test_einsum.py.
+
+    def optimize_compare(self, equation, operands=None):
+        "Compares numpy einsum and ONNX."
+        with self.subTest(equation=equation):
+            if operands is not None:
+                inputs = operands
+            else:
+                eqs = equation.split("->")[0].split(",")
+                inputs = []
+                for d, eq in enumerate(eqs):
+                    i = np.arange(2 ** len(eq)).reshape(
+                        (2,) * len(eq)).astype(np.float32)
+                    inputs.append(
+                        i + np.array([3 ** d], dtype=np.float32))
+
+            exp = np.einsum(equation, *inputs)
+            shapes = [m.shape for m in inputs]
+
+            seq = decompose_einsum_equation(equation, *shapes)
+            got = self.apply_einsum_sequence(seq, *inputs)
+            assert_almost_equal(exp, got, decimal=5)
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_numpy_test_hadamard_like_products(self):
+        self.optimize_compare('a,ab,abc->abc')
+        self.optimize_compare('a,b,ab->ab')
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_np_test_np_test_collapse(self):
+        self.optimize_compare('ab,ab,cd,cd->ac')
+        self.optimize_compare('ab,ab,c->c')
+        self.optimize_compare('ab,ab,cd,cd->cd')
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_np_test_index_transformations(self):
+        self.optimize_compare('ea,fb,gc,hd,abcd->efgh')
+        self.optimize_compare('ea,fb,abcd,gc,hd->efgh')
+        self.optimize_compare('abcd,ea,fb,gc,hd->efgh')
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_np_test_expand(self):
+        self.optimize_compare('ab,cd,ef->abcdef')
+        self.optimize_compare('ab,cd,ef->acdf')
+        self.optimize_compare('ab,cd,de->abcde')
+        self.optimize_compare('ab,cd,de->be')
+        self.optimize_compare('ab,bcd,cd->abcd')
+        self.optimize_compare('ab,bcd,cd->abd')
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_np_test_edge_cases1(self):
+        self.optimize_compare('efc,dbc,acf,fd->abe')
+        self.optimize_compare(
+            'eac->ace', operands=[np.arange(24).reshape((2, 3, 4))])
+        self.optimize_compare('eac->ace')
+        self.optimize_compare('bd,db,eac->ace')
+        self.optimize_compare('ba,ac,da->bcd')
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_np_test_edge_cases2(self):
+        self.optimize_compare(
+            'eac->ace', operands=[np.arange(24).reshape((2, 3, 4))])
+        self.optimize_compare('eb,cb,fb->cef')
+
+    @unittest.skipIf(True, "diagonal still not converted into ONNX")
+    def test_np_test_random_cases(self):
+        self.optimize_compare('aab,fa,df,ecc->bde')
+        self.optimize_compare('bb,ff,be->e')
+        self.optimize_compare('afd,ba,cc,dc->bf')
+        self.optimize_compare('bbd,bda,fc,db->acf')
+        self.optimize_compare('dba,ead,cad->bce')
+        self.optimize_compare('aef,fbc,dca->bde')
+
+    def test_np_test_combined_views_mapping(self):
+        a = np.arange(9).reshape(1, 1, 3, 1, 3)
+        b = np.einsum('bbcdc->d', a)
+        assert_almost_equal(b, [12])
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_np_test_broadcasting_dot_cases1(self):
+        a = np.random.rand(1, 5, 4)
+        b = np.random.rand(4, 6)
+        c = np.random.rand(5, 6)
+        d = np.random.rand(10)
+        self.optimize_compare('ijk,kl,jl,i->i', operands=[a, b, c, d])
+        e = np.random.rand(1, 1, 5, 4)
+        f = np.random.rand(7, 7)
+        self.optimize_compare('abjk,kl,jl,ab->ab', operands=[e, b, c, f])
+
+    @check_opset_min_version(13, "Squeeze")
+    def test_np_test_broadcasting_dot_cases2(self):
+        f = np.arange(7 * 55).reshape(7, 11, 5)
+        g = np.arange(30).reshape(2, 3, 5)
+        self.optimize_compare('obk,ijk->ioj', operands=[f, g])
+
+    def np_test_complex(self):
+        self.optimize_compare('acdf,jbje,gihb,hfac,gfac,gifabc,hfac')
+        self.optimize_compare('acdf,jbje,gihb,hfac,gfac,gifabc,hfac')
+        self.optimize_compare('cd,bdhe,aidb,hgca,gc,hgibcd,hgac')
+        self.optimize_compare('abhe,hidj,jgba,hiab,gab')
+        self.optimize_compare('bde,cdh,agdb,hica,ibd,hgicd,hiac')
+        self.optimize_compare('chd,bde,agbc,hiad,hgc,hgi,hiad')
+        self.optimize_compare('chd,bde,agbc,hiad,bdi,cgh,agdb')
+        self.optimize_compare('bdhe,acad,hiab,agac,hibd')
+
+    def np_test_inner_product(self):
+        self.optimize_compare('ab,ab')
+        self.optimize_compare('ab,ba')
+        self.optimize_compare('abc,abc')
+        self.optimize_compare('abc,bac')
+        self.optimize_compare('abc,cba')
+
+    @unittest.skipIf(True, reason="diagonal still not converted into ONNX")
+    def test_np_test_random_cases_difficult(self):
+        "unit test"
+        self.optimize_compare('db,bc,cfc->d')
+        self.optimize_compare('cac,c,h->h')
+        self.optimize_compare('cfc,c,h->h')
+        self.optimize_compare('cfc,c,d->d')
+        self.optimize_compare('c,cfc,d->d')
+        self.optimize_compare('d,c,cfc->d')
+        self.optimize_compare('d,bc,cfc->d')
+        self.optimize_compare('adb,bc,cfc->d')
+        self.optimize_compare('adb,bc,fa,cfc->d')
+        self.optimize_compare('ecb,fef,bad,ed->ac')
+        self.optimize_compare('fdf,cdd,ccd,afe->ae')
+        self.optimize_compare('adb,cfc->d')
+
+    @unittest.skipIf(True, "diagonal still not converted into ONNX")
+    def test_np_test_edge_cases_duplicate_indices(self):
+        self.optimize_compare('dd,fb,be,cdb->cef')
+        self.optimize_compare('dcc,fce,ea,dbf->ab')
+        self.optimize_compare('ed,fcd,ff,bcf->be')
+        self.optimize_compare('baa,dcf,af,cde->be')
+        self.optimize_compare('fff,fae,bef,def->abd')
+
+
+if __name__ == "__main__":
+    unittest.main()