def two_qubit_gate_product_tabulation(
    base_gate: np.ndarray,
    max_infidelity: float,
    *,
    sample_scaling: int = 50,
    allow_missed_points: bool = True,
    random_state: cirq.RANDOM_STATE_OR_SEED_LIKE = None,
) -> TwoQubitGateTabulation:
    r"""Generate a TwoQubitGateTabulation for a base two qubit unitary.

    Args:
        base_gate: The base gate of the tabulation.
        max_infidelity: Sets the desired density of tabulated product unitaries.
            The typical nearest neighbor Euclidean spacing (of the KAK vectors)
            will be on the order of $\sqrt{max\_infidelity}$. Thus the number of
            tabulated points will scale as $max\_infidelity^{-3/2}$.
        sample_scaling: Relative number of random gate products to use in the
            tabulation. The total number of random local unitaries scales as
            ~ $max\_infidelity^{-3/2} * sample\_scaling$. Must be positive.
        random_state: Random state or random state seed.
        allow_missed_points: If True, the tabulation is allowed to conclude
            even if not all points in the Weyl chamber are expected to be
            compilable using 2 or 3 base gates. Otherwise, an error is raised
            in this case.

    Returns:
        A TwoQubitGateTabulation object used to compile new two-qubit gates from
        products of the base gate with 1-local unitaries.

    Raises:
        ValueError: If `allow_missing_points` is False and not all the points
            in the Weyl chamber are compilable using 2 or 3 base gates.
    """
    rng = value.parse_random_state(random_state)

    assert 1 / 2 > max_infidelity > 0
    spacing = np.sqrt(max_infidelity / 3)
    mesh_points = weyl_chamber_mesh(spacing)

    # Number of random gate products to sample over in constructing the
    # tabulation. This has to be at least the number of mesh points, as
    # a single product can only be associated with one mesh point.
    assert sample_scaling > 0, 'Input sample_scaling must positive.'
    num_mesh_points = mesh_points.shape[0]
    num_samples = num_mesh_points * sample_scaling

    # include the base gate itself
    kak_vecs = [cirq.kak_vector(base_gate, check_preconditions=False)]
    sq_cycles: list[tuple[_SingleQubitGatePair, ...]] = [()]

    # Tabulate gates that are close to gates in the mesh
    u_locals_0 = random_qubit_unitary((num_samples,), rng=rng)
    u_locals_1 = random_qubit_unitary((num_samples,), rng=rng)

    tabulated_kak_inds = np.zeros((num_mesh_points,), dtype=bool)

    tabulation_cutoff = 0.5 * spacing
    out = _tabulate_kak_vectors(
        already_tabulated=tabulated_kak_inds,
        base_gate=base_gate,
        max_dist=tabulation_cutoff,
        kak_mesh=mesh_points,
        local_unitary_pairs=[(u_locals_0, u_locals_1)],
    )
    kak_vecs.extend(out.kept_kaks)
    sq_cycles.extend(out.kept_cycles)

    # Will be used later for getting missing KAK vectors.
    kak_vecs_single = np.array(kak_vecs)
    sq_cycles_single = list(sq_cycles)

    summary = (
        f'Fraction of Weyl chamber reached with 2 gates'
        f': {tabulated_kak_inds.sum() / num_mesh_points :.3f}'
    )

    # repeat for double products
    # Multiply by the same local unitary in the gate product
    out = _tabulate_kak_vectors(
        already_tabulated=tabulated_kak_inds,
        base_gate=base_gate,
        max_dist=tabulation_cutoff,
        kak_mesh=mesh_points,
        local_unitary_pairs=[(u_locals_0, u_locals_1)] * 2,
    )

    kak_vecs.extend(out.kept_kaks)
    sq_cycles.extend(out.kept_cycles)

    summary += (
        f'\nFraction of Weyl chamber reached with 2 gates and 3 gates'
        f'(same single qubit): '
        f'{tabulated_kak_inds.sum() / num_mesh_points :.3f}'
    )

    # If all KAK vectors in the mesh have been tabulated, return.
    missing_vec_inds = np.logical_not(tabulated_kak_inds).nonzero()[0]

    if not np.any(missing_vec_inds):  # pragma: no cover
        return TwoQubitGateTabulation(
            base_gate, np.array(kak_vecs), sq_cycles, max_infidelity, summary, ()
        )

    # Run through remaining KAK vectors that don't have products and try to
    # correct them

    u_locals_0p = random_qubit_unitary((100,), rng=rng)
    u_locals_1p = random_qubit_unitary((100,), rng=rng)
    u_locals = vector_kron(u_locals_0p, u_locals_1p)

    # Loop through the mesh points that have not yet been tabulated.
    # Consider their nonlocal parts A and compute products of the form
    # base_gate^\dagger k A
    # Compare the KAK vector of any of those products to the already tabulated
    # KAK vectors from single products of the form
    # base_gate k0 base_gate.
    # If they are close, then base_gate^\dagger k A  ~ base_gate k0 base_gate
    # So we may compute the outer local unitaries kL, kR such that
    #    base_gate^\dagger k A = kL base_gate k0 base_gate kR
    #    A = k^\dagger base_gate kL base_gate k0 base_gate kR
    #    the single-qubit unitary kL is the one we need to get the desired
    #    KAK vector.
    missed_points = []
    base_gate_dag = base_gate.conj().T
    for ind in missing_vec_inds:
        missing_vec = mesh_points[ind]
        # Unitary A we wish to solve for
        missing_unitary = kak_vector_to_unitary(missing_vec)

        # Products of the from base_gate^\dagger k A
        products = np.einsum('ab,...bc,cd', base_gate_dag, u_locals, missing_unitary)
        # KAK vectors for these products
        kaks = cirq.kak_vector(products, check_preconditions=False)
        kaks = kaks[..., np.newaxis, :]

        # Check if any of the product KAK vectors are close to a previously
        # tabulated KAK vector
        dists2 = np.sum((kaks - kak_vecs_single) ** 2, axis=-1)
        min_dist_inds = np.unravel_index(dists2.argmin(), dists2.shape)
        min_dist = np.sqrt(dists2[min_dist_inds])
        if min_dist < tabulation_cutoff:
            # If so, compute the single qubit unitary k_L such that
            # base_gate^\dagger k A = kL base_gate k0 base_gate kR
            # where k0 is the old (previously tabulated) single qubit unitary
            # and k is one of the single qubit unitaries used above.
            # Indices below are for k, k0 respectively
            new_ind, old_ind = min_dist_inds

            # Special case where the RHS is just base_gate (no single qubit
            # gates yet applied). I.e. base_gate^\dagger k A ~  base_gate
            # which implies  base_gate^\dagger k A = k_L base_gate k_R
            new_product = products[new_ind]
            if old_ind == 0:
                assert not sq_cycles_single[old_ind]
                base_product = base_gate
                _, kL, actual = _outer_locals_for_unitary(new_product, base_product)
                # Add to the enumeration
                sq_cycles.append((kL,))
            else:  # typical case mentioned above
                assert len(sq_cycles_single[old_ind]) == 1
                old_sq_cycle = sq_cycles_single[old_ind][0]
                old_k = np.kron(*old_sq_cycle)
                base_product = base_gate @ old_k @ base_gate
                _, kL, actual = _outer_locals_for_unitary(new_product, base_product)
                # Add to the enumeration
                sq_cycles.append((old_sq_cycle, kL))

            kak_vecs.append(cirq.kak_vector(base_gate @ actual, check_preconditions=False))
        elif not allow_missed_points:
            raise ValueError(f'Failed to tabulate a KAK vector near {missing_vec}')
        else:
            missed_points.append(missing_vec)

    kak_vecs_arr = np.array(kak_vecs)
    summary += (
        f'\nFraction of Weyl chamber reached with 2 gates and 3 gates '
        f'(after patchup)'
        f': {(len(kak_vecs_arr) - 1) / num_mesh_points :.3f}'
    )

    return TwoQubitGateTabulation(
        base_gate, kak_vecs_arr, sq_cycles, max_infidelity, summary, tuple(missed_points)
    )
Generate a TwoQubitGateTabulation for a base two qubit unitary.

    Args:
        base_gate: The base gate of the tabulation.
        max_infidelity: Sets the desired density of tabulated product unitaries.
            The typical nearest neighbor Euclidean spacing (of the KAK vectors)
            will be on the order of $\sqrt{max\_infidelity}$. Thus the number of
            tabulated points will scale as $max\_infidelity^{-3/2}$.
        sample_scaling: Relative number of random gate products to use in the
            tabulation. The total number of random local unitaries scales as
            ~ $max\_infidelity^{-3/2} * sample\_scaling$. Must be positive.
        random_state: Random state or random state seed.
        allow_missed_points: If True, the tabulation is allowed to conclude
            even if not all points in the Weyl chamber are expected to be
            compilable using 2 or 3 base gates. Otherwise, an error is raised
            in this case.

    Returns:
        A TwoQubitGateTabulation object used to compile new two-qubit gates from
        products of the base gate with 1-local unitaries.

    Raises:
        ValueError: If `allow_missing_points` is False and not all the points
            in the Weyl chamber are compilable using 2 or 3 base gates.
    
cirq-core/cirq/transformers/heuristic_decompositions/two_qubit_gate_tabulation.py