internal code quality attempts 3

Author: refraction-ray

.pylintrc

@@ -56,24 +56,6 @@ max-module-lines=2000
 overgeneral-exceptions=BaseException,
                        Exception
 
-
-[BASIC]
-
-good-names=Bell_pair_block,
-           QAOA_block,
-           Grid2D_block,
-           dtypestr,
-           libjax,
-           jnp,
-           jsp,
-           torchlib,
-           N,
-           M,
-           U,
-           S,
-           V,
-           R,
-           Q,
-           A,
-           E,
-           P
+# Note how codecc doesn't accept goddnames in pylintrc, and how we use pylint disable invalid name per file instead
+# it is not neat at all, but this is codecc's badness :(
+# stupid to check variable name convention when you are a scientist dealing with lots of N, Pauli or QAOA

tensorcircuit/backends/ops.py

@@ -1,6 +1,7 @@
 """
 customized ops for ML framework
 """
+# pylint: disable=invalid-name
 
 from typing import Any, Tuple, Sequence
 

tensorcircuit/backends/pytorch_backend.py

@@ -5,6 +5,7 @@
 import logging
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 
+import tensornetwork
 from tensornetwork.backends.pytorch import pytorch_backend
 
 try:  # old version tn compatiblity
@@ -30,6 +31,20 @@
 # To be added once pytorch backend is ready
 
 
+def _sum_torch(
+    self: Any,
+    tensor: Tensor,
+    axis: Optional[Sequence[int]] = None,
+    keepdims: bool = False,
+) -> Tensor:
+    if axis is None:
+        axis = tuple([i for i in range(len(tensor.shape))])
+    return torchlib.sum(tensor, dim=axis, keepdim=keepdims)
+
+
+tensornetwork.backends.pytorch.pytorch_backend.PyTorchBackend.sum = _sum_torch
+
+
 class PyTorchBackend(pytorch_backend.PyTorchBackend):  # type: ignore
     def __init__(self) -> None:
         super(PyTorchBackend, self).__init__()
@@ -155,33 +170,39 @@ def cond(
         return false_fun()
 
     def switch(self, index: Tensor, branches: Sequence[Callable[[], Tensor]]) -> Tensor:
-        return branches[index]()
+        return branches[index.numpy()]()
 
     def grad(
         self, f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0
     ) -> Callable[..., Any]:
         def wrapper(*args: Any, **kws: Any) -> Any:
-            x = []
-            if isinstance(argnums, int):
-                argnumsl = [argnums]
-                # if you also call lhs as argnums, something weird may happen
-                # the reason is that python then take it as local vars
-            else:
-                argnumsl = argnums  # type: ignore
-            for i, arg in enumerate(args):
-                if i in argnumsl:
-                    x.append(arg.requires_grad_(True))
-                else:
-                    x.append(arg)
-            y = f(*x, **kws)
-            y.backward()
-            gs = [x[i].grad for i in argnumsl]
-            if len(gs) == 1:
-                gs = gs[0]
-            return gs
+            _, gr = self.value_and_grad(f, argnums)(*args, **kws)
+            return gr
 
         return wrapper
 
+        # def wrapper(*args: Any, **kws: Any) -> Any:
+        #     x = []
+        #     if isinstance(argnums, int):
+        #         argnumsl = [argnums]
+        #         # if you also call lhs as argnums, something weird may happen
+        #         # the reason is that python then take it as local vars
+        #     else:
+        #         argnumsl = argnums  # type: ignore
+        #     for i, arg in enumerate(args):
+        #         if i in argnumsl:
+        #             x.append(arg.requires_grad_(True))
+        #         else:
+        #             x.append(arg)
+        #     y = f(*x, **kws)
+        #     y.backward()
+        #     gs = [x[i].grad for i in argnumsl]
+        #     if len(gs) == 1:
+        #         gs = gs[0]
+        #     return gs
+
+        # return wrapper
+
     def value_and_grad(
         self, f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0
     ) -> Callable[..., Tuple[Any, Any]]:

tests/test_backends.py

@@ -1,3 +1,5 @@
+# pylint: disable=invalid-name
+
 import sys
 import os
 from functools import partial
@@ -54,6 +56,16 @@ def test_vmap_torch(torchb):
     assert r.numpy()[0, 0] == 3.0
 
 
+def test_grad_torch(torchb):
+    a = tc.backend.ones([2], dtype="float32")
+
+    @tc.backend.grad
+    def f(x):
+        return tc.backend.sum(x)
+
+    np.testing.assert_allclose(f(a), np.ones([2]), atol=1e-5)
+
+
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
 def test_backend_scatter(backend):
     assert np.allclose(

tests/test_circuit.py

@@ -573,8 +573,6 @@ def f(param, max_singular_values=None, max_truncation_err=None, fixed_choice=Non
     s3, g3 = f_vag(tc.backend.ones([4, n]), max_singular_values=2, fixed_choice=1)
 
     np.testing.assert_allclose(s1, s3, atol=1e-5)
-    print(g1[:, :])
-    print(g3[:, :])
     # DONE(@refraction-ray): nan on jax backend?
     # i see, complex value SVD is not supported on jax for now :)
     # I shall further customize complex SVD, finally it has applications
@@ -615,13 +613,9 @@ def f(param, max_singular_values=None, max_truncation_err=None, fixed_choice=Non
                 c.rx(i, theta=param[2 * j + 1, i])
         loss = c.expectation(
             (
-                tc.gates.z(),
+                tc.gates.x(),
                 [1],
             ),
-            (
-                tc.gates.z(),
-                [2],
-            ),
         )
         return tc.backend.real(loss)
 

tests/test_dmcircuit.py

@@ -109,55 +109,25 @@ def test_inputs_and_kraus():
 
 
 def test_gate_depolarizing():
-    c = tc.DMCircuit(2)
-    c.H(0)
-    c.apply_general_kraus(depolarizingchannel(0.1, 0.1, 0.1), [(1,)])
-    np.testing.assert_allclose(
-        c.densitymatrix(check=True),
-        np.array(
-            [[0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1], [0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1]]
-        ),
-        atol=1e-5,
+    ans = np.array(
+        [[0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1], [0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1]]
     )
 
-    c = tc.DMCircuit(2)
-    c.H(0)
-    kraus = depolarizingchannel(0.1, 0.1, 0.1)
-    # kraus = [k.tensor for k in kraus]
-    c.depolarizing(1, px=0.1, py=0.1, pz=0.1)
-    np.testing.assert_allclose(
-        c.densitymatrix(check=True),
-        np.array(
-            [[0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1], [0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1]]
-        ),
-        atol=1e-5,
-    )
-
-    c = tc.DMCircuit2(2)
-    c.H(0)
-    kraus = depolarizingchannel(0.1, 0.1, 0.1)
-    # kraus = [k.tensor for k in kraus]
-    c.apply_general_kraus(kraus, 1)
-    np.testing.assert_allclose(
-        c.densitymatrix(check=True),
-        np.array(
-            [[0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1], [0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1]]
-        ),
-        atol=1e-5,
-    )
+    def check_template(c, api="v1"):
+        c.H(0)
+        if api == "v1":
+            kraus = depolarizingchannel(0.1, 0.1, 0.1)
+            c.apply_general_kraus(kraus, [(1,)])
+        else:
+            c.depolarizing(1, px=0.1, py=0.1, pz=0.1)
+        np.testing.assert_allclose(
+            c.densitymatrix(check=True),
+            ans,
+            atol=1e-5,
+        )
 
-    c = tc.DMCircuit2(2)
-    c.H(0)
-    kraus = depolarizingchannel(0.1, 0.1, 0.1)
-    # kraus = [k.tensor for k in kraus]
-    c.depolarizing(1, px=0.1, py=0.1, pz=0.1)
-    np.testing.assert_allclose(
-        c.densitymatrix(check=True),
-        np.array(
-            [[0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1], [0.4, 0, 0.4, 0], [0, 0.1, 0, 0.1]]
-        ),
-        atol=1e-5,
-    )
+    for c, v in zip([tc.DMCircuit(2), tc.DMCircuit2(2)], ["v1", "v2"]):
+        check_template(c, v)
 
 
 @pytest.mark.parametrize("backend", [lf("jaxb"), lf("tfb")])

tests/test_mpscircuit.py

@@ -1,4 +1,6 @@
 # pylint: disable=unused-variable
+# pylint: disable=invalid-name
+
 import sys
 import os
 import numpy as np

tests/test_quantum.py

@@ -1,3 +1,5 @@
+# pylint: disable=invalid-name
+
 from functools import partial
 import os
 import sys