diff --git a/lib.py b/lib.py
index a9d5e65..8409c66 100644
--- a/lib.py
+++ b/lib.py
@@ -643,7 +643,7 @@ class UnitaryOperator(LinearTransformation, UnitaryMatrix):
             assert 1 <= len(which_qbit) < total_qbits
             assert len(which_qbit) == len(self.partials)
             assert all([q < total_qbits for q in which_qbit])
-            this_gate_len = 2**(len(self.partials) - 1)
+            this_gate_len = 2 ** (len(self.partials) - 1)
             bins = generate_bins(this_gate_len)
             extended_m, next_partial = [[] for _ in bins], 0
             for qbit in range(total_qbits):
@@ -656,7 +656,8 @@ class UnitaryOperator(LinearTransformation, UnitaryMatrix):
                     if this_partial == C_partial:
                         for i in range(this_gate_len):
                             bin_dig = bins[i][next_partial]
-                            extended_m[i].append(s("|" + bin_dig + "><" + bin_dig + "|"))
+                            extended_m[i].append(
+                                s("|" + bin_dig + "><" + bin_dig + "|"))
                     else:
                         for i in range(this_gate_len - 1):
                             extended_m[i].append(I)
@@ -664,9 +665,11 @@ class UnitaryOperator(LinearTransformation, UnitaryMatrix):
                     next_partial += 1
             new_m = sum([np.prod(e).m for e in extended_m])
         else:
-            raise Exception("which_qubit needs to be either an int of N-th qubit or list")
+            raise Exception(
+                "which_qubit needs to be either an int of N-th qubit or list")
 
-        extended_op = UnitaryOperator(new_m, name=self.name, partials=self.partials)
+        extended_op = UnitaryOperator(new_m, name=self.name,
+                                      partials=self.partials)
         return State(np.dot(extended_op.m, other.m))
 
     def __repr__(self):
@@ -854,6 +857,7 @@ T = lambda phi: Gate([[1, 0],
 
 C_partial = Gate(I.m, name="C")
 x_partial = Gate(X.m, name=REPR_TARGET)
+z_partial = Gate(Z.m, name=REPR_TARGET)
 
 CNOT = Gate([
     [1, 0, 0, 0],
@@ -862,6 +866,22 @@ CNOT = Gate([
     [0, 0, 1, 0],
 ], name="CNOT", partials=[C_partial, x_partial])
 
+CZ = Gate([
+    [1, 0, 0, 0],
+    [0, 1, 0, 0],
+    [0, 0, 1, 0],
+    [0, 0, 0, -1],
+], name="CZ", partials=[C_partial, z_partial])
+
+# TODO: These are not the correct partials
+SWAP = Gate([
+    [1, 0, 0, 0],
+    [0, 0, 1, 0],
+    [0, 1, 0, 0],
+    [0, 0, 0, 1],
+], name="SWAP", partials=[C_partial, C_partial])
+
+
 TOFF = Gate([
     [1, 0, 0, 0, 0, 0, 0, 0],
     [0, 1, 0, 0, 0, 0, 0, 0],
@@ -1017,6 +1037,12 @@ def test_partials():
     # apply on 0, 3 of 4qbit state
     assert CNOT.on(s("|1000>"), which_qbit=[0, 3]) == s("|1001>")
 
+    # test SWAP gate
+    assert SWAP.on(s("|10>")) == s("|01>")
+    assert SWAP.on(s("|01>")) == s("|10>")
+    # TODO:
+    # assert SWAP.on(s("|001>"), which_qbit=[0, 2]) == s("|100>")
+
     # test Toffoli gate
     assert TOFF.on(s("|000>")) == s("|000>")
     assert TOFF.on(s("|100>")) == s("|100>")