add mat prod vmap example

Author: refraction-ray

CHANGELOG.md

@@ -6,6 +6,8 @@
 
 - Add multiple GPU VQE examples using jax pmap
 
+- Add benchmark example showcasing new way of implementing matrix product using vmap
+
 ## 0.10.0
 
 ### Added

examples/matprod_vmap.py

@@ -0,0 +1,41 @@
+"""
+matrix product: a new twist
+rewrite matrix product in a vmap style
+"""
+from functools import partial
+
+import numpy as np
+import tensorcircuit as tc
+
+for bk in ["jax", "tensorflow"]:
+    with tc.runtime_backend(bk) as K:
+        print("~~~~~~~~~~~~~~~~~~~~~")
+        print(f"using {K.name} backend")
+
+        @partial(K.jit, jit_compile=True)
+        def mul(a, b):
+            return a @ b
+
+        def ij(i, j):
+            """
+            Inner product
+            """
+            return K.tensordot(i, j, 1)
+
+        vij = K.vmap(ij, vectorized_argnums=1)
+        vvij = K.vmap(vij, vectorized_argnums=0)
+
+        @partial(K.jit, jit_compile=True)
+        def mul2(a, b):
+            b = K.transpose(b)
+            return vvij(a, b)
+
+        for shape in [(256, 4096), (4096, 256), (2048, 2048)]:
+            print(shape)
+            a = K.implicit_randn(shape)
+            b = K.implicit_randn([shape[1], shape[0]])
+            print("plain matprod")
+            r1, _, _ = tc.utils.benchmark(mul, a, b, tries=10)
+            print("vmap matprod")
+            r2, _, _ = tc.utils.benchmark(mul2, a, b, tries=10)
+            np.testing.assert_allclose(r1, r2, atol=1e-5)