SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

SAI replication review · Referee report

Summary

SoftJAX and SoftTorch are companion libraries that provide differentiable, drop-in replacements for a large family of "hard" primitives in JAX and PyTorch — comparisons, rounding, indexing, sorting, ranking, quantiles, top- $k$ , argmax, and the associated index-selection helpers — together with straight-through estimation utilities. The conceptual move is to unify a fragmented ecosystem of soft-relaxation tricks under a single API organised as a taxonomy of methods (OT, SoftSort, NeuralSort, FastSoftSort, SmoothSort, sorting networks) crossed with smoothness modes (smooth giving $C^{\infty}$ ; c0/c1/c2 giving the corresponding piecewise-polynomial regularity via $p$ -norm regularisers). The libraries expose "softness" as a single knob $\tau$ that consistently recovers the hard operator as $\tau \to 0^{+}$ . This is a valuable infrastructure contribution: the soft-relaxation literature is scattered across a decade of papers whose implementations are often incompatible, and a well-tested, feature-uniform library lowers the barrier for practitioners in robotics, differentiable simulation, structured prediction, and scientific applications who currently reimplement these operators from scratch. Alongside the consolidation, the paper introduces genuine new content: a clear articulation of an "STE pitfall" in multiplicative compositions and its composite-STE remedy; a novel SmoothSort variant that adds entropic regularisation to the dual permutahedron LP to obtain $C^{\infty}$ smoothness; systematic extensions of NeuralSort and sorting networks to c0/c1/c2 modes; and Theorem 1, which pins down the $C^{k}$ regularity of $p$ -norm regularised simplex and OT projections. An executed reproduction confirms that the library itself is solid: every runnable behaviour claim (soft-vs-hard gradients in the top- $k$ and rule-classifier examples, the sj.argmax simplex output, softsort's $\tau \to \infty$ limit, the $10\tau$ heaviside transition width, Theorem 1's smoothness classes, the chunked $O(n)$ memory / $O(1)$ compile behaviour, and the MJX collision case study's informative-gradient contrast) reproduces exactly on the shipped code with no source patches. The limitations sit around the empirical and demonstrative material: the two MNIST tables have no training pipeline in either repository and could not be attempted; the GPU-specific benchmark numbers ( $1.0$ ms, $\sim 3.8\times$ , SoftSort $16$ ms, NeuralSort $37$ ms) could not be audited on the reproduction hardware; and the paper contains several small internal inconsistencies (notation, transition-width wording, FastSoftSort capabilities, rasteriser tensor shapes) that leave the "principled unification" narrative less thoroughly demonstrated than the API surface suggests.

Strengths

Consolidation of a fragmented toolbox. The paper knits together OT-based, simplex-projection-based, permutahedron-projection-based, and sorting-network approaches under one API with a shared softness knob and mode family. The Table 2 matrix (methods × modes × operators) makes the coverage explicit and honestly labels which cells are prior work vs. derived.
Novel unifying content. The "STE pitfall" section and its composite-STE remedy is a useful pedagogical point; the SmoothSort variant is a legitimately new $C^{\infty}$ permutahedron-projection method (dual entropic + log-sum-exp bound smoothing); Theorem 1 gives a clean $C^{k}$ characterisation that ties c0/c1/c2 mode names to $p$ -norm regularisers, including the sharpness of the exponent. An independent finite-difference probe over sj.sigmoidal at $n=2$ reproduces exactly the $C^{0}$ / $C^{1}$ / $C^{2}$ / $C^{\infty}$ ordering asserted by the theorem, and the shipped test_gradients.py tests exercising all four modes pass with no source modifications.
Systematic mode extensions. NeuralSort and sorting networks are re-derived with c0/c1/c2 regularisers not present in the originals. The reduction of $n=2$ axiswise operators back to the elementwise sigmoidals (Section D.3) is a nice consistency check that ties the elementwise and axiswise pillars of the paper together.
Extensive benchmarking. Figures 8 and 16–23 sweep runtime, JIT-compile time, and peak memory across sizes, methods, and modes; Section F distils this into practical guidance (choose NeuralSort/Sorting Networks by default; FastSoftSort for memory; smooth for gradient conditioning). The measurement harness itself is bundled and, under the executed reproduction, regenerates Figure 8's two log-log panels and all seven curves from benchmarks/plot_paper.py --figure main.
Two-library parity. Shipping both a JAX and a PyTorch implementation with a documented naming map (Table 1) meaningfully broadens the audience.
Reproducible library core. Small self-contained examples (quick_example.py, the manifold-points notebook, the paper examples script) run end-to-end from the released repositories and reproduce the printed outputs of the feature-selection and rule-classifier code blocks byte-for-byte, and the informative-gradient contrast in the MJX collision case study. This is a strong positive signal for the core library implementation.
Practical case studies. The MJX collision subroutine, the top- $k$ feature selector, the rule-based classifier, and the mesh renderer sketch how the primitives compose into real pipelines and are all short enough to read end-to-end.

Weaknesses

Empirical downstream evidence is thin. Tables 3 and 4 report single-run percentages/MSE with no seeds, no variance, no confidence intervals, and no direct comparison to the published NeuralSort/SoftSort/FastSoftSort/Sorting Network numbers on the same benchmark. Several table entries are dashed (–) without explanation (Sorting Network at $n=32$ , OT for $n \in \{3,5,7\}$ ), which reads as "did not fit / did not run" rather than a design choice.
Notational drift around soft permutation matrices. Section 4.4 tracks a soft, doubly-stochastic $P$ through the sorting network but then writes the exact equality $\operatorname{sort}(\mathbf{x}) = P\mathbf{x}$ with the hard sort operator. Elsewhere the paper carefully distinguishes $P^{\star}$ (hard) from $P_{\tau}$ (soft), so this passage is inconsistent with the paper's own conventions and could mislead a reader into thinking the sorting-network surrogate is exact.
Elementwise transition-width claim conflicts with the piecewise-support description. Section F states "all modes share an effective transition width of approximately $10\tau$ ", but the definitions in Section 3 give c0/c1/c2 modes exactly-compact support outside a $\pm 5\tau$ window, while the smooth mode still has strictly non-zero gradient there. The wording invites users to set $\tau$ uniformly across modes without realising the tail behaviours differ substantively.
FastSoftSort capability description is internally inconsistent. Section 4.3 says "this approach does not allow for an argsort or argrank operator", yet Section F says FastSoftSort "produces sorted values and ranks". A reader who wants a soft rank cannot tell from these two passages whether FastSoftSort supports it, and Table 2's SmoothSort rank entry inherits the same ambiguity.
"Feature-complete" is overstated. Section 4 openly excludes stochastic relaxations (Gumbel-softmax, perturbed optimizers, black-box differentiation), and the exclusion is reasonable, but framing the library as "feature-complete" in the abstract and Conclusion overreaches given this whole family of established techniques is out of scope.
Rendering appendix has under-specified methodology. In Section B.3 the inside-triangle probability equation and the soft-argmax equation appear to render blank in the manuscript; the array-shape convention for $P$ oscillates between $N \times H \times W$ and $H \times W \times N$ ; the "inverse depth" transform is claimed to send nearest to $1$ and farthest to $0$ without stating the actual map; and $b \exp(a) = \exp(a + \log b)$ is invoked without addressing $b = 0$ , which is exactly the boundary the soft-Bool $H_{\tau_1}$ approaches for pixels outside a triangle. Any of these individually is a nit; taken together, the pipeline cannot be reproduced from the paper.
Small terminological inaccuracy in the renderer derivation. The paper describes $b_a$ as "corresponds to the half signed area spanned by $c-b$ and $l_i-b$ ", but a barycentric coordinate is a ratio to the full triangle area, not the sub-area itself. This does not affect the inside/outside sign test but does matter for the depth/colour interpolation that later relies on the same coordinates.
Feature-selection recovery claim overstates the printed evidence. Section B.4 claims the smooth variant "recovers the correct features in just a single step and preserves their ranking", but the printed one-step Soft grad assigns the true feature 4 (index 3, weight $-1.5$ ) a magnitude of $1126$ while every uninformative feature receives an identical magnitude of $2268$ . The [8, 2, 4] selection is reported after five gradient steps, not one; only two of three true features are unambiguously separated by magnitude at step 1. The executed reproduction confirms the printed numbers, so the issue is the narrative, not the code.
"Many initial positions did not converge" turns the renderer case study into anecdote. Section B.3 hedges Figure 14's convergence with a note that many initialisations failed. Without a success rate over random initialisations or a small softness / step-size ablation, the figure documents a single successful run rather than a benchmark.
Polygon-size narrative jump in the MJX case study. Section B.1 introduces a seven-sided polygon (Figure 10) and then switches to a five-vertex example (Figure 11) whose text reads as a general claim about the algorithm rather than a distinct example. The "gradient null for one polygon point" argument is specific to $n = 5$ ; the seven-sided case would leave three vertices with null gradients.
STE-pitfall novelty needs a small hedge. The claim "To our knowledge, this 'STE pitfall' and its remedy have not been explicitly discussed before" is plausible but the phenomenon is folklore in the discrete-VAE / STE literature; a brief pointer to related discussions of composite STE would strengthen the claim.
Axiswise softness convention leaves cross-mode calibration to the user. Section D.3 acknowledges that the same $\tau$ produces different effective softness across modes and sizes and recommends users measure $H(\mathbf{P})/\log n$ themselves. Given how central mode selection is to the library, a shipped helper for this calibration would close a practical gap.
Memory reduction for the SmoothSort bounds is stated without the recomputation cost. Section D.6 cites optimal online checkpointing to reduce memory from $O(n^{2})$ to $O(n\sqrt{n})$ for the smoothed bounds, but does not report the resulting recomputation overhead; the same paragraph states runtime is $O(n^{2})$ without clarifying whether that is pre- or post-checkpointing.
Section 4.3 packaging obscures SmoothSort's asymptotic. SmoothSort is introduced in the same subsection that motivates permutahedron methods as $O(n \log n)$ via PAV, but it itself carries $O(n^{2})$ preprocessing plus an L-BFGS solve. That is stated later, but the top-level framing of Section 4.3 invites readers to pick SmoothSort under the wrong asymptotic expectations.
Underspecified all-as-product independence assumption. The default all operator is defined as $\prod_i p_i$ and glossed as "the joint probability that independent Boolean events $p_1, \dots, p_n$ equate to True". The SoftBools arising in the paper's own case studies (three barycentric coordinates of the same pixel, comparisons over the same feature) are structurally correlated, so the product is a coordinate-wise composition, not a genuine joint probability. Either flag the independence assumption or drop the "joint probability" framing.
Table 2 attribution for c1/c2 OT is opaque. Table 2 assigns c1/c2 OT regularization to Paty & Cuturi (2020), but the text describes only entropic/Euclidean/p-norm OT and never spells out which regularizer + exponent choice yields c1/c2 in the OT setting, so the tie to Theorem 1's exponent bookkeeping is left implicit.

Reproducibility & code

The executed reproduction ran the shipped code with zero source-code fixes required and confirmed most of the paper's runnable claims; the two consequential gaps are the MNIST tables and the exact GPU benchmark numbers.

Library-behaviour claims all reproduced exactly. The quick_example.py top- $k$ feature-selection block emitted the same hard-grad zeros, the same soft grad values ([2268.416, -2371.1378, 2268.416, 1126.1998, ...]), and the same 5-step selection [8 2 4] as printed in the paper. The rule-classifier block emitted identical Soft grad: [-4.2777 1.4152] and Learned [l, h]: [0.2925 0.5999], with training-set accuracy hitting $1.0$ within $\sim 10$ steps as shown in Figure 15(c). The sj.argmax example reproduced [0.004 0.042 0.953] and its expected soft-indexed value $\approx 0.78$ . A $\tau$ -sweep of softsort on a length-6 vector shows monotonic convergence of both the outputs to the mean and the Jacobian entries to $1/n$ , matching Figure 7.
Theorem 1 smoothness classes are directly checkable and hold. The shipped tests/test_gradients.py smoothness tests pass (48/48 in the relevant subset). An independent finite-difference probe of sj.sigmoidal near saturation reproduces the ordering: c0 has a first-derivative jump ( $\sim 10^{-1}$ ), c1 a second-derivative jump, c2 a third-derivative jump, and smooth is smooth at every order.
Chunked memory / flat JIT behaviour matches the description. projections_permutahedron.py and utils.py._map_in_chunks implement the lax.scan-based chunking described in Section 4.3 and are unmodified; the JIT-compile times in the reproduced CSV move only $\sim 1.5$ – $2.6\times$ over a $16\times$ array-size increase (roughly flat), and peak memory of the chunked methods scales near-linearly, while the non-chunked OT baseline scales quadratically. The FastSoftSort $\sim 420$ kB @ $n = 4096$ number extrapolates cleanly from the reproduced $n \leq 1024$ data (log-log slope $\approx 0.95$ , extrapolated $\sim 435$ – $454$ kB, within a few percent).
MJX collision subroutine reproduces qualitatively. The shipped docs/examples/manifold_points.ipynb runs end-to-end, and an added script that reuses the notebook's own functions confirms Figure 11's central claim: the hard version has exactly zero gradient at 12/16 non-selected polygon vertices, whereas the soft version has non-zero gradients at all 16. Two caveats: (i) the notebook default polygon is 8-sided, not the 7- or 5-sided geometries used by Figures 10 and 11; and (ii) the explicit hard-vs-soft gradient comparison plotted in Figure 11 (right) is not itself a notebook cell, so the exact figure cannot be regenerated from what is shipped.
Figure 8 benchmark reproduces qualitatively; absolute GPU numbers are not verifiable on the reproduction hardware. benchmarks/benchmark.py and benchmarks/plot_paper.py --figure main regenerate the two log-log panels of Figure 8 with all seven curves and the same axis conventions. Under the executed reproduction (shared A100 was contended and forced a CPU run at sizes $2^6 \dots 2^{10}$ ), the qualitative core holds: SortingNet is the fastest soft method at the largest size reached ( $\approx 3.28$ ms vs $\approx 3.77$ / $3.92$ ms for NeuralSort / SoftSort, with a $\sim 2.44\times$ ratio to the hard baseline). The paper's absolute GPU-specific numbers ( $1.0$ ms at $n = 4096$ , $\sim 3.8\times$ hard baseline, SoftSort $16$ ms, NeuralSort $37$ ms) cannot be audited without matching hardware; no reference CSV is bundled.
"SmoothSort and OT are the two slowest" is only half-supported by the harness. Under the same shipped plot_paper.py metric (steady-state forward+backward time), SmoothSort is unambiguously the slowest at every reproduced size, but FastSoftSort emerges as the second slowest ( $\approx 2650$ ms at $n=1024$ ), leaving OT in third place ( $\approx 28$ ms). This may well be a CPU artefact of FastSoftSort's steady-state timing that would reorder on the paper's GPU, but as stated the qualitative ranking is not confirmed by the reproduction on the shipped harness.
MNIST sorting/quantile tables (Tables 3, 4) are not reproducible from the shipped code. A search of both repositories for MNIST data loaders, argsort + cross-entropy training loops, softness schedules, or a Petersen-et-al training port returned only a citation inside a docstring — there is no training script, no data loader, no seed information. Reproducing Tables 3 and 4 requires reimplementing the full Appendix B.2 pipeline from prose (a $100$ k-step Adam training loop for each mode / method / $n$ configuration).
Differentiable renderer (Section B.3, Figure 14) is not shipped. The rasteriser code, the sphere and "Dave" meshes, the initialisation policies, and the AdamW loop are all described textually only. Combined with the "many initial positions did not converge" caveat, this makes Figure 14's convergence traces effectively unauditable.
Code-availability sentence points at live repositories rather than a paper-associated release. Both repositories exist and ran cleanly, but the sentence does not pin a version tag or archived DOI. Given the reproducibility gaps above, an explicit release aligned with the ICML 2026 version would let readers audit the exact commit that produced the paper's numbers.

Recommended Changes

Essential

Add seeds, error bars, and cross-method calibration for Tables 3 and 4. Run each MNIST sorting/quantile configuration over multiple seeds, report mean $\pm$ standard deviation (or per-seed points), and add side-by-side numbers for the same $n$ from the original NeuralSort/SoftSort/FastSoftSort/Sorting Network papers so a reader can see whether SoftJAX's implementations match published performance. Also state why the dashed cells were not run. This maps to the "Empirical downstream evidence is thin" weakness.
Ship the MNIST training pipeline. Release the JAX/Equinox port of the Petersen et al. (2022a) code that produced Tables 3 and 4, including data loader, softness schedule, and seeds. This maps to the "MNIST sorting/quantile tables are not reproducible from the shipped code" bullet under Reproducibility & code.
Ship the differentiable renderer or downgrade the case study. Release the rasteriser code, mesh assets, initialisation policy, and AdamW loop used for Figure 14, or clearly label the case study as illustrative and drop the specific convergence traces. This maps to the "Differentiable renderer is not shipped" bullet and to the "many initial positions did not converge" weakness.
Reconcile the notational drift in Section 4.4. Replace $\operatorname{sort}(\mathbf{x}) = P\mathbf{x}$ with $\operatorname{sort}_{\tau}(\mathbf{x}) = P_{\tau}\mathbf{x}$ and denote the tracked matrix $P_{\tau}$ so it matches the paper's own $P^{\star}$ / $P_{\tau}$ convention. Add the explicit row-mixing formula so the new c0/c1/c2 sorting-network variants are self-contained.
Fix the " $10\tau$ transition width" claim in Section F. Either qualify it to the transition width of the smooth mode, or explicitly point out that c1/c2 modes remain compact-support inside this window while smooth does not — so users understand what the shared width does and does not mean.
Reconcile the FastSoftSort rank-vs-argrank statements between Section 4.3 and Section F. State once and consistently whether FastSoftSort supports rank (value-space soft rank via permutahedron projection) and, if so, distinguish that from soft argrank / soft permutation matrices; likewise clarify SmoothSort's rank entry in Table 2.
Ship a reference benchmark CSV and a hardware manifest, and re-verify the "SmoothSort and OT are the two slowest" ranking on the paper's GPU. So the " $1.0$ ms / $\sim 3.8\times$ hard baseline" claim and the two-slowest ordering in Section 5 are auditable; the reproduction on the shipped harness could confirm the SortingNet-fastest ordering but not the OT vs. FastSoftSort tail ranking.

Suggested

Fix the rendering appendix (Section B.3). Make Equations (53) and (54) explicit (fuzzy-AND aggregation of $H_{\tau_1}(b_a)$ , $H_{\tau_1}(b_b)$ , $H_{\tau_1}(b_c)$ if that is what is meant); reconcile the $P$ tensor shape between rasterisation and z-buffering; state the exact "inverse depth" transform (per-pixel min–max, or the actual formula); and specify the clamp / floor that guards $b \exp(a) = \exp(a + \log b)$ when $b \to 0$ .
Correct the barycentric-coordinate description so that $b_a$ is the sub-triangle area divided by the full triangle area, not the raw signed area.
Clarify the MJX polygon-size narrative. Explicitly mark Figure 10 (seven-sided) and Figure 11 (five-vertex) as distinct examples, and reword "always chooses four points out of the polygon's five vertices" so it does not read as a general statement about the algorithm.
Extend the shipped manifold-points notebook to include the exact 7- and 5-sided polygon geometries and the Figure 11 gradient-perturbation panel, so the case-study figures can be regenerated from a single notebook.
Tone down the one-step feature-selection claim. Report the actual step at which [8, 2, 4] first appears in the top- $k$ selection, or restate what "recovers in a single step" means (e.g. correct set rather than correct magnitude ranking).
Qualify the "feature-complete" framing and the STE-pitfall novelty claim. Explicitly exclude stochastic relaxations from the "feature-complete" scope in the abstract/Conclusion, and add one or two references acknowledging prior awareness of STE-in-products behaviour.
Provide a softness-calibration helper. Since Section D.3 recommends users compute $H(\mathbf{P})/\log n$ themselves for axiswise operators, expose this as a library utility so cross-mode comparisons are practical.
Clarify the SmoothSort asymptotic upfront in Section 4.3 so readers do not confuse it with the $O(n \log n)$ FastSoftSort family, and state whether the $O(n^{2})$ runtime for the smoothed bounds is pre- or post-checkpointing.
State the independence assumption behind all = product (or drop the "joint probability" framing), and give a default and rule-of-thumb for the $K$ parameter of the soft round definition so Section 3.1 is self-contained.
Introduce the gated_grad flag at first use in the top- $k$ feature-selection code block, not only in Appendix D.5, so the example is self-contained.
Spell out Table 2's c1/c2 OT attribution by naming the regularizer + exponent choice that yields c1/c2 in the OT setting and tying it to Theorem 1's $q = p/(p-1)$ bookkeeping.
Tag the release corresponding to the ICML 2026 version of the paper in both repositories (and cite the tag/DOI in the code-availability statement) so reviewers and users can audit the exact commit that produced the paper's numbers.