SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

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.

Modern machine learning models are trained with gradients, small numerical signals that tell each part of a model how to improve. But many basic operations that programs rely on, such as sorting a list, picking the largest value, or comparing two numbers, produce zero or undefined gradients and block learning. For decades, researchers have invented "soft" workarounds for these operations one paper at a time, leaving the techniques scattered across the literature and hard to combine. SoftJAX and SoftTorch are open-source libraries that bring these soft, gradient-friendly versions of common operations together as drop-in replacements in the two most widely used machine learning frameworks. By unifying years of fragmented work into a single, easy-to-use package, they lower the barrier to applying gradient-based optimization in domains like robotics, scientific simulation, and structured prediction, where discrete decisions have traditionally required specialized expertise to handle.