cs.LGMar 9, 2026

SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

Authors: Anselm Paulus, A. René Geist, Vít Musil, Sebastian Hoffmann, Onur Beker, Georg Martius

Organizations: University of Tübingen, Germany · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Masaryk University, Czechia · Max Planck Institute for Biogeochemistry, Germany

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Automatic differentiation (AD) frameworks such as JAX and PyTorch have enabled gradient-based optimization for a wide range of scientific fields. Yet, many "hard" primitives in these libraries such as thresholding, Boolean logic, discrete indexing, and sorting operations yield zero or undefined gradients that are not useful for optimization. While numerous "soft" relaxations have been proposed that provide informative gradients, the respective implementations are fragmented across projects, making them difficult to combine and compare. This work introduces SoftJAX and SoftTorch, open-source, feature-complete libraries for soft differentiable programming. These libraries provide a variety of soft functions as drop-in replacements for their hard JAX and PyTorch counterparts. This includes (i) elementwise operators such as clip or abs, (ii) utility methods for manipulating Booleans and indices via fuzzy logic, (iii) axiswise operators such as sort or rank -- based on optimal transport or permutahedron projections, and (iv) offer full support for straight-through gradient estimation. Overall, SoftJAX and SoftTorch make the toolbox of soft relaxations easily accessible to differentiable programming, as demonstrated through benchmarking and a practical case study. Code is available at github.com/a-paulus/softjax and github.com/a-paulus/softtorch.

Figures & tables

Figure 1: Top: SoftJAX and SoftTorch provide differentiable surrogates for discrete operations i. e., rank⁡τ\operatorname{rank}_{\tau} instead of rank⁡\operatorname{rank} . Numerous soft approximations can be obtained by tuning the softness parameter τ\tau , selecting the softening method (e. g., “neuralsort” or “softsort”), and choosing a smoothness mode (e. g., smooth for C∞\mathcal{C}^{\infty} or c1 for C1\mathcal{C}^{1} ). Bottom: The soft surrogates shown above are instantiated using just three lines of code.
Figure 2: The function relu_st resorts to the straight-through trick to use the hard function jax.nn.relu in the forward pass while using the soft function sj.relu for gradient computation.
Figure 3: Arrows depict the normalized gradient of the product of two functions f(x,y)=relu⁡(x)f(x,y)=\operatorname{relu}(x) and g(x,y)=relu⁡(y)g(x,y)=\operatorname{relu}(y) . Applying straight-through estimation on each function individually causes the gradient ∇(fSTE⋅gSTE)\nabla(f_{\text{STE}}\cdot g_{\text{STE}}) to be zero if x<0x<0 , whereas the gradient ∇(f⋅g)STE\nabla(f\cdot g)_{\text{STE}} is non-zero.
Figure 4: Top: The heaviside function surrogates HτH_{\tau} are used to derive the sign⁡\operatorname{sign} , round⁡\operatorname{round} , abs⁡\operatorname{abs} , and clip⁡\operatorname{clip} functions. Bottom: Comparison operations compare values x,y∈Rx,y\in\mathbb{R} by interpreting HτH_{\tau} as a CDF defining the probability that x>0x>0 .
Figure 5: OT-based \argtopk\argtopk computes a (scaled) optimal transport plan Pτ⋆P^{\star}_{\tau} whose entries Pij⋆P_{ij}^{\star} contain the probability mass transported from the jj -th entry of x\mathbf{x} to the ii -th entry of the anchor y\mathbf{y} .
Figure 6: Illustration of the Softsort algorithm.
Figure 7: Top: Output of the sort⁡\operatorname{sort} operator using “softsort” for an array of n=6n=6 values as a function of softness τ\tau . Elements converge to the mean value as τ→∞\tau\to\infty . Bottom: Partial derivatives ∂yj/∂xi{\partial{y_{j}}}/{\partial{x_{i}}} . Here, jj denotes the sorted position of input element xix_{i} . As τ→∞\tau\to\infty , the derivatives converge to 1/n1/n .
Figure 8: Computation time ( left ) and peak memory consumption ( right ) for a full forward-backward pass of sj.sort using smooth regularization. All six soft sorting methods are compared against the hard baseline. Extended results across all modes are in Figure 14 .
Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

SoftJAXSoftTorch
clipclamp
equaleq
top_ktopk
take_along_axistake_along_dim
dynamic_index_in_dimindex_select
dynamic_slice_in_dimnarrow
Appendix
Table 1: Naming differences between SoftJAX and SoftTorch. SoftTorch follows PyTorch conventions. Functions listed as “—” are not available in SoftTorch (no PyTorch equivalent).
Figure 10: Manifold points computed by the algorithm’s SoftJAX reimplementation over different softness values τ\tau . While a small softness causes the algorithm to choose points close to the vertices of the polygon, for large softness values, all four points collapse into a single point.
Figure 11: Left: The manifold_points_mjx algorithm selects the four vertices of a polygon (here n=5n=5 ) that maximize the enclosed area. In contrast, manifold_points_softjax selects the four points by computing an expectation using probability distributions over vertex indices. For small softness τ=0.1\tau=0.1 , SoftJAX selects points similar to MJX. Right: Perturbing the x-position of an individual vertex viv_{i} by Δxi\Delta x_{i} changes the mean over all coordinates of the selected points. Notably, in the mode smooth with τ=0.1\tau=0.1 and τ=1\tau=1 , the SoftJAX variant yields smooth gradients of the mean with respect to Δxi\Delta x_{i} . By contrast, MJX yields non-zero gradients solely for the vertices that have currently been selected by the algorithm.
More in this topic

Straight-Through Estimator

See all 11 papers

Explore similar work

CardsList

  1. asdex: Automatic Sparse Differentiation in JAX

    Oct 8, 2026Adrian Hill, Guillaume DalleAutomatic DifferentiationStructured Sparsity

    BIFOLD – Berlin Institute for the Foundations of Learning and Data, Berlin, Germany · Machine Learning Group, Technical University of Berlin, Berlin, Germany · LVMT, ENPC, Institut Polytechnique de Paris, Univ Gustave Eiffel, Marne-la-Vallée, France

  2. Differentiable Zero-One Loss via Hypersimplex Projections

    Feb 26, 2026Camilo Gomez, Pengyang Wang, Liansheng TangClassificationDeep Learning Optimization

    School of Data, Mathematical, and Statistical Sciences, University of Central Florida, Orlando, USA · Department of CIS, University of Macau, Macao, China

  3. JAX-Privacy: A library for differentially private machine learning

    Feb 19, 2026Ryan McKenna, Galen Andrew, Borja Balle +6Privacy-Preserving MLDifferential Privacy