Function tensor_batch

ot/batch/_quadratic.py:17–152

r""" Compute the Gromov-Wasserstein cost tensor for a batch of problems. The Gromov-Wasserstein distance can be expressed as: .. math:: \text{GW}(\mathbf{T}, \mathbf{C}_1, \mathbf{C}_2) = \sum_{ijkl} T_{ik} T_{jl} \ell(C_{1,ij}, C_{2,kl}) = \langle \mathcal{L} \times \mathbf{T}

(
    a, b, C1, C2, symmetric=True, nx=None, loss="sqeuclidean", logits=None
)

def tensor_batch(
    a, b, C1, C2, symmetric=True, nx=None, loss="sqeuclidean", logits=None
):
    r"""
    Compute the Gromov-Wasserstein cost tensor for a batch of problems.

    The Gromov-Wasserstein distance can be expressed as:

    .. math::
        \text{GW}(\mathbf{T}, \mathbf{C}_1, \mathbf{C}_2) = \sum_{ijkl} T_{ik} T_{jl} \ell(C_{1,ij}, C_{2,kl}) = \langle \mathcal{L} \times \mathbf{T}, \mathbf{T} \rangle

    where :math:`\mathcal{L}` is a 4D tensor with elements :math:`\mathcal{L}[i,j,k,l] = \ell(C_{1,ij}, C_{2,kl})`.

    For loss functions of the form :math:`\ell(a,b) = f_1(a) + f_2(b) - \langle h_1(a), h_2(b) \rangle`,
    the tensor product :math:`\mathcal{L} \times \mathbf{T}` can be computed efficiently without explicitly computing :math:`\mathcal{L}` [12].

    This function precomputes all matrices that implicitly define `\mathcal{L}` for various loss functions.

    Parameters
    ----------
    a : array-like, shape (B, n)
        Source distributions for each problem in the batch.
    b : array-like, shape (B, m)
        Target distributions for each problem in the batch.
    C1 : array-like, shape (B, n, n) or (B, n, n, d)
        Source cost matrices for each problem. Can be a 3D array for scalar costs or a 4D array for vector-valued costs (edge features).
    C2 : array-like, shape (B, m, m) or (B, n, n, d)
        Target cost matrices for each problem. Can be a 3D array for scalar costs or a 4D array for vector-valued costs (edge features).
    symmetric : bool, optional
        Whether the cost matrices are symmetric. Default is True.
    nx : backend object, optional
        Numerical backend to use for computations. If None, the default backend is used.
    loss : str, optional
        Loss function to use. Supported values: 'sqeuclidean', 'kl'.
        Default is 'sqeuclidean'.
    logits : bool, optional
        For KL divergence, whether inputs are logits (unnormalized log probabilities).
        If True, inputs are treated as logits. Default is None.

    Returns
    -------
    dict
        Dictionary containing:
        - constC : array-like, shape (B, n, m)
            Constant term in the tensor product.
        - hC1 : array-like, shape (B, n, n, d) or (B, n, n)
        - hC2 : array-like, shape (B, m, m, d) or (B, m, m)
        - fC1 : array-like, shape (B, n, n)
        - fC2 : array-like, shape (B, m, m)


    Supported loss functions:
    ------------------------------

    **Squared Euclidean loss**:

    .. math::
        \ell(a, b) = \|a - b\|_2^2 = \sum_i (a_i - b_i)^2