torchpsort 0.0.1__py3-none-any.whl

torchpsort/__init__.py

@@ -0,0 +1,3 @@
+from .ops import soft_rank, soft_sort
+
+__all__ = ["soft_sort", "soft_rank"]

torchpsort/isotonic.py

@@ -0,0 +1,180 @@
+#  Copyright 2007-2020 The scikit-learn developers.
+#  Copyright 2020 Google LLC.
+#  Copyright 2021 Teddy Koker.
+#  Copyright 2026 RektPunk
+#
+#  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#    a. Redistributions of source code must retain the above copyright notice,
+#       this list of conditions and the following disclaimer.
+#    b. Redistributions in binary form must reproduce the above copyright
+#       notice, this list of conditions and the following disclaimer in the
+#       documentation and/or other materials provided with the distribution.
+#    c. Neither the name of the Scikit-learn Developers  nor the names of
+#       its contributors may be used to endorse or promote products
+#       derived from this software without specific prior written
+#       permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+#  ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE FOR
+#  ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+#  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+#  OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+#  DAMAGE.
+
+import torch
+from torch import Tensor
+
+
+@torch.jit.script
+def isotonic_l2_forward(x: Tensor) -> Tensor:
+    batch_size, p = x.shape
+    solution = x.clone()
+    block_sums = x.clone()
+    block_counts = torch.ones_like(x)
+    block_end = (
+        torch.arange(p, device=x.device, dtype=torch.long)
+        .expand(batch_size, -1)
+        .clone()
+    )
+    for b in range(batch_size):
+        block = 0
+        while block < p:
+            next_block = block_end[b, block] + 1
+            if next_block == p:
+                break
+            if solution[b, block] > solution[b, next_block]:
+                block = next_block
+                continue
+
+            block_sum = block_sums[b, block]
+            block_count = block_counts[b, block]
+
+            while True:
+                last_value = solution[b, next_block]
+                block_sum += block_sums[b, next_block]
+                block_count += block_counts[b, next_block]
+                next_block = block_end[b, next_block] + 1
+
+                if next_block == p or last_value > solution[b, next_block]:
+                    solution[b, block] = block_sum / block_count
+                    block_sums[b, block] = block_sum
+                    block_counts[b, block] = block_count
+                    block_end[b, block] = next_block - 1
+                    block_end[b, next_block - 1] = block
+
+                    if block > 0:
+                        block = block_end[b, block - 1]
+                    break
+
+        block = 0
+        while block < p:
+            next_block = block_end[b, block] + 1
+            solution[b, (block + 1) : next_block] = solution[b, block]
+            block = next_block
+
+    return solution
+
+
+@torch.jit.script
+def isotonic_kl_forward(x: Tensor, w: Tensor) -> Tensor:
+    batch_size, p = x.shape
+    solution = x - w
+    block_lse_x = x.clone()
+    block_lse_w = w.clone()
+    block_end = (
+        torch.arange(p, device=x.device, dtype=torch.long)
+        .expand(batch_size, -1)
+        .clone()
+    )
+
+    for b in range(batch_size):
+        block = 0
+        while block < p:
+            next_block = block_end[b, block] + 1
+            if next_block == p:
+                break
+            if solution[b, block] > solution[b, next_block]:
+                block = next_block
+                continue
+
+            curr_lse_x = block_lse_x[b, block]
+            curr_lse_w = block_lse_w[b, block]
+
+            while True:
+                last_value = solution[b, next_block]
+                curr_lse_x = torch.logaddexp(curr_lse_x, block_lse_x[b, next_block])
+                curr_lse_w = torch.logaddexp(curr_lse_w, block_lse_w[b, next_block])
+                next_block = block_end[b, next_block] + 1
+
+                if next_block == p or last_value > solution[b, next_block]:
+                    solution[b, block] = curr_lse_x - curr_lse_w
+                    block_lse_x[b, block] = curr_lse_x
+                    block_lse_w[b, block] = curr_lse_w
+                    block_end[b, block] = next_block - 1
+                    block_end[b, next_block - 1] = block
+
+                    if block > 0:
+                        block = block_end[b, block - 1]
+                    break
+
+        block = 0
+        while block < p:
+            next_block = block_end[b, block] + 1
+            solution[b, (block + 1) : next_block] = solution[b, block]
+            block = next_block
+
+    return solution
+
+
+@torch.jit.script
+def isotonic_l2_backward(sol: Tensor, grad_output: Tensor) -> Tensor:
+    batch_size, p = sol.shape
+    grad_input = torch.empty_like(grad_output)
+    tol = 1e-6 if sol.dtype == torch.float32 else 1e-12
+    for b in range(batch_size):
+        start = 0
+        while start < p:
+            end = start + 1
+            while end < p and torch.abs(sol[b, end] - sol[b, start]) < tol:
+                end += 1
+
+            val = 1.0 / float(end - start)
+
+            grad_sum = grad_output[b, start:end].sum()
+            grad_input[b, start:end] = val * grad_sum
+
+            start = end
+
+    return grad_input
+
+
+@torch.jit.script
+def isotonic_kl_backward(s: Tensor, sol: Tensor, grad_output: Tensor) -> Tensor:
+    batch_size, p = sol.shape
+    grad_input = torch.empty_like(grad_output)
+    tol = 1e-6 if sol.dtype == torch.float32 else 1e-12
+    for b in range(batch_size):
+        start = 0
+        while start < p:
+            end = start + 1
+            while end < p and torch.abs(sol[b, end] - sol[b, start]) < tol:
+                end += 1
+
+            block_s = s[b, start:end]
+            softmax_weights = torch.softmax(block_s, dim=0)
+
+            grad_sum = grad_output[b, start:end].sum()
+            grad_input[b, start:end] = softmax_weights * grad_sum
+
+            start = end
+
+    return grad_input

torchpsort/ops.py

@@ -0,0 +1,160 @@
+# Copyright 2020 Google LLC
+# Copyright 2021 Teddy Koker
+# Copyright 2026 RektPunk
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+from .isotonic import (
+    isotonic_kl_backward,
+    isotonic_kl_forward,
+    isotonic_l2_backward,
+    isotonic_l2_forward,
+)
+
+
+@torch.jit.script
+def _index_range(x: Tensor) -> Tensor:
+    return torch.arange(
+        x.shape[1],
+        dtype=x.dtype,
+        device=x.device,
+    ).expand(x.shape[0], -1)
+
+
+@torch.jit.script
+def _descending_ranks(x: Tensor) -> Tensor:
+    return torch.arange(
+        x.shape[1],
+        0,
+        -1,
+        dtype=x.dtype,
+        device=x.device,
+    ).expand_as(x)
+
+
+@torch.jit.script
+def _inv_permutation(permutation: Tensor) -> Tensor:
+    inv_permutation = torch.empty_like(permutation)
+    inv_permutation.scatter_(1, permutation, _index_range(permutation))
+    return inv_permutation
+
+
+def _validate_inputs(x: Tensor, reg: Literal["l2", "kl"], tau: float) -> None:
+    if x.ndim != 2:
+        raise ValueError("x must be a 2D tensor")
+
+    if reg not in ("l2", "kl"):
+        raise ValueError(f"Unknown reg: {reg}")
+
+    if tau <= 0:
+        raise ValueError("tau must be positive")
+
+
+class SoftRank(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: Tensor, reg: Literal["l2", "kl"] = "l2", tau: float = 1.0):
+        _validate_inputs(x, reg, tau)
+
+        ctx.scale = 1.0 / tau
+        ctx.is_l2 = reg == "l2"
+        w = _descending_ranks(x)
+        theta = x * ctx.scale
+        s, permutation = torch.sort(theta, descending=True)
+        inv_permutation = _inv_permutation(permutation)
+        if ctx.is_l2:
+            dual_sol = isotonic_l2_forward(s - w)
+            ret = (s - dual_sol).gather(1, inv_permutation)
+            ctx.save_for_backward(dual_sol, permutation, inv_permutation)
+        else:
+            dual_sol = isotonic_kl_forward(s, torch.log(w))
+            ret = torch.exp((s - dual_sol).gather(1, inv_permutation))
+            ctx.save_for_backward(ret, s, dual_sol, permutation, inv_permutation)
+
+        return ret
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        if ctx.is_l2:
+            dual_sol, permutation, inv_permutation = ctx.saved_tensors
+            grad_iso = isotonic_l2_backward(
+                dual_sol,
+                grad_outputs[0].gather(1, permutation),
+            )
+            grad = grad_outputs[0] - grad_iso.gather(1, inv_permutation)
+        else:
+            ret, s, dual_sol, permutation, inv_permutation = ctx.saved_tensors
+            grad_iso = isotonic_kl_backward(
+                s,
+                dual_sol,
+                (grad_outputs[0] * ret).gather(1, permutation),
+            )
+            grad = (grad_outputs[0] * ret) - grad_iso.gather(1, inv_permutation)
+
+        return grad * ctx.scale, None, None
+
+
+class SoftSort(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: Tensor, reg: Literal["l2", "kl"] = "l2", tau: float = 1.0):
+        _validate_inputs(x, reg, tau)
+
+        ctx.is_l2 = reg == "l2"
+        w = _descending_ranks(x) / tau
+        s, permutation = torch.sort(-x, descending=True)
+        inv_permutation = _inv_permutation(permutation)
+        if ctx.is_l2:
+            sol = isotonic_l2_forward(w - s)
+            ctx.save_for_backward(sol, inv_permutation)
+        else:
+            sol = isotonic_kl_forward(w, s)
+            ctx.save_for_backward(s, sol, inv_permutation)
+
+        return sol - w
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        if ctx.is_l2:
+            sol, inv_permutation = ctx.saved_tensors
+            grad = isotonic_l2_backward(sol, grad_outputs[0])
+        else:
+            s, sol, inv_permutation = ctx.saved_tensors
+            grad = isotonic_kl_backward(s, sol, grad_outputs[0])
+
+        return grad.gather(1, inv_permutation), None, None
+
+
+def soft_rank(x: Tensor, reg: Literal["l2", "kl"] = "l2", tau: float = 1.0) -> Tensor:
+    """Differentiable approximation of the rank operator.
+
+    Args:
+        x: Input tensor of shape (batch, p)
+        reg: Regularization type ("l2" or "kl")
+        tau: Regularization strength (temperature)
+    """
+    return SoftRank.apply(x, reg, tau)
+
+
+def soft_sort(x: Tensor, reg: Literal["l2", "kl"] = "l2", tau: float = 1.0) -> Tensor:
+    """Differentiable approximation of the sort operator.
+
+    Args:
+        x: Input tensor of shape (batch, p)
+        reg: Regularization type ("l2" or "kl")
+        tau: Regularization strength (temperature)
+    """
+    return SoftSort.apply(x, reg, tau)

torchpsort-0.0.1.dist-info/METADATA

@@ -0,0 +1,26 @@
+Metadata-Version: 2.3
+Name: torchpsort
+Version: 0.0.1
+Summary: Fast, differentiable sorting in pure PyTorch without C++ or CUDA
+Author: RektPunk
+Author-email: RektPunk <rektpunk@gmail.com>
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Dist: numpy>=2.5.1
+Requires-Dist: torch>=2.13.0
+Requires-Python: >=3.12, <=3.14
+Project-URL: repository, https://github.com/RektPunk/torchpsort
+Description-Content-Type: text/markdown
+
+<div style="text-align: center;">
+  <img src="https://capsule-render.vercel.app/api?type=transparent&fontColor=0047AB&text=torchpsort&height=120&fontSize=90">
+</div>
+
+Fast, differentiable sorting and ranking in **pure PyTorch without C++ or CUDA**. This is a lightweight implementation of [Fast Differentiable Sorting and Ranking (Blondel et al.)](https://arxiv.org/abs/2002.08871) and inspired by [torchsort](https://github.com/teddykoker/torchsort). Unlike the [torchsort](https://github.com/teddykoker/torchsort), this version contains **no C++ or CUDA extensions**, making it easy to install, portable, and compatible with any environment where PyTorch runs. While the original C++/CUDA implementation may have a performance edge for extremely large batch sizes, this pure PyTorch version is optimized to be efficient for standard deep learning workflows. Try it here:
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/RektPunk/torchpsort/blob/main/examples/notebook.ipynb)
+
+
+## Installation
+
+```bash
+pip install torchpsort
+```

torchpsort-0.0.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+torchpsort/__init__.py,sha256=uGQfsZ-ZTg1gCbIi43YefATxKxwD-1M8jAtx6d-5lOQ,76
+torchpsort/isotonic.py,sha256=7T3KsXHz-8T5IScTIc9fVACrTk4H2snJa_u7q8NqAqI,6445
+torchpsort/ops.py,sha256=UJeRUreEXEaLr7KjMH7zFO4M9ArFk8z0YsdnmamuVmk,5074
+torchpsort/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torchpsort-0.0.1.dist-info/WHEEL,sha256=YR4QUxGCbsyoOnWBVTv-p1I4yc3seKGPev46mvmc8Cg,81
+torchpsort-0.0.1.dist-info/METADATA,sha256=VU1zI_3qQ7vsYcS5bnXHq0oL97cWxAWuUE9u1RIY9GM,1556
+torchpsort-0.0.1.dist-info/RECORD,,

torchpsort-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.11.29
+Root-Is-Purelib: true
+Tag: py3-none-any