mlsanityData to 14 Oct 2026

SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

arXiv 2603.08824 ↗First posted 9 Mar 2026cs.LGcited by 2

Not on the map. Why?

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.

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

38 of 59 references.

References · 38

  1. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
  2. 20233D Gaussian Splatting for Real-Time Radiance Field Renderingcited by 6,226
  3. 2016Categorical Reparameterization with Gumbel-Softmaxcited by 2,739
  4. 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,316
  5. 2013Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computationcited by 2,079
  6. 2016The Concrete Distribution: A Continuous Relaxation of Discrete Random Variablescited by 1,148
  7. 2017OptNet: Differentiable Optimization as a Layer in Neural Networkscited by 599
  8. 2019Differentiable Convex Optimization Layerscited by 465
  9. 2016From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classificationcited by 341
  10. 2019Differentiable Surface Splatting for Point-based Geometry Processingcited by 239
  11. 2019Differentiation of Blackbox Combinatorial Solverscited by 226
  12. 2021Efficient and Modular Implicit Differentiationcited by 203
26 more
  1. 2020Fast Differentiable Sorting and Rankingcited by 197
  2. 2018Differentiable MPC for End-to-end Planning and Controlcited by 189
  3. 2019Differentiable Ranks and Sorting using Optimal Transportcited by 113
  4. 2019Stochastic Optimization of Sorting Networks via Continuous Relaxationscited by 111
  5. 2022Optimal Transport Tools (OTT): A JAX Toolbox for all things Wassersteincited by 101
  6. 2021Implicit MLE: Backpropagating Through Discrete Exponential Family Distributionscited by 98
  7. 2017Smooth and Sparse Optimal Transportcited by 83
  8. 2020Learning with Differentiable Perturbed Optimizerscited by 81
  9. 2020Deep Graph Matching via Blackbox Differentiation of Combinatorial Solverscited by 81
  10. 2013Projection onto the probability simplex: An efficient algorithm with a simple proof, and an applicationcited by 67
  11. 2020Gradient Estimation with Stochastic Softmax Trickscited by 67
  12. 2020SoftSort: A Continuous Relaxation for the argsort Operatorcited by 66
  13. 2019Optimizing Rank-based Metrics with Blackbox Differentiationcited by 58
  14. 2021Differentiable Sorting Networks for Scalable Sorting and Ranking Supervisioncited by 31
  15. 2022Backpropagation through Combinatorial Algorithms: Identity with Projection Workscited by 29
  16. 2019The Differentiable Cross-Entropy Methodcited by 27
  17. 2023Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspectivecited by 26
  18. 2024Optimistix: modular optimisation in JAX and Equinoxcited by 19
  19. 2020Regularized Optimal Transport is Ground Cost Adversarialcited by 17
  20. 2022GenDR: A Generalized Differentiable Renderercited by 12
  21. 2024LPGD: A General Framework for Backpropagation through Embedded Optimization Layerscited by 6
  22. 2025Differentiable Simulation of Hard Contacts with Soft Gradients for Learning and Controlcited by 5
  23. 2020Universal Activation Function For Machine Learningcited by 3
  24. 2025A Smooth Analytical Formulation of Collision Detection and Rigid Body Dynamics With Contactcited by 3
  25. 2025Smooth Approximations of the Rounding Functioncited by 3
  26. 2026Smoothly Differentiable and Efficiently Vectorizable Contact Manifold Generationcited by 1

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 2.

2 papers.

  1. 2026Emulating Non-Differentiable Metrics via Knowledge-Guided Learning: Introducing the Minkowski Image Losscited by 0
  2. 2026Critical evaluation of PINN for FWD inverse analysis and differentiable FEM as an alternativecited by 0