Refactor einsum files

.gitignore

@@ -313,3 +313,4 @@ _unittests/ut_tools/**/*.npz
 _unittests/ut_tools/**/*.pb
 _unittests/ut_onnxrt/onnxruntime_profile*.json
 _doc/notebooks/onnxruntime_profile*.json
+_doc/sphinxdoc/source/phdoc_static/embed*.js

_doc/sphinxdoc/source/api/index.rst

@@ -30,6 +30,7 @@ This is a summary of functions this modules provides.
  asv
  validation
  tools
+ onnx_tools
 
 **Outside ONNX world**
 

_doc/sphinxdoc/source/api/onnx_tools.rst

@@ -0,0 +1,36 @@
+
+Tools to play with ONNX graphs
+==============================
+
+.. contents::
+ :local:
+
+Accessor
+++++++++
+
+.. autosignature:: mlprodict.onnx_tools.onnx_tools.find_node_input_name
+
+.. autosignature:: mlprodict.onnx_tools.onnx_tools.find_node_name
+
+.. autosignature:: mlprodict.onnx_tools.onnx_tools.insert_node
+
+Optimisation
+++++++++++++
+
+The following functions reduce the number of ONNX operators in a graph
+while keeping the same results. The optimized graph
+is left unchanged.
+
+.. autosignature:: mlprodict.onnx_tools.optim.onnx_optimisation.onnx_remove_node
+
+.. autosignature:: mlprodict.onnx_tools.optim.onnx_optimisation_identity.onnx_remove_node_identity
+
+.. autosignature:: mlprodict.onnx_tools.optim.onnx_optimisation_redundant.onnx_remove_node_redundant
+
+.. autosignature:: mlprodict.onnx_tools.optim.onnx_remove_unused.onnx_remove_node_unused
+
+Validation
+++++++++++
+
+.. autosignature:: mlprodict.onnx_tools.model_checker.onnx_shaker
+

_doc/sphinxdoc/source/api/onnxrt.rst

@@ -46,8 +46,6 @@ Validation
 
 .. autosignature:: mlprodict.onnxrt.validate.validate_summary.summary_report
 
-.. autosignature:: mlprodict.onnx_tools.model_checker.onnx_shaker
-
 .. autosignature:: mlprodict.onnxrt.validate.validate_graph.plot_validate_benchmark
 
 C++ classes
@@ -115,21 +113,6 @@ C++ classes
 
 .. autosignature:: mlprodict.onnxrt.ops_cpu._op_onnx_numpy.topk_element_fetch_int64
 
-Optimisation
-++++++++++++
-
-The following functions reduce the number of ONNX operators in a graph
-while keeping the same results. The optimized graph
-is left unchanged.
-
-.. autosignature:: mlprodict.onnx_tools.optim.onnx_optimisation.onnx_remove_node
-
-.. autosignature:: mlprodict.onnx_tools.optim.onnx_optimisation_identity.onnx_remove_node_identity
-
-.. autosignature:: mlprodict.onnx_tools.optim.onnx_optimisation_redundant.onnx_remove_node_redundant
-
-.. autosignature:: mlprodict.onnx_tools.optim.onnx_remove_unused.onnx_remove_node_unused
-
 Shapes
 ++++++
 

_doc/sphinxdoc/source/api/testing.rst

@@ -18,27 +18,31 @@ Experimental implementations for algorithm.
 Einsum
 ^^^^^^
 
-.. autosignature:: mlprodict.testing.einsum_impl.analyse_einsum_equation
+.. autosignature:: mlprodict.testing.einsum.einsum_fct.einsum
 
-.. autosignature:: mlprodict.testing.einsum_impl.apply_sequence
+.. autosignature:: mlprodict.testing.einsum.einsum_fct.CachedEinsum
 
-.. autosignature:: mlprodict.testing.einsum_impl.decompose_einsum_equation
+.. autosignature:: mlprodict.testing.einsum.einsum_impl.analyse_einsum_equation
+
+.. autosignature:: mlprodict.testing.einsum.einsum_impl.apply_sequence
+
+.. autosignature:: mlprodict.testing.einsum.einsum_impl.decompose_einsum_equation
 
 .. autosignature:: mlprodict.testing.experimental_c.custom_einsum_float
 
 .. autosignature:: mlprodict.testing.experimental_c.custom_einsum_double
 
-.. autosignature:: mlprodict.testing.einsum_bench.einsum_benchmark
+.. autosignature:: mlprodict.testing.einsum.einsum_bench.einsum_benchmark
 
-.. autosignature:: mlprodict.testing.einsum_impl_ext.numpy_diagonal
+.. autosignature:: mlprodict.testing.einsum.einsum_impl_ext.numpy_diagonal
 
-.. autosignature:: mlprodict.testing.einsum_impl_ext.numpy_extended_dot
+.. autosignature:: mlprodict.testing.einsum.einsum_impl_ext.numpy_extended_dot
 
-.. autosignature:: mlprodict.testing.einsum_impl_ext.numpy_extended_dot_python
+.. autosignature:: mlprodict.testing.einsum.einsum_impl_ext.numpy_extended_dot_python
 
-.. autosignature:: mlprodict.testing.einsum_impl_ext.numpy_extended_dot_matrix
+.. autosignature:: mlprodict.testing.einsum.einsum_impl_ext.numpy_extended_dot_matrix
 
-.. autosignature:: mlprodict.testing.einsum_impl_ext.numpy_extended_dot_ouput_shape
+.. autosignature:: mlprodict.testing.einsum.einsum_impl_ext.numpy_extended_dot_ouput_shape
 
 Pad
 ^^^

_doc/sphinxdoc/source/blog/2021/2021-05-05_numpyapionnx1.rst

@@ -67,6 +67,7 @@
 
  .. runpython::
  :showcode:
+ :process:
 
  from typing import Any
  import numpy as np
@@ -103,6 +104,7 @@
 
  .. runpython::
  :showcode:
+ :process:
 
  from typing import Any
  import numpy as np

_doc/sphinxdoc/source/blog/2021/2021-05-05_numpyapionnx2.rst

@@ -28,6 +28,7 @@
 
  .. runpython::
  :showcode:
+ :process:
 
  import numpy
  from pandas import DataFrame

_unittests/ut_testing/test_blas_lapack.py

@@ -5,7 +5,7 @@
 import numpy
 from scipy.linalg.blas import sgemm # pylint: disable=E0611
 from pyquickhelper.pycode import ExtTestCase
-from mlprodict.testing.blas_lapack import gemm_dot, pygemm
+from mlprodict.testing.einsum.blas_lapack import gemm_dot, pygemm
 
 
 class TestBlasLapack(ExtTestCase):

_unittests/ut_testing/test_einsum.py

@@ -10,9 +10,9 @@
 from onnxruntime import (
  InferenceSession, GraphOptimizationLevel, SessionOptions)
 from pyquickhelper.pycode import ExtTestCase
-from mlprodict.testing.einsum_impl_ext import (
+from mlprodict.testing.einsum.einsum_impl_ext import (
  numpy_diagonal, numpy_extended_dot, numpy_extended_dot_python)
-from mlprodict.testing.einsum_impl import (
+from mlprodict.testing.einsum.einsum_impl import (
  analyse_einsum_equation, decompose_einsum_equation, EinsumSubOp,
  apply_einsum_sequence)
 from mlprodict.onnxrt import OnnxInference

_unittests/ut_testing/test_einsum_benchmark.py

@@ -3,7 +3,7 @@
 """
 import unittest
 from pyquickhelper.pycode import ExtTestCase
-from mlprodict.testing.einsum_bench import einsum_benchmark
+from mlprodict.testing.einsum.einsum_bench import einsum_benchmark
 
 
 class TestEinsumBenchmark(ExtTestCase):

_unittests/ut_testing/test_einsum_einsum.py

@@ -0,0 +1,54 @@
+"""
+@brief test log(time=8s)
+"""
+import unittest
+import numpy
+from pyquickhelper.pycode import ExtTestCase
+from mlprodict.testing.einsum import einsum
+from mlprodict.testing.einsum.einsum_fct import enumerate_cached_einsum
+from mlprodict.tools.asv_options_helper import get_opset_number_from_onnx
+
+
+class TestEinsumEinsum(ExtTestCase):
+
+ def common_test(self, equation, runtime=None, opset=None, N=5,
+ optimize=False):
+ if opset is None:
+ opset = get_opset_number_from_onnx()
+ inps = equation.split('->')[0].split(',')
+ lens = [len(s) for s in inps]
+ inputs = [numpy.random.randn(N ** d).reshape((N,) * d)
+ for d in lens]
+ if runtime is None:
+ runtime = ['batch_dot', 'python', 'onnxruntime1']
+ elif isinstance(runtime, str):
+ runtime = [runtime]
+ for rt in runtime:
+ for dtype in [numpy.float32, numpy.float64]:
+ decimal = 5 if dtype == numpy.float32 else 8
+ with self.subTest(dt=dtype, rt=rt,
+ eq=equation, opset=opset,
+ opt=optimize):
+ typed = [i.astype(dtype) for i in inputs]
+ kwargs = dict(runtime=rt, opset=opset, optimize=optimize)
+ if __name__ == "__main__":
+ kwargs["verbose"] = 1
+ exp = numpy.einsum(equation, *typed)
+ got = einsum(equation, *typed, **kwargs)
+ self.assertEqualArray(exp, got, decimal=decimal)
+ got = einsum(equation, *typed, **kwargs)
+ self.assertEqualArray(exp, got, decimal=decimal)
+
+ def test_einsum(self):
+ self.common_test("abc,cd->abd")
+ self.common_test("abc,cd,de->abe")
+ res = list(enumerate_cached_einsum())
+ self.assertGreater(len(res), 2)
+ self.assertIn('CachedEinsum', str(res))
+
+ def test_einsum_optimize(self):
+ self.common_test("abc,cd->abd", optimize=True)
+
+
+if __name__ == "__main__":
+ unittest.main()

_unittests/ut_testing/test_einsum_generic_dot.py

@@ -6,7 +6,7 @@
 from contextlib import redirect_stdout
 import numpy
 from pyquickhelper.pycode import ExtTestCase
-from mlprodict.testing.einsum_impl_ext import (
+from mlprodict.testing.einsum.einsum_impl_ext import (
  numpy_extended_dot, numpy_extended_dot_python,
  numpy_extended_dot_matrix, numpy_diagonal,
  _numpy_extended_dot_equation,

mlprodict/cli/einsum.py

@@ -37,7 +37,7 @@ def einsum_test(equation="abc,cd->abd", shape="30", perm=False,
 
  python -m mlprodict einsum_test --equation="abc,cd->abd" --output=res.csv
  """
- from ..testing.einsum_bench import einsum_benchmark # pylint: disable=E0402
+ from ..testing.einsum.einsum_bench import einsum_benchmark # pylint: disable=E0402
 
  perm = perm in ('True', '1', 1, True)
  if "(" not in shape:

mlprodict/testing/einsum/__init__.py

@@ -0,0 +1,14 @@
+"""
+@file
+@brief Shortcut to *testing.einsum*.
+"""
+
+from .einsum_bench import einsum_benchmark
+from .einsum_fct import einsum
+from .einsum_impl import decompose_einsum_equation, apply_einsum_sequence
+from .einsum_impl_classes import EinsumSubOp, GraphEinsumSubOp
+from .einsum_impl_ext import (
+ numpy_extended_dot,
+ numpy_extended_dot_python,
+ numpy_extended_dot_matrix,
+ numpy_diagonal)

mlprodict/testing/einsum_bench.py → mlprodict/testing/einsum/einsum_bench.py

@@ -6,8 +6,8 @@
 import numpy
 from onnx import helper, TensorProto
 from onnxruntime import InferenceSession
-from ..onnxrt import OnnxInference
-from .bench_helper import measure_time
+from ...onnxrt import OnnxInference
+from ..bench_helper import measure_time
 from .einsum_impl import decompose_einsum_equation, apply_einsum_sequence