# SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

arXiv: https://arxiv.org/abs/2603.08824 · HTML: https://mlsanity.com/paper/2603.08824/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 59 references that are in the map.
- First posted on arXiv 9 Mar 2026.

## Builds on (most cited first, 30 of 38)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [3D Gaussian Splatting for Real-Time Radiance Field Rendering](https://mlsanity.com/paper/2308.04079/) (2023)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification](https://mlsanity.com/paper/1602.02068/) (2016)
- [Differentiable Surface Splatting for Point-based Geometry Processing](https://mlsanity.com/paper/1906.04173/) (2019)
- [Differentiation of Blackbox Combinatorial Solvers](https://mlsanity.com/paper/1912.02175/) (2019)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Fast Differentiable Sorting and Ranking](https://mlsanity.com/paper/2002.08871/) (2020)
- [Differentiable MPC for End-to-end Planning and Control](https://mlsanity.com/paper/1810.13400/) (2018)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions](https://mlsanity.com/paper/2106.01798/) (2021)
- [Smooth and Sparse Optimal Transport](https://mlsanity.com/paper/1710.06276/) (2017)
- [Learning with Differentiable Perturbed Optimizers](https://mlsanity.com/paper/2002.08676/) (2020)
- [Deep Graph Matching via Blackbox Differentiation of Combinatorial Solvers](https://mlsanity.com/paper/2003.11657/) (2020)
- [Projection onto the probability simplex: An efficient algorithm with a simple proof, and an application](https://mlsanity.com/paper/1309.1541/) (2013)
- [Gradient Estimation with Stochastic Softmax Tricks](https://mlsanity.com/paper/2006.08063/) (2020)
- [SoftSort: A Continuous Relaxation for the argsort Operator](https://mlsanity.com/paper/2006.16038/) (2020)
- [Optimizing Rank-based Metrics with Blackbox Differentiation](https://mlsanity.com/paper/1912.03500/) (2019)
- [Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision](https://mlsanity.com/paper/2105.04019/) (2021)
- [Backpropagation through Combinatorial Algorithms: Identity with Projection Works](https://mlsanity.com/paper/2205.15213/) (2022)
- [The Differentiable Cross-Entropy Method](https://mlsanity.com/paper/1909.12830/) (2019)
- [Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective](https://mlsanity.com/paper/2302.01425/) (2023)
- [Optimistix: modular optimisation in JAX and Equinox](https://mlsanity.com/paper/2402.09983/) (2024)

## Built on by (most cited first, 2 of 2)

- [Emulating Non-Differentiable Metrics via Knowledge-Guided Learning: Introducing the Minkowski Image Loss](https://mlsanity.com/paper/2604.11422/) (2026)
- [Critical evaluation of PINN for FWD inverse analysis and differentiable FEM as an alternative](https://mlsanity.com/paper/2606.03210/) (2026)

## Abstract (arXiv)

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.