Repository navigation
Conversation
Removed TORCH_CUDA_ARCH_LIST environment variable from the workflow.
…, batched (n, d) API - Use the standard convention K = exp(-C / epsilon); POT baselines need no parameter inversion anymore. - Keep log potentials in log space end-to-end: the log-domain CUDA solver returns (log_u, log_v), gradients are computed as f = eps * log_u, and the transport plan is reconstructed as exp(log_u - C/eps + log_v) for all backends so it cannot underflow at small epsilon. - Every backend takes strictly (n, d) batched histograms sharing one (d, d) cost matrix; the CUDA kernels are batched over a 2-D launch grid, removing the previous B=C=1 restriction. - Merge the four kernels into two (scaling_update, log_scaling_update), fixing the uncoalesced K^T reads; collapse entry points into solve(); delete the dead sinkhorn__nobatch and loss_test. - CUDA hygiene: current-stream launches, CUDAGuard, TORCH_CHECK input validation, kernel-launch checks, and a NaN guard for all -inf LSE rows. - Export SinkhornSolver / build_cost_matrix from the package; tests import it directly, run the torch backend on CPU, and add a small-epsilon stability test against POT's sinkhorn_log. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…dispatch - SinkhornSolver now inherits torch.nn.Module: forward(source, target, cost_matrix) returns the (n,) transport distances, and a custom autograd.Function implements backward via the centered dual potentials (envelope theorem), chained through the differentiable normalization. - Replace the confused torch/cuda/cuda_log backend strings with one orthogonal option: log_space=True/False selects the algorithm, while the implementation is picked from the input automatically (fused CUDA kernels for float32 CUDA tensors, pure torch ops elsewhere - including a logsumexp-based log-space path that also runs on CPU and float64). - Convenience methods plan() and potentials() expose the transport plan and dual potentials; extra_repr() reports the module configuration. - Tests cover both log_space modes across forward shapes, marginals, POT baselines, fused-vs-CPU consistency, small-epsilon stability, and autograd-vs-numeric gradients; benchmark updated to the module API. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Previously most tests picked a single device (CUDA when available), so the CPU paths went untested on GPU machines, POT baselines only covered the plain iteration, and backward never ran through the fused kernels. Now: - Shared DEVICES parametrization runs every applicable test on both CPU and CUDA (skipping cleanly when CUDA is absent). - POT distance/plan baselines and the small-epsilon stability test cover both log_space modes and both devices (torch logsumexp path on CPU, fused log kernel on CUDA). - New tests: dtype/device consistency (float64 CUDA exercises the torch fallback on GPU), backward through the fused float32 kernels vs CPU, fused-vs-CPU potentials, and threshold early stopping reaching the same solution on both algorithms. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Kernel changes (both memory-bandwidth-bound, so all wins target matrix traffic and launch latency): - Batch tiling: one block now streams a matrix row once and applies it to up to 8 batch items held in registers, cutting d^2 matrix reads by up to 8x for batched inputs. A TILE=1 template instantiation keeps the n=1 fast path free of predication overhead. - Online logsumexp: the log-domain update now computes the LSE in a single pass over the row (running max with rescaled running sum, NaN-free -inf merging), halving matrix traffic vs the two-pass max-then-sum version. - Warp-shuffle reductions replace the shared-memory tree (one barrier instead of eight per reduction). Python-side: - Long fused runs (max_iter >= 100) capture a 25-iteration chunk into a torch.cuda.CUDAGraph and replay it, hiding the 2-launches-per-iteration overhead that dominates small problems; falls back to plain launches when capture is unavailable (extra fixed-point iterations are harmless). - Skip the (d, d) transpose copy when the cost/kernel matrix is symmetric, which build_cost_matrix outputs always are. Tests: fused-vs-CPU now covers single-item, partial-tile, and multi-tile batches (n = 1, 3, 20) plus a short-run case below the CUDA-graph threshold to exercise the plain-launch fused path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The backward pass already returned gradients for both marginals (the centered dual potentials f and g), but the tests only verified the source side. The numeric-derivative test now perturbs source and target independently, and the fused-kernel backward test compares both gradients against CPU. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
zeakey
added a commit
that referenced
this pull request
Jul 18, 2026
* feature:add grey erosion kernel * ENG * dilation opening closing funciton * rename * gradient black/white tophat * iterate structure * grey laplace * CUDA binary erosion dilation * delete convnd * fix origin bug * binary propagation * format * update readme * readme * binary_fill_hole and binary_hit_or_miss * fix bugs and add validation tests * update .gitignore * back validation and refactor distance transform * delete JFA argoithm * add CUDAGuard * add validation test * optimal transport (#6) * optimal transport * debug ot function and add test * add basic benchmark * decouple the iteration from the derivation * Implement Sinkhorn iterations in CUDA * Add Sinkhorn CUDA benchmark and correctness tests * Add CUDA support for Sinkhorn log-domain iterations and related tests. * Refactor Sinkhorn functions and test gradient accuracy * more test * try test * test on cuda kernal * tighten numeric gradient test and support more dimensions * cancel unnecessary change in test.yml Removed TORCH_CUDA_ARCH_LIST environment variable from the workflow. * Rework Sinkhorn OT: standard epsilon convention, log-space end-to-end, batched (n, d) API - Use the standard convention K = exp(-C / epsilon); POT baselines need no parameter inversion anymore. - Keep log potentials in log space end-to-end: the log-domain CUDA solver returns (log_u, log_v), gradients are computed as f = eps * log_u, and the transport plan is reconstructed as exp(log_u - C/eps + log_v) for all backends so it cannot underflow at small epsilon. - Every backend takes strictly (n, d) batched histograms sharing one (d, d) cost matrix; the CUDA kernels are batched over a 2-D launch grid, removing the previous B=C=1 restriction. - Merge the four kernels into two (scaling_update, log_scaling_update), fixing the uncoalesced K^T reads; collapse entry points into solve(); delete the dead sinkhorn__nobatch and loss_test. - CUDA hygiene: current-stream launches, CUDAGuard, TORCH_CHECK input validation, kernel-launch checks, and a NaN guard for all -inf LSE rows. - Export SinkhornSolver / build_cost_matrix from the package; tests import it directly, run the torch backend on CPU, and add a small-epsilon stability test against POT's sinkhorn_log. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Apply black formatting to optimal transport files Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Make SinkhornSolver a differentiable nn.Module with automatic device dispatch - SinkhornSolver now inherits torch.nn.Module: forward(source, target, cost_matrix) returns the (n,) transport distances, and a custom autograd.Function implements backward via the centered dual potentials (envelope theorem), chained through the differentiable normalization. - Replace the confused torch/cuda/cuda_log backend strings with one orthogonal option: log_space=True/False selects the algorithm, while the implementation is picked from the input automatically (fused CUDA kernels for float32 CUDA tensors, pure torch ops elsewhere - including a logsumexp-based log-space path that also runs on CPU and float64). - Convenience methods plan() and potentials() expose the transport plan and dual potentials; extra_repr() reports the module configuration. - Tests cover both log_space modes across forward shapes, marginals, POT baselines, fused-vs-CPU consistency, small-epsilon stability, and autograd-vs-numeric gradients; benchmark updated to the module API. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Expand Sinkhorn tests to a full device x log_space x dtype matrix Previously most tests picked a single device (CUDA when available), so the CPU paths went untested on GPU machines, POT baselines only covered the plain iteration, and backward never ran through the fused kernels. Now: - Shared DEVICES parametrization runs every applicable test on both CPU and CUDA (skipping cleanly when CUDA is absent). - POT distance/plan baselines and the small-epsilon stability test cover both log_space modes and both devices (torch logsumexp path on CPU, fused log kernel on CUDA). - New tests: dtype/device consistency (float64 CUDA exercises the torch fallback on GPU), backward through the fused float32 kernels vs CPU, fused-vs-CPU potentials, and threshold early stopping reaching the same solution on both algorithms. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Optimize Sinkhorn CUDA path: batch tiling, online logsumexp, CUDA graphs Kernel changes (both memory-bandwidth-bound, so all wins target matrix traffic and launch latency): - Batch tiling: one block now streams a matrix row once and applies it to up to 8 batch items held in registers, cutting d^2 matrix reads by up to 8x for batched inputs. A TILE=1 template instantiation keeps the n=1 fast path free of predication overhead. - Online logsumexp: the log-domain update now computes the LSE in a single pass over the row (running max with rescaled running sum, NaN-free -inf merging), halving matrix traffic vs the two-pass max-then-sum version. - Warp-shuffle reductions replace the shared-memory tree (one barrier instead of eight per reduction). Python-side: - Long fused runs (max_iter >= 100) capture a 25-iteration chunk into a torch.cuda.CUDAGraph and replay it, hiding the 2-launches-per-iteration overhead that dominates small problems; falls back to plain launches when capture is unavailable (extra fixed-point iterations are harmless). - Skip the (d, d) transpose copy when the cost/kernel matrix is symmetric, which build_cost_matrix outputs always are. Tests: fused-vs-CPU now covers single-item, partial-tile, and multi-tile batches (n = 1, 3, 20) plus a short-run case below the CUDA-graph threshold to exercise the plain-launch fused path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Test gradients w.r.t. both source and target The backward pass already returned gradients for both marginals (the centered dual potentials f and g), but the tests only verified the source side. The numeric-derivative test now perturbs source and target independently, and the fused-kernel backward test compares both gradients against CPU. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: maokovski <1498743918@qq.com> Co-authored-by: maokovski <“1498743918@qq.com”> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> * refactor: share validation, cut duplication, drop redundant work - Add torchmorph/_validation.py with shared validate_bcs_input / validate_output; remove three near-identical copies from binary, grey, and distance_transform (error messages preserved). - binary: drop per-iteration tensor .clone() in the stabilize loop (step/where already return fresh tensors) and skip the redundant bool conversion when the value is already bool + contiguous. - Move _validate_origin to structure.py; grey now reuses _normalize_origin/_validate_origin instead of hand-rolling it. - grey: pack the shared structuring-element args via _element_kwargs across the composite operators. Net ~-108 lines; behavior unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * update requirements * smaller head file --------- Co-authored-by: Yuhan Deng <yuhandeng@corgi.intcomp.org> Co-authored-by: Yuhandeng <moliflower@233.com> Co-authored-by: Kai Zhao <694691@qq.com> Co-authored-by: maokovski <1498743918@qq.com> Co-authored-by: maokovski <“1498743918@qq.com”> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implementation of the entropy-regularized Wasserstein distance and its derivative.
References