Optimize targeted_left_multiply for onehot source

cirq-core/cirq/linalg/transformations.py

@@ -14,7 +14,7 @@
 
 """Utility methods for transforming matrices or vectors."""
 
-from typing import Tuple, Optional, Sequence, List, Union
+from typing import Any, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -141,6 +141,26 @@ def targeted_left_multiply(
 
     k = len(target_axes)
     d = len(right_target.shape)
+    nonzeros = np.flatnonzero(left_matrix)
+    if len(nonzeros) == 1:
+        # This is just moving a slice from one place to another with everything else zeros, and an
+        # optional rescale. More efficient to operate on slices directly than do a full einsum.
+        index = np.unravel_index(nonzeros[0], left_matrix.shape)
+        if out is None:
+            out = np.zeros_like(right_target)
+        else:
+            out[...] = 0
+        source_slices: List[Any] = [slice(None)] * d
+        target_slices: List[Any] = [slice(None)] * d
+        for i in range(k):
+            source_slices[target_axes[i]] = index[k + i]
+            target_slices[target_axes[i]] = index[i]
+        sleis = right_target[tuple(source_slices)]
+        if left_matrix[index] != 1:
+            sleis *= left_matrix[index]
+        out[tuple(target_slices)] = sleis
+        return out
+
     work_indices = tuple(range(k))
     data_indices = tuple(range(k, k + d))
     used_data_indices = tuple(data_indices[q] for q in target_axes)

cirq-core/cirq/linalg/transformations_test.py

@@ -192,7 +192,7 @@ def test_targeted_conjugate_simple():
         np.array([1, 2, 3, 4,
                   5, 6, 7, 8,
                   9, 10, 11, 12,
-                  13, 14, 15, 16]),
+                  13, 14, 15, 16], dtype=complex),
         (2,) * 4
     )
     expected = np.reshape(