SoftJAX: Soft Differentiable Programming

• SoftJAX is a JAX-native library that provides soft differentiable programming by replacing hard primitives with smooth, informative gradient surrogates.
• It supports a wide range of operations including thresholding, Boolean logic, and sorting using multiple smoothness regimes and straight-through estimation techniques.
• Designed for integration with JAX idioms like jit and vmap, SoftJAX facilitates end-to-end differentiation in optimization pipelines while addressing gradient limitations.

SoftJAX is a JAX-native library for soft differentiable programming that replaces “hard” primitives with soft surrogates that yield informative gradients while remaining compatible with gradient-based optimization (Paulus et al., 9 Mar 2026). Introduced alongside SoftTorch in “SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients” (Paulus et al., 9 Mar 2026), it targets a recurrent limitation of automatic differentiation frameworks: thresholding, Boolean logic, discrete indexing, sorting, ranking, and related operations often have zero gradients on large domains or undefined subgradients at nondifferentiable points. SoftJAX addresses this by providing drop-in soft counterparts for hard JAX operators, together with straight-through estimation (STE), multiple smoothness regimes, and JAX-specific implementation features such as jit, vmap, lax.scan, and custom vector-Jacobian products (Paulus et al., 9 Mar 2026).

1. Motivation and conceptual scope

Automatic differentiation frameworks perform well when the computational graph is smooth and gradients are informative, but many hard primitives are discontinuous or flat almost everywhere (Paulus et al., 9 Mar 2026). The motivating examples listed for SoftJAX include thresholding and Boolean logic through Heaviside-type operations, discrete indexing and branching through argmax and where, and axiswise global operations such as sort, rank, top-k, quantile, and median (Paulus et al., 9 Mar 2026). These operations can return zero gradients on large domains or undefined subgradients at nondifferentiable points, including ties in sorting and the point $0$ in ReLU (Paulus et al., 9 Mar 2026).

SoftJAX operationalizes soft differentiable programming by systematically replacing a hard operator with a soft surrogate $f_\tau$ that is continuous, often smooth, produces informative gradients, and recovers the hard operator as $\tau \to 0^+$ (Paulus et al., 9 Mar 2026). This design is not restricted to one class of primitive. The library covers elementwise operators, fuzzy-logic-style Boolean and indexing utilities, axiswise ordering and ranking operators, and STE wrappers (Paulus et al., 9 Mar 2026).

The paper characterizes the library as open-source and feature-complete, with SoftJAX serving the JAX ecosystem and SoftTorch providing cross-framework parity for PyTorch-like workflows (Paulus et al., 9 Mar 2026). A plausible implication is that the project is intended not merely as a collection of ad hoc relaxations, but as a unified interface for comparing and composing existing soft operators under a common API.

2. Core design: soft surrogates and straight-through estimation

The central design principle is replacement of hard primitives by soft surrogates. In the forward model, these surrogates approximate the hard operation; in the backward pass, they provide nontrivial derivatives suitable for optimization (Paulus et al., 9 Mar 2026). SoftJAX exposes a softness parameter $\tau$ through softness=..., and supports modes "smooth", "c0", "c1", and "c2" to trade off smoothness, sparsity, and faithfulness to the hard primitive (Paulus et al., 9 Mar 2026).

For applications where the forward behavior must remain hard, SoftJAX provides principled STE. Its JAX implementation is given by

$$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$

so the forward computation equals the hard primitive $f$, while the backward pass uses $\nabla f_\tau$ (Paulus et al., 9 Mar 2026). The library includes a decorator sj.st to apply STE safely to composite functions (Paulus et al., 9 Mar 2026).

A central warning in the paper is the “STE pitfall” in composite expressions (Paulus et al., 9 Mar 2026). If subexpressions are wrapped separately, multiplicative interactions of hard outputs can zero out gradients:

$$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$

The recommended remedy is to apply STE to the composite instead:

$$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$

Accordingly, the paper recommends decorating the outer function with @sj.st and using soft primitives internally (Paulus et al., 9 Mar 2026).

This architecture makes SoftJAX usable in two distinct regimes. Fully soft computation is intended for settings where the relaxed forward model is acceptable and dense gradients are useful. STE is intended for settings where exact forward semantics are necessary, such as simulators and control, but informative surrogate gradients are still desired (Paulus et al., 9 Mar 2026).

3. Families of supported primitives

SoftJAX groups its primitives into elementwise, Boolean/indexing, and axiswise operators, plus STE wrappers (Paulus et al., 9 Mar 2026). The following summary reflects the catalog described in the paper.

Family	Examples	Mechanisms
Elementwise	clip, abs, round, ReLU	Heaviside surrogates, Softplus, SiLU
Boolean and indexing	comparisons, where, take_along_axis	fuzzy logic, SoftBool, SoftIndex
Axiswise	sort, rank, top-k, quantile, median	OT, simplex projections, permutahedron, sorting networks

Elementwise operators are built from softened Heaviside functions and related constructions (Paulus et al., 9 Mar 2026). The catalog includes Heaviside surrogates in several smoothness modes, soft sign, soft absolute value, soft round, and two forms of soft ReLU and clipping: an integration-based construction yielding Softplus and a gating-based construction yielding SiLU (Paulus et al., 9 Mar 2026). SoftJAX also provides autograd-safe wrappers for sqrt, arcsin, arccos, div, log, and norm to clamp gradients near singularities (Paulus et al., 9 Mar 2026).

Boolean and indexing utilities are formulated through fuzzy logic (Paulus et al., 9 Mar 2026). Comparisons such as greater-than, less-than, equality, and isclose produce SoftBool outputs in $[0,1]$ (Paulus et al., 9 Mar 2026). Logical all, any, and, or, xor, and not are supported, with default conjunction implemented as the product t-norm and disjunction derived through De Morgan’s law (Paulus et al., 9 Mar 2026). The paper also notes optional geometric mean scaling for all (Paulus et al., 9 Mar 2026). Soft selection and indexing are provided through where, dynamic_index_in_dim, take_along_axis, and choose, all using SoftIndex weights on the simplex $f_\tau$0 (Paulus et al., 9 Mar 2026).

Axiswise operators are the most elaborate component. SoftJAX implements argsort, sort, rank, top-k, quantile, and median through three main method families: optimal transport over the Birkhoff polytope, unit-simplex projection methods such as SoftSort and NeuralSort, and permutahedron projection methods including FastSoftSort and the newly introduced SmoothSort (Paulus et al., 9 Mar 2026). It also provides bitonic sorting networks with soft compare-and-swap based on softened Heaviside operators, and these networks can produce soft permutation matrices for argsort and argmax (Paulus et al., 9 Mar 2026).

The library’s method selection is operator-dependent. The paper states that defaults depend on the operator, with NeuralSort as a default for sort, rank, median, and quantile, and SoftSort as a default for argmax, argmin, max, and min (Paulus et al., 9 Mar 2026).

4. Mathematical formulations and smoothness regimes

SoftJAX adopts a canonical family of Heaviside relaxations (Paulus et al., 9 Mar 2026). In "smooth" mode, the softened Heaviside is the logistic sigmoid,

$f_\tau$1

with derivative

$f_\tau$2

For compactly supported piecewise modes on $f_\tau$3, the interpolants are defined through $f_\tau$4:

• $f_\tau$5: $f_\tau$6,
• $f_\tau$7: $f_\tau$8,
• $f_\tau$9: $\tau \to 0^+$0 (Paulus et al., 9 Mar 2026).

The paper notes that in the library, c0/c1/c2 inputs are rescaled by $\tau \to 0^+$1 to match the effective transition width of the smooth mode, since $\tau \to 0^+$2 (Paulus et al., 9 Mar 2026).

Soft sign and soft absolute value are derived from $\tau \to 0^+$3:

$\tau \to 0^+$4

In smooth mode, this yields

$\tau \to 0^+$5

The paper additionally mentions an alternative softsign convention,

$\tau \to 0^+$6

and an alternative soft absolute value used in the literature,

$\tau \to 0^+$7

while emphasizing that SoftJAX’s built-in abs uses $\tau \to 0^+$8 (Paulus et al., 9 Mar 2026).

For ReLU, the paper distinguishes two canonical constructions. The integration-style variant is

$\tau \to 0^+$9

which in smooth mode becomes Softplus,

$\tau$0

The gating-style variant is

$\tau$1

which yields SiLU-like derivatives (Paulus et al., 9 Mar 2026). Soft clipping is then composed from soft ReLU:

$\tau$2

The fuzzy-logic layer defines SoftBool probabilities in $\tau$3 (Paulus et al., 9 Mar 2026). The default conjunction is the product t-norm,

$\tau$4

with disjunction through De Morgan,

$\tau$5

and negation

$\tau$6

(Paulus et al., 9 Mar 2026). Alternative Łukasiewicz and Gödel/min t-norm and t-conorm families are discussed conceptually but are not the defaults (Paulus et al., 9 Mar 2026).

Soft comparisons use softened Heaviside functions on pairwise differences. For example,

$\tau$7

with analogous definitions for soft greater-or-equal, less-or-equal, and isclose (Paulus et al., 9 Mar 2026). In smooth mode,

$\tau$8

Soft indexing is attention-style (Paulus et al., 9 Mar 2026). For values $\tau$9 and an index variable $$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$0, SoftJAX constructs weights on the simplex,

$$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$1

with distances such as $$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$2 or $$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$3 (Paulus et al., 9 Mar 2026). The associated gradients are those of softmax composed with the chosen distance function.

5. Axiswise ordering methods and computational structure

The optimal-transport family treats soft sorting and ranking as regularized transport on the transport polytope $$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$4 (Paulus et al., 9 Mar 2026). The objective minimizes

$$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$5

subject to $$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$6 and $$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$7, with costs $$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$8 and balanced marginals (Paulus et al., 9 Mar 2026). A soft permutation matrix is recovered as

$$f_{\mathrm{ST}}(x) = \mathrm{sg}(f(x)) + f_\tau(x) - \mathrm{sg}(f_\tau(x)),$$9

from which

$f$0

(Paulus et al., 9 Mar 2026).

For entropic regularization, SoftJAX uses Sinkhorn iterations and states the standard scaling steps

$f$1

with $f$2 (Paulus et al., 9 Mar 2026). Entropic regularization yields $f$3 everywhere, dense Jacobians, and a smooth $f$4 map (Paulus et al., 9 Mar 2026). Euclidean regularization promotes sparsity and is piecewise smooth; $f$5-norm regularization with $f$6 promotes sparsity while retaining controlled smoothness, and the library provides c1 and c2 settings with regularity guarantees under connected support graphs (Paulus et al., 9 Mar 2026).

The unit-simplex projection family includes SoftSort and NeuralSort (Paulus et al., 9 Mar 2026). SoftSort approximates Sinkhorn through row-wise projection against hard anchors,

$f$7

where $f$8 denotes row projection to $f$9 using entropic, Euclidean, or $\nabla f_\tau$0-norm regularization (Paulus et al., 9 Mar 2026). NeuralSort forms argsort distributions by projecting a linear combination of $\nabla f_\tau$1 and sums of soft absolute differences. Its construction uses a matrix $\nabla f_\tau$2 with entries $\nabla f_\tau$3 and row formulas depending on the target rank, argmax, argmin, or quantile (Paulus et al., 9 Mar 2026). Because SoftJAX uses soft abs, NeuralSort becomes fully smooth in "smooth" mode (Paulus et al., 9 Mar 2026).

The permutahedron family includes FastSoftSort and SmoothSort (Paulus et al., 9 Mar 2026). FastSoftSort projects onto the permutahedron under Euclidean or log-KL objectives, solving the resulting isotonic regression problem by pool adjacent violators in $\nabla f_\tau$4 time and $\nabla f_\tau$5 memory (Paulus et al., 9 Mar 2026). It returns values rather than permutation matrices, so argsort and argrank are unavailable (Paulus et al., 9 Mar 2026). SmoothSort, introduced in SoftJAX, adds entropic regularization in a dual formulation and replaces hard majorization bounds with smooth log-sum-exp bounds,

$\nabla f_\tau$6

yielding $\nabla f_\tau$7 differentiability and dense Jacobians (Paulus et al., 9 Mar 2026). The dual is solved with L-BFGS and a custom VJP, with $\nabla f_\tau$8 preprocessing for the bounds but without materializing $\nabla f_\tau$9 matrices (Paulus et al., 9 Mar 2026).

SoftJAX also includes bitonic sorting networks in which compare-and-swap is replaced by soft comparisons $$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$0 (Paulus et al., 9 Mar 2026). The associated soft minimum and maximum are

$$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$1

By propagating a soft permutation matrix through the network, the method can produce both sorted outputs and soft argsort values (Paulus et al., 9 Mar 2026).

The paper provides explicit complexity claims. Sinkhorn-based soft sort has $$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$2 time and $$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$3 memory, though SoftJAX reduces memory through chunked scans with lax.scan and uses implicit differentiation to avoid storing all intermediates (Paulus et al., 9 Mar 2026). FastSoftSort has $$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$4 time and $$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$5 memory (Paulus et al., 9 Mar 2026). Soft indexing operates in $$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$6 time and $$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$7 memory and is fully compatible with vmap (Paulus et al., 9 Mar 2026).

6. API, numerical stability, and practical use

SoftJAX is described as drop-in with respect to JAX idioms: all functions support jax.grad, vmap, jit, and batched axes (Paulus et al., 9 Mar 2026). Methods are selected through method=..., smoothness through mode=..., and softness through softness=... (Paulus et al., 9 Mar 2026). For axiswise operators, standardize_and_squash can be enabled to improve numerical stability and scale-independence (Paulus et al., 9 Mar 2026).

The paper defines standardize-and-squash as

$$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$8

with inverse transformation for value outputs by logit and destandardization (Paulus et al., 9 Mar 2026). This is recommended for OT and simplex-based methods, but the paper explicitly advises disabling it for SmoothSort because values may leave the convex hull (Paulus et al., 9 Mar 2026).

Several numerical stability recommendations are stated directly. Too small a softness parameter leads to hard behavior and large gradients; too large a value leads to overly smooth, mean-like behavior (Paulus et al., 9 Mar 2026). Log-domain implementations are recommended for Sinkhorn to avoid underflow and overflow, and gradient clipping is recommended for very small $$\nabla(f_{\mathrm{ST}}\cdot g_{\mathrm{ST}}) = (\nabla f_\tau)\cdot g + f\cdot(\nabla g_\tau).$$9 (Paulus et al., 9 Mar 2026). For safe wrappers around singular functions such as sqrt, arcsin, arccos, div, log, and norm, gradients are clamped near boundaries (Paulus et al., 9 Mar 2026).

The guidance for choosing $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$0 is also explicit (Paulus et al., 9 Mar 2026). For axiswise operators after standardize-and-squash, the recommended starting range is approximately $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$1 to $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$2, followed by annealing downward during training (Paulus et al., 9 Mar 2026). The paper further proposes calibrating softness by the normalized entropy of the soft permutation, $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$3, to compare behavior across methods and problem sizes (Paulus et al., 9 Mar 2026).

The API examples in the paper illustrate three usage modes: fully soft sorting and ranking through methods such as NeuralSort, soft Boolean masking through greater, less, logical_and, and logical_or, and soft indexing through argmax followed by dynamic_index_in_dim (Paulus et al., 9 Mar 2026). These examples underscore that the library is intended for compositional use rather than isolated operator substitution.

7. Benchmarks, case study, related work, and limitations

The reported benchmarks were obtained on an RTX 3060 GPU in smooth mode for the forward-backward pass of sj.sort (Paulus et al., 9 Mar 2026). At $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$4, the sorting network is reported as the fastest soft method at approximately $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$5 ms, about $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$6 the hard baseline; SoftSort is approximately $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$7 ms; NeuralSort is approximately $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$8 ms (Paulus et al., 9 Mar 2026). FastSoftSort is identified as the most memory-efficient method and is reported to scale roughly linearly, using approximately $$\nabla(f\cdot g)_{\mathrm{ST}} = \nabla(f_\tau\cdot g_\tau) = (\nabla f_\tau)\cdot g_\tau + f_\tau\cdot(\nabla g_\tau).$$9 KB at $[0,1]$0 because it avoids permutation matrices (Paulus et al., 9 Mar 2026). OT and SmoothSort are the slowest among the tested methods, due respectively to iterative solves and $[0,1]$1 preprocessing, but they provide full smoothness and dense Jacobians (Paulus et al., 9 Mar 2026).

The principal application example is a practical case study in MJX collision detection (Paulus et al., 9 Mar 2026). There, replacing hard argmax and abs with sj.argmax and sj.abs, together with soft indexing, yields smooth, non-zero gradients at all polygon vertices and enables end-to-end differentiation of contact manifold selection (Paulus et al., 9 Mar 2026). Applying STE to the outer routine preserves the hard forward pass while using informative gradients in backpropagation (Paulus et al., 9 Mar 2026).

The paper situates SoftJAX relative to several research lines (Paulus et al., 9 Mar 2026). It states that SoftJAX generalizes SoftSort across entropic, Euclidean, and $[0,1]$2-norm regularizers; adds NeuralSort with fully smooth soft abs, ranking, and sorting networks; and introduces SmoothSort with $[0,1]$3 behavior through smooth majorization bounds (Paulus et al., 9 Mar 2026). It notes that its simplex projections reproduce softmax under entropic regularization and Sparsemax under Euclidean regularization, while extending to $[0,1]$4-norm projections with C^1 and C^2 guarantees (Paulus et al., 9 Mar 2026). It also distinguishes deterministic relaxations, which are the focus of SoftJAX, from stochastic relaxations such as Gumbel-Sinkhorn, which are described as complementary but outside the library’s scope (Paulus et al., 9 Mar 2026). Among external dependencies, the paper identifies ott-jax for Sinkhorn computations and Optimistix for L-BFGS with implicit differentiation (Paulus et al., 9 Mar 2026).

The limitations are presented in operational terms rather than as a single theoretical objection (Paulus et al., 9 Mar 2026). Ties and discrete symmetries can produce non-unique soft permutations; the paper recommends entropic regularization for stable gradients or adding jitter (Paulus et al., 9 Mar 2026). Very small $[0,1]$5 can produce large gradient magnitudes, and Euclidean or $[0,1]$6-norm modes can yield sparse Jacobians that slow signal propagation (Paulus et al., 9 Mar 2026). Fully soft methods optimize a biased surrogate objective but offer low-variance gradients; STE reduces forward bias but still propagates surrogate gradients and can remain biased (Paulus et al., 9 Mar 2026). This suggests that SoftJAX is best understood not as a universal replacement for hard operators, but as a library of controlled relaxations whose suitability depends on the interaction among forward semantics, optimization geometry, and memory-computation constraints.