torchcell.losses.multi_dim_nan_tolerant.FastSoftSort

class FastSoftSort(*args, **kwargs)[source]

Bases: Function

Fast differentiable sorting following the original paper implementation.

Uses isotonic regression for the forward pass and custom gradients.

static forward(ctx: Any, values: Tensor, regularization_strength: float = 1.0) → Tensor[source]

Soft sort using isotonic regression.

Parameters:

  • ctx – Autograd context used to stash tensors for the backward pass.

  • values – 1D tensor to sort

  • regularization_strength – Controls smoothness

static backward(ctx: Any, grad_output: Tensor) → tuple[Tensor, None][source]

Compute gradients using the Jacobian of isotonic regression.