transport.samples.kernels.streaming_sqeuclid
module torchmatch.transport.samples.kernels.streaming_sqeuclidSinkhorn OT with streaming softmax and shifted potentials.
This module implements a reformulated Sinkhorn algorithm that aligns exactly with
streaming-softmax's interface, enabling potential future integration with optimized
streaming-softmax kernels.
Members
function
- compute_bias_fCompute pre-scaled bias for f-update: u = (ĝ + δ)/ε.
- compute_bias_gCompute pre-scaled bias for g-update: v = (f̂ + γ)/ε.
- precompute_sinkhorn_inputsPrecompute static bias components.
- shifted_to_standard_potentialsConvert shifted potentials back to standard form.
- sinkhorn_lseCompute shifted potential using the streaming-softmax kernel.
- sinkhorn_lse_fusedFused LSE kernel that computes bias in SRAM (matches symmetric kernel interface).
- sinkhorn_symmetric_stepFused symmetric Sinkhorn step: computes both f and g updates in ONE kernel.
- standard_to_shifted_potentialsConvert standard potentials to shifted form.
Source: torchmatch/transport/samples/kernels/streaming_sqeuclid.py:1