diff --git a/lib.py b/lib.py
index 7fb1243..a9d5e65 100644
--- a/lib.py
+++ b/lib.py
@@ -643,20 +643,24 @@ 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])
-            extended_m, next_partial = [[], []], 0
+            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):
                 if qbit not in which_qbit:
-                    extended_m[0].append(I)
-                    extended_m[1].append(I)
+                    for i in range(this_gate_len):
+                        extended_m[i].append(I)
                 else:
                     this_partial = self.partials[next_partial]
                     # TODO: This works only with C_partial :(((
                     if this_partial == C_partial:
-                        extended_m[0].append(s("|0><0|"))
-                        extended_m[1].append(s("|1><1|"))
+                        for i in range(this_gate_len):
+                            bin_dig = bins[i][next_partial]
+                            extended_m[i].append(s("|" + bin_dig + "><" + bin_dig + "|"))
                     else:
-                        extended_m[0].append(I)
-                        extended_m[1].append(this_partial)
+                        for i in range(this_gate_len - 1):
+                            extended_m[i].append(I)
+                        extended_m[-1].append(this_partial)
                     next_partial += 1
             new_m = sum([np.prod(e).m for e in extended_m])
         else:
@@ -842,6 +846,12 @@ T = lambda phi: Gate([[1, 0],
 # np.kron(np.outer(_0.m, _0.m), np.eye(2)) + np.kron(np.outer(_1.m, _1.m), X.m)
 # _0.x(_0) * Matrix(I.m) + _1.x(_1) * Matrix(X.m)
 
+# Additionally for Toffoli: on a 4qubit system with C-CX (qbits 0,2,3)
+# s("|0><0|") * I * s("|0><0|") * I \
+# + s("|0><0|") * I * s("|1><1|") * I \
+# + s("|1><1|") * I * s("|0><0|") * I \
+# + s("|1><1|") * I * s("|1><1|") * X
+
 C_partial = Gate(I.m, name="C")
 x_partial = Gate(X.m, name=REPR_TARGET)
 
@@ -988,17 +998,22 @@ def test_partials():
                   "Operating dim-4 operator on a dim-8 state. Please specify "
                   "which_qubit to operate on",
                   CNOT.on, s("|100>"))
+
     # apply on 0, 1 of 3qbit state
     assert CNOT.on(s("|000>"), which_qbit=[0, 1]) == s("|000>")
     assert CNOT.on(s("|100>"), which_qbit=[0, 1]) == s("|110>")
+
     # apply on 1, 2 of 3qbit state
     assert CNOT.on(s("|000>"), which_qbit=[1, 2]) == s("|000>")
     assert CNOT.on(s("|010>"), which_qbit=[1, 2]) == s("|011>")
+
     # apply on 0, 2 of 3qbit state
     assert CNOT.on(s("|000>"), which_qbit=[0, 2]) == s("|000>")
     assert CNOT.on(s("|100>"), which_qbit=[0, 2]) == s("|101>")
+
     # apply on 0, 2 of 4qbit state
     assert CNOT.on(s("|1000>"), which_qbit=[0, 2]) == s("|1010>")
+
     # apply on 0, 3 of 4qbit state
     assert CNOT.on(s("|1000>"), which_qbit=[0, 3]) == s("|1001>")
 
@@ -1008,8 +1023,17 @@ def test_partials():
     assert TOFF.on(s("|110>")) == s("|111>")
     assert TOFF.on(s("|111>")) == s("|110>")
 
-    # TODO: This doesn't work :(((
-    # assert TOFF.on(s("|1100>"), which_qbit=[0, 1, 3]) == s("|1101>")
+    # controls on 0, 1 target on 3 for 4
+    assert TOFF.on(s("|1100>"), which_qbit=[0, 1, 3]) == s("|1101>")
+    # controls on 0, 2 target on 3 for 4
+    assert TOFF.on(s("|1010>"), which_qbit=[0, 2, 3]) == s("|1011>")
+    assert TOFF.on(s("|1011>"), which_qbit=[0, 2, 3]) == s("|1010>")
+    assert TOFF.on(s("|1111>"), which_qbit=[0, 2, 3]) == s("|1110>")
+    # controls on 0, 2 target on 4 for 5
+    assert TOFF.on(s("|10101>"), which_qbit=[0, 2, 4]) == s("|10100>")
+    assert TOFF.on(s("|10111>"), which_qbit=[0, 2, 4]) == s("|10110>")
+    assert TOFF.on(s("|11111>"), which_qbit=[0, 2, 4]) == s("|11110>")
+    assert TOFF.on(s("|10100>"), which_qbit=[0, 2, 4]) == s("|10101>")
 
 
 def test():