def _kak_equivalent_vectors(kak_vec) -> np.ndarray:
    """Generates all KAK vectors equivalent under single qubit unitaries."""

    # Technically this is not all equivalent vectors, but a subset of vectors
    # which are not guaranteed to give the same answer under the infidelity
    # formula.

    kak_vec = np.asarray(kak_vec, dtype=float)

    # Apply all permutations, then all negations, then all shifts.

    out = np.einsum('pab,...b->...pa', _perms_123, kak_vec)  # (...,6,3)
    out = np.einsum('nab,...b->...na', _negations, out)  # (...,6,4,3)

    # (...,8,6,4,3)
    out = out[..., np.newaxis, :, :, :] + _offsets[:, np.newaxis, np.newaxis, :]

    # Merge indices
    return np.reshape(out, out.shape[:-4] + (192, 3))
Generates all KAK vectors equivalent under single qubit unitaries.
cirq-core/cirq/transformers/heuristic_decompositions/gate_tabulation_math_utils.py