- class ott.geometry.low_rank.LRCGeometry(cost_1, cost_2, bias=0.0, scale_factor=1.0, scale_cost=1.0, batch_size=None, **kwargs)[source]#
Geometry whose cost is defined by product of two low-rank matrices.
Implements geometries that are defined as low rank products, i.e. for which there exists two matrices \(A\) and \(B\) of \(r\) columns such that the cost of the geometry equals \(AB^T\). Apart from being faster to apply to a vector, these geometries are characterized by the fact that adding two such geometries should be carried out by concatenating factors, i.e. if \(C = AB^T\) and \(D = EF^T\) then \(C + D = [A,E][B,F]^T\)
Array) – jnp.ndarray<float>[num_a, r]
Array) – jnp.ndarray<float>[num_b, r]
float) – constant added to entire cost matrix.
scale – Value used to rescale the factors of the low-rank geometry.
'max_cost']]) – option to rescale the cost matrix. Implemented scalings are ‘max_bound’, ‘mean’ and ‘max_cost’. Alternatively, a float factor can be given to rescale the cost such that
cost_matrix /= scale_cost. If True, use ‘mean’.
int]) – optional size of the batch to compute online (without instantiating the matrix) the scale factor
scale_cost = 'max_cost'. If None, the batch size is set to 1024 or to the largest number of samples between
cost_2if smaller than 1024.
apply_cost(arr[, axis, fn])
cost_matrixto array (vector or matrix).
apply_kernel(scaling[, eps, axis])
kernel_matrixon positive scaling vector.
apply_lse_kernel(f, g, eps[, vec, axis])
kernel_matrixin log domain.
Apply elementwise-square of cost matrix to array (vector or matrix).
apply_transport_from_potentials(f, g, vec[, ...])
Apply transport matrix computed from potentials to a (batched) vec.
apply_transport_from_scalings(u, v, vec[, axis])
Apply transport matrix computed from scalings to a (batched) vec.
Compute the maximum of the
Copy the epsilon parameters from another geometry.
marginal_from_potentials(f, g[, axis])
Output marginal of transportation matrix from potentials.
marginal_from_scalings(u, v[, axis])
Output marginal of transportation matrix from scalings.
mask(src_mask, tgt_mask[, mask_value])
Mask rows or columns of a geometry.
Compute dual potential vector from scaling vector.
Instantiate 2 (or 3) geometries to compute a Sinkhorn divergence.
Compute scaling vector from dual potential.
Modify how to rescale of the
subset(src_ixs, tgt_ixs, **kwargs)
Subset rows or columns of a geometry.
to_LRCGeometry([rank, tol, rng, scale])
Output transport matrix from potentials.
Output transport matrix from pair of scalings.
update_potential(f, g, log_marginal[, ...])
Carry out one Sinkhorn update for potentials, i.e. in log space.
update_scaling(scaling, marginal[, ...])
Carry out one Sinkhorn update for scalings, using kernel directly.
Constant offset added to the entire
Check quickly if casting geometry as LRC makes sense.
First factor of the
Second factor of the
Materialize the cost matrix.
Output rank of cost matrix, if any was provided.
The data type.
Epsilon regularization value.
Compute and return inverse of scaling factor for cost matrix.
Whether geometry cost/kernel should be recomputed on the fly.
Whether cost is computed by taking squared Euclidean distance.
Whether geometry cost/kernel is a symmetric matrix.
Mean of the
Median of the
Shape of the geometry.
Mask of shape
Mask of shape