update nested_vmap example

Author: refraction-ray

examples/nested_vmap_grad.py

@@ -10,7 +10,7 @@
     with tc.runtime_backend(backend) as K:
         L = 2
         inputs = K.cast(K.ones([3, 2]), tc.rdtypestr)
-        weights = K.cast(K.ones([2]), tc.rdtypestr)
+        weights = K.cast(K.ones([3, 2]), tc.rdtypestr)
 
         def ansatz(thetas, alpha):
             c = tc.Circuit(L)
@@ -27,24 +27,42 @@ def f(thetas, alpha):
             observables = K.stack([K.real(c.expectation_ps(z=[i])) for i in range(L)])
             return K.mean(observables)
 
-        # f_vmap = K.vmap(f, vectorized_argnums=0)
-
-        print("grad", K.grad(f)(inputs[0], weights))
-        print("vmap", K.vmap(f)(inputs, weights))
-        print("vmap over grad", K.vmap(K.grad(f))(inputs, weights))
+        print("grad_0", K.grad(f)(inputs[0], weights[0]))
+        print("grad_1", K.grad(f, argnums=1)(inputs[0], weights[0]))
+        print("vmap_0", K.vmap(f)(inputs, weights[0]))
+        print("vmap_1", K.vmap(f, vectorized_argnums=1)(inputs[0], weights))
+        print("vmap over grad_0", K.vmap(K.grad(f))(inputs, weights[0]))
         # wrong in tf due to https://github.com/google/TensorNetwork/issues/940
         # https://github.com/tensorflow/tensorflow/issues/52148
-        print("vmap over jacfwd", K.vmap(K.jacfwd(f))(inputs, weights))
-        print("jacfwd over vmap", K.jacfwd(K.vmap(f))(inputs, weights))
-        r = K.vmap(K.jacrev(f))(inputs, weights)
-        print("vmap over jacrev", r)
+        print("vmap over grad_1", K.vmap(K.grad(f, argnums=1))(inputs, weights[0]))
+        # wrong in tf
+        print("vmap over jacfwd_0", K.vmap(K.jacfwd(f))(inputs, weights[0]))
+        print("jacfwd_0 over vmap", K.jacfwd(K.vmap(f))(inputs, weights[0]))
+        print("vmap over jacfwd_1", K.vmap(K.jacfwd(f, argnums=1))(inputs, weights[0]))
+        print("jacfwd_1 over vmap", K.jacfwd(K.vmap(f), argnums=1)(inputs, weights[0]))
+        r = K.vmap(K.jacrev(f))(inputs, weights[0])
+        print("vmap over jacrev0", r)
+        # wrong in tf
+        r = K.jacrev(K.vmap(f))(inputs, weights[0])
+        print("jacrev0 over vmap", r)
+        r = K.vmap(K.jacrev(f, argnums=1))(inputs, weights[0])
+        print("vmap over jacrev1", r)
+        # wrong in tf
+        r = K.jacrev(K.vmap(f), argnums=1)(inputs, weights[0])
+        print("jacrev1 over vmap", r)
+        r = K.vmap(K.jacrev(f, argnums=1), vectorized_argnums=1)(inputs[0], weights)
+        print("vmap1 over jacrev1", r)
+        r = K.jacrev(K.vmap(f, vectorized_argnums=1), argnums=1)(inputs[0], weights)
+        print("jacrev1 over vmap1", r)
+        r = K.vmap(K.hessian(f))(inputs, weights[0])
+        print("vmap over hess0", r)
         # wrong in tf
-        r = K.jacrev(K.vmap(f))(inputs, weights)
-        print("jacrev over vmap", r)
-        r = K.vmap(K.hessian(f))(inputs, weights)
-        print("vmap over hess", r)
+        r = K.hessian(K.vmap(f))(inputs, weights[0])
+        print("hess0 over vmap", r)
+        r = K.vmap(K.hessian(f, argnums=1))(inputs, weights[0])
+        print("vmap over hess1", r)
         # wrong in tf
-        r = K.hessian(K.vmap(f))(inputs, weights)
-        print("hess over vmap", r)
+        r = K.hessian(K.vmap(f), argnums=1)(inputs, weights[0])
+        print("hess1 over vmap", r)
 
 # lessons: never put vmap outside gradient function in tf