---
title: 'SoftJAX: Soft Differentiable Programming'
url: https://www.emergentmind.com/topics/softjax
type: topic
---

# SoftJAX: Soft Differentiable Programming

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 [2603.08824]. Introduced alongside SoftTorch in “SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients” [2603.08824], 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 [2603.08824].

## 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 [2603.08824]. 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` [2603.08824]. 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 [2603.08824].

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^+$ [2603.08824]. 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 [2603.08824].

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 [2603.08824]. 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 [2603.08824]. 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 [2603.08824].

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$ [2603.08824]. The library includes a decorator `sj.st` to apply STE safely to composite functions [2603.08824].

A central warning in the paper is the “STE pitfall” in composite expressions [2603.08824]. 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 [2603.08824].

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 [2603.08824].

## 3. Families of supported primitives

SoftJAX groups its primitives into elementwise, Boolean/indexing, and axiswise operators, plus STE wrappers [2603.08824]. 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 [2603.08824]. 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 [2603.08824]. SoftJAX also provides autograd-safe wrappers for `sqrt`, `arcsin`, `arccos`, `div`, `log`, and `norm` to clamp gradients near singularities [2603.08824].

Boolean and indexing utilities are formulated through fuzzy logic [2603.08824]. Comparisons such as greater-than, less-than, equality, and `isclose` produce `SoftBool` outputs in $[0,1]$ [2603.08824]. 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 [2603.08824]. The paper also notes optional geometric mean scaling for `all` [2603.08824]. Soft selection and indexing are provided through `where`, `dynamic_index_in_dim`, `take_along_axis`, and `choose`, all using `SoftIndex` weights on the simplex $\Delta_n$ [2603.08824].

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 [2603.08824]. 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` [2603.08824].

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` [2603.08824].

## 4. Mathematical formulations and smoothness regimes

SoftJAX adopts a canonical family of Heaviside relaxations [2603.08824]. In `"smooth"` mode, the softened Heaviside is the logistic sigmoid,
$$
H_\tau(x)=\sigma(x/\tau)=\frac{1}{1+e^{-x/\tau}},
$$
with derivative
$$
\frac{\partial H_\tau}{\partial x}=\frac{1}{\tau}\sigma(x/\tau)\bigl(1-\sigma(x/\tau)\bigr).
$$
For compactly supported piecewise modes on $[-\tau,\tau]$, the interpolants are defined through $s=x/\tau$:
- $c0$: $g(s)=1/2+s/2$,
- $c1$: $g(s)=1/2+3s/4-s^3/4$,
- $c2$: $g(s)=1/2+15s/16-5s^3/8+3s^5/16$ [2603.08824].

The paper notes that in the library, `c0/c1/c2` inputs are rescaled by $1/5$ to match the effective transition width of the smooth mode, since $\sigma(\pm 5)\approx 0.007/0.993$ [2603.08824].

Soft sign and soft absolute value are derived from $H_\tau$:
$$
\operatorname{sign}_\tau(x)=2H_\tau(x)-1,\qquad \operatorname{abs}_\tau(x)=\operatorname{sign}_\tau(x)\cdot x.
$$
In smooth mode, this yields
$$
\operatorname{sign}_\tau(x)=2\sigma(x/\tau)-1=\tanh(x/(2\tau)).
$$
The paper additionally mentions an alternative softsign convention,
$$
\operatorname{softsign}_\tau(x)=\tanh(x/\tau),\qquad \frac{\partial \operatorname{softsign}_\tau}{\partial x}=\frac{1}{\tau}\operatorname{sech}^2(x/\tau),
$$
and an alternative soft absolute value used in the literature,
$$
\operatorname{softabs}_\tau(x)=\tau\log\bigl(2\cosh(x/\tau)\bigr),\qquad \frac{\partial \operatorname{softabs}_\tau}{\partial x}=\tanh(x/\tau),
$$
while emphasizing that SoftJAX’s built-in `abs` uses $\operatorname{sign}_\tau(x)\cdot x$ [2603.08824].

For ReLU, the paper distinguishes two canonical constructions. The integration-style variant is
$$
\operatorname{relu}_\tau(x)=\int_0^x H_\tau(t)\,dt,
$$
which in smooth mode becomes Softplus,
$$
\operatorname{relu}_\tau(x)=\tau\log(1+e^{x/\tau}),\qquad \frac{\partial \operatorname{relu}_\tau}{\partial x}=\sigma(x/\tau).
$$
The gating-style variant is
$$
\operatorname{relu}_\tau(x)=x\cdot H_\tau(x)=x\sigma(x/\tau),
$$
which yields SiLU-like derivatives [2603.08824]. Soft clipping is then composed from soft ReLU:
$$
\operatorname{softclip}_{\tau,a,b}(x)=a+\operatorname{relu}_\tau(x-a)-\operatorname{relu}_\tau(x-b).
$$

The fuzzy-logic layer defines `SoftBool` probabilities in $[0,1]$ [2603.08824]. The default conjunction is the product t-norm,
$$
a\wedge b = a\cdot b,
$$
with disjunction through De Morgan,
$$
a\vee b = 1-(1-a)(1-b)=a+b-a\cdot b,
$$
and negation
$$
\neg a = 1-a
$$
[2603.08824]. Alternative Łukasiewicz and Gödel/min t-norm and t-conorm families are discussed conceptually but are not the defaults [2603.08824].

Soft comparisons use softened Heaviside functions on pairwise differences. For example,
$$
\operatorname{softgt}_\tau(x,y)=\sigma((x-y)/\tau),\qquad
\operatorname{softlt}_\tau(x,y)=\sigma((y-x)/\tau),
$$
with analogous definitions for soft greater-or-equal, less-or-equal, and `isclose` [2603.08824]. In smooth mode,
$$
\frac{\partial \operatorname{softgt}_\tau}{\partial x}
=
\frac{1}{\tau}\sigma((x-y)/\tau)\bigl(1-\sigma((x-y)/\tau)\bigr),\qquad
\frac{\partial \operatorname{softgt}_\tau}{\partial y}
=
-\frac{\partial \operatorname{softgt}_\tau}{\partial x}.
$$

Soft indexing is attention-style [2603.08824]. For values $v_i$ and an index variable $x$, SoftJAX constructs weights on the simplex,
$$
w_i=\operatorname{softmax}(-d(x,i)/\tau),\qquad y=\sum_i w_i v_i,
$$
with distances such as $|x-i|$ or $(x-i)^2$ [2603.08824]. 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 $U(a,b)$ [2603.08824]. The objective minimizes
$$
\langle \Gamma, C(x)\rangle + \epsilon \sum_{ij}\Gamma_{ij}(\log \Gamma_{ij}-1)
$$
subject to $\Gamma \mathbf{1}=a$ and $\Gamma^\top \mathbf{1}=b$, with costs $C_{ij}=(x_i-y_j)^2$ and balanced marginals [2603.08824]. A soft permutation matrix is recovered as
$$
P_\tau := n\Gamma_\tau^\top \in B_n,
$$
from which
$$
\operatorname{argsort}_\tau(x)=P_\tau,\qquad
\operatorname{sort}_\tau(x)=P_\tau x,\qquad
\operatorname{rank}_\tau(x)=P_\tau^\top[n,\dots,1]^\top
$$
[2603.08824].

For entropic regularization, SoftJAX uses Sinkhorn iterations and states the standard scaling steps
$$
u \leftarrow a/(Kv),\qquad v \leftarrow b/(K^\top u),
$$
with $K=\exp(-C/\epsilon)$ [2603.08824]. Entropic regularization yields $\Gamma_\tau>0$ everywhere, dense Jacobians, and a smooth $C^\infty$ map [2603.08824]. Euclidean regularization promotes sparsity and is piecewise smooth; $p$-norm regularization with $1<p\le 2$ promotes sparsity while retaining controlled smoothness, and the library provides `c1` and `c2` settings with regularity guarantees under connected support graphs [2603.08824].

The unit-simplex projection family includes SoftSort and NeuralSort [2603.08824]. SoftSort approximates Sinkhorn through row-wise projection against hard anchors,
$$
\operatorname{argsort}_\tau(x)=\Pi_\tau\bigl(-|\operatorname{sort}(x)\mathbf{1}^\top-\mathbf{1}x^\top|\bigr),
$$
where $\Pi_\tau$ denotes row projection to $\Delta_n$ using entropic, Euclidean, or $p$-norm regularization [2603.08824]. NeuralSort forms argsort distributions by projecting a linear combination of $x$ and sums of soft absolute differences. Its construction uses a matrix $A_\tau(x)$ with entries $(A_\tau)_{ij}=\operatorname{abs}_\tau(x_i-x_j)$ and row formulas depending on the target rank, argmax, argmin, or quantile [2603.08824]. Because SoftJAX uses soft abs, NeuralSort becomes fully smooth in `"smooth"` mode [2603.08824].

The permutahedron family includes FastSoftSort and SmoothSort [2603.08824]. FastSoftSort projects onto the permutahedron under Euclidean or log-KL objectives, solving the resulting isotonic regression problem by pool adjacent violators in $O(n\log n)$ time and $O(n)$ memory [2603.08824]. It returns values rather than permutation matrices, so `argsort` and `argrank` are unavailable [2603.08824]. SmoothSort, introduced in SoftJAX, adds entropic regularization in a dual formulation and replaces hard majorization bounds with smooth log-sum-exp bounds,
$$
\tilde b_k=\tau \log e_k(\exp(z/\tau)),
$$
yielding $C^\infty$ differentiability and dense Jacobians [2603.08824]. The dual is solved with L-BFGS and a custom VJP, with $O(n^2)$ preprocessing for the bounds but without materializing $n\times n$ matrices [2603.08824].

SoftJAX also includes bitonic sorting networks in which compare-and-swap is replaced by soft comparisons $\sigma=H_\tau(a-b)$ [2603.08824]. The associated soft minimum and maximum are
$$
\operatorname{soft\_min}(a,b)=\sigma\cdot b+(1-\sigma)\cdot a,\qquad
\operatorname{soft\_max}(a,b)=\sigma\cdot a+(1-\sigma)\cdot b.
$$
By propagating a soft permutation matrix through the network, the method can produce both sorted outputs and soft `argsort` values [2603.08824].

The paper provides explicit complexity claims. Sinkhorn-based soft sort has $O(Tn^2)$ time and $O(n^2)$ memory, though SoftJAX reduces memory through chunked scans with `lax.scan` and uses implicit differentiation to avoid storing all intermediates [2603.08824]. FastSoftSort has $O(n\log n)$ time and $O(n)$ memory [2603.08824]. Soft indexing operates in $O(n)$ time and $O(n)$ memory and is fully compatible with `vmap` [2603.08824].

## 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 [2603.08824]. Methods are selected through `method=...`, smoothness through `mode=...`, and softness through `softness=...` [2603.08824]. For axiswise operators, `standardize_and_squash` can be enabled to improve numerical stability and scale-independence [2603.08824].

The paper defines standardize-and-squash as
$$
\mu=\operatorname{mean}(x),\qquad s=\operatorname{std}(x),\qquad \tilde x=\sigma((x-\mu)/s),
$$
with inverse transformation for value outputs by logit and destandardization [2603.08824]. 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 [2603.08824].

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 [2603.08824]. Log-domain implementations are recommended for Sinkhorn to avoid underflow and overflow, and gradient clipping is recommended for very small $\tau$ [2603.08824]. For safe wrappers around singular functions such as `sqrt`, `arcsin`, `arccos`, `div`, `log`, and `norm`, gradients are clamped near boundaries [2603.08824].

The guidance for choosing $\tau/\epsilon$ is also explicit [2603.08824]. For axiswise operators after standardize-and-squash, the recommended starting range is approximately $0.1$ to $0.5$, followed by annealing downward during training [2603.08824]. The paper further proposes calibrating softness by the normalized entropy of the soft permutation, $H(P)/\log n$, to compare behavior across methods and problem sizes [2603.08824].

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` [2603.08824]. 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` [2603.08824]. At $n=4096$, the sorting network is reported as the fastest soft method at approximately $1.0$ ms, about $3.8\times$ the hard baseline; SoftSort is approximately $16$ ms; NeuralSort is approximately $37$ ms [2603.08824]. FastSoftSort is identified as the most memory-efficient method and is reported to scale roughly linearly, using approximately $420$ KB at $n=4096$ because it avoids permutation matrices [2603.08824]. OT and SmoothSort are the slowest among the tested methods, due respectively to iterative solves and $O(n^2)$ preprocessing, but they provide full smoothness and dense Jacobians [2603.08824].

The principal application example is a practical case study in MJX collision detection [2603.08824]. 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 [2603.08824]. Applying STE to the outer routine preserves the hard forward pass while using informative gradients in backpropagation [2603.08824].

The paper situates SoftJAX relative to several research lines [2603.08824]. It states that SoftJAX generalizes SoftSort across entropic, Euclidean, and $p$-norm regularizers; adds NeuralSort with fully smooth soft abs, ranking, and sorting networks; and introduces SmoothSort with $C^\infty$ behavior through smooth majorization bounds [2603.08824]. It notes that its simplex projections reproduce softmax under entropic regularization and Sparsemax under Euclidean regularization, while extending to $p$-norm projections with `C^1` and `C^2` guarantees [2603.08824]. 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 [2603.08824]. Among external dependencies, the paper identifies `ott-jax` for Sinkhorn computations and `Optimistix` for L-BFGS with implicit differentiation [2603.08824].

The limitations are presented in operational terms rather than as a single theoretical objection [2603.08824]. Ties and discrete symmetries can produce non-unique soft permutations; the paper recommends entropic regularization for stable gradients or adding jitter [2603.08824]. Very small $\tau$ can produce large gradient magnitudes, and Euclidean or $p$-norm modes can yield sparse Jacobians that slow signal propagation [2603.08824]. 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 [2603.08824]. 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.

Source: https://www.emergentmind.com/topics/softjax