SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

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

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

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., $\operatorname{rank}_{\tau}$ instead of $\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 $\mathcal{C}^{\infty}$ or c1 for $\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)=\operatorname{relu}(x)

and

g(x,y)=\operatorname{relu}(y)

. Applying straight-through estimation on each function individually causes the gradient

\nabla(f_{\text{STE}}\cdot g_{\text{STE}})

to be zero if

x<0

, whereas the gradient

\nabla(f\cdot g)_{\text{STE}}

is non-zero.

Figure 4: Top: The heaviside function surrogates

H_{\tau}

are used to derive the

\operatorname{sign}

\operatorname{round}

\operatorname{abs}

, and

\operatorname{clip}

functions. Bottom: Comparison operations compare values

x,y\in\mathbb{R}

by interpreting

H_{\tau}

as a CDF defining the probability that

x>0

Figure 6: Illustration of the Softsort algorithm.

Figure 7: Top: Output of the

\operatorname{sort}

operator using “softsort” for an array of

n=6

values as a function of softness

\tau

. Elements converge to the mean value as

\tau\to\infty

. Bottom: Partial derivatives

{\partial{y_{j}}}/{\partial{x_{i}}}

. Here,

j

denotes the sorted position of input element

x_{i}

. As

\tau\to\infty

, the derivatives converge to

1/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

SoftJAX	SoftTorch
clip	clamp
equal	eq
top_k	topk
take_along_axis	take_along_dim
dynamic_index_in_dim	index_select
dynamic_slice_in_dim	narrow

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=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

\tau=0.1

, SoftJAX selects points similar to MJX. Right: Perturbing the x-position of an individual vertex

v_{i}

\Delta x_{i}

changes the mean over all coordinates of the selected points. Notably, in the mode smooth with

\tau=0.1

and

\tau=1

, the SoftJAX variant yields smooth gradients of the mean with respect to

\Delta x_{i}

. By contrast, MJX yields non-zero gradients solely for the vertices that have currently been selected by the algorithm.

Straight-Through Estimator

See all 11 papers

Oct 8, 2026

asdex: Automatic Sparse Differentiation in JAX

Many tasks in scientific computing and machine learning require the Jacobian or Hessian matrix of a function. Automatic differentiation (AD) computes these derivatives to machine precision, but materializing a dense

m \times n

Jacobian requires

n

forward-mode or

m

reverse-mode AD passes, one per column or row. For a large class of functions, each output depends on only a few inputs, making the derivative matrix sparse. Automatic sparse differentiation (ASD) exploits this structure in four steps: detection of the input-agnostic sparsity pattern, coloring of a graph to group columns or rows that can share an AD pass, compressed differentiation to compute a compressed derivative matrix with one AD pass per color, and finally decompression into the original sparsity pattern. The number of colors, and hence of AD passes, is often independent of the problem dimension: a banded Jacobian with

b

contiguous bands, for instance, only ever requires

b

colors, regardless of its size. asdex offers the first standalone ASD toolkit in the popular JAX ecosystem. With

\texttt{asdex.jacobian}

and

\texttt{asdex.hessian}

, it provides sparse drop-in replacements for

\texttt{jax.jacobian}

and

\texttt{jax.hessian}

Adrian Hill, Guillaume Dalle

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

Automatic Differentiation Structured Sparsity

Ask PDF

Feb 26, 2026

Differentiable Zero-One Loss via Hypersimplex Projections

Recent advances in machine learning have emphasized the integration of structured optimization components into end-to-end differentiable models, enabling richer inductive biases and tighter alignment with task-specific objectives. In this work, we introduce a novel differentiable approximation to the zero-one loss-long considered the gold standard for classification performance, yet incompatible with gradient-based optimization due to its non-differentiability. Our method constructs a smooth, order-preserving projection onto the n,k-dimensional hypersimplex through a constrained optimization framework, leading to a new operator we term Soft-Binary-Argmax. After deriving its mathematical properties, we show how its Jacobian can be efficiently computed and integrated into binary and multiclass learning systems. Empirically, our approach achieves significant improvements in generalization under large-batch training by imposing geometric consistency constraints on the output logits, thereby narrowing the performance gap traditionally observed in large-batch training.

Camilo Gomez, Pengyang Wang, Liansheng Tang

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

Classification Deep Learning Optimization Differentiable Optimization

Ask PDF

Feb 19, 2026

JAX-Privacy: A library for differentially private machine learning

JAX-Privacy is a library designed to simplify the deployment of robust and performant mechanisms for differentially private machine learning. Guided by design principles of usability, flexibility, and efficiency, JAX-Privacy serves both researchers requiring deep customization and practitioners who want a more out-of-the-box experience. The library provides verified, modular primitives for critical components for all aspects of the mechanism design including batch selection, gradient clipping, noise addition, accounting, and auditing, and brings together a large body of recent research on differentially private ML.

Ryan McKenna, Galen Andrew, Borja Balle +6

Privacy-Preserving ML Differential Privacy Differentially Private Stochastic Gradient Descent

Ask PDF

Abstract

Straight-Through Estimator

Explore similar work

asdex: Automatic Sparse Differentiation in JAX

Differentiable Zero-One Loss via Hypersimplex Projections

JAX-Privacy: A library for differentially private machine learning