Writing study: SoftJAX & SoftTorch: Empowering Automatic Differentiation Libraries with Informative Gradients

ICML overviewCorpus atlasBy contribution typeVisual galleryRules & templateWriting patternsPaper studies

← ICML 2026 oral writing corpus

Anselm Paulus, Andreas René Geist, Vı́t Musil, Sebastian Hoffmann, Georg Martius · Official ICML 2026 oral record

Observed writing choices, not acceptance rules. This study examines the inspected manuscript version. It does not establish why the paper was accepted or that its presentation choices should be copied unchanged.

Figure, table and section ledger

Editorial contribution type: Efficiency, inference & systems (primary).

Layout extraction, not a required paper template. Counts are distinct caption labels; a multi-panel figure counts once. Section spans are estimated from heading positions, including floats and whitespace. Method and limitations.

Canonical PDF

35 physical pages

Before references / appendix

≈ 9.29 page-equivalents

Main figures / tables

8 / 0

Attached appendix figures / tables

15 / 4

Attached appendix span

≈ 21.91 page-equivalents

Before references / appendix: ≈ 26.54%References / intervening material: ≈ 10.86%Attached appendix: ≈ 62.6%Estimated layout envelopes, not text/artwork area. Unresolved appendix boundaries remain within the undivided post-body material.

Section-by-section lengths and source pages

The title/abstract/front matter precede the first numbered heading. A span ends at the next detected top-level heading, not at the last sentence; end-page whitespace is included. References and statements are listed separately. These are not measurements of text-only or artwork area.

Top-level heading	Starts on PDF page	Estimated span	Share of entire PDF
1. Introduction	1	≈ 1.31 pp.	3.74%
2. Softening and Straight-Through Estimation	2	≈ 0.97 pp.	2.77%
3. Elementwise Operators	3	≈ 1.24 pp.	3.54%
4. Axiswise Operators	4	≈ 4.53 pp.	12.94%
5. Benchmarks & Case Study	9	≈ 0.3 pp.	0.86%
6. Conclusion	9	≈ 0.32 pp.	0.91%
Acknowledgements	10	≈ 0.12 pp.	0.34%
Impact Statement	10	≈ 0.12 pp.	0.34%
References	10	≈ 3.8 pp.	10.86%
A. Library Overview	14	≈ 1.61 pp.	4.6%
B. Case Studies	15	≈ 6.39 pp.	18.26%
C. Input Standardization and Squashing	22	≈ 0.37 pp.	1.06%
D. Extended Axiswise Operators	22	≈ 7.16 pp.	20.46%
E. Additional Benchmark Experiments	29	≈ 4.89 pp.	13.97%
F. Choosing a Method and Mode	34	≈ 1.5 pp.	4.29%

All detected figure and table locations

Main text

• Figure 1 · p. 1
• Figure 2 · p. 2
• Figure 3 · p. 3
• Figure 4 · p. 3
• Figure 5 · p. 5
• Figure 6 · p. 7
• Figure 7 · p. 8
• Figure 8 · p. 9

Attached appendix

• Figure 9 · p. 14
• Table 1 · p. 14
• Table 2 · p. 15
• Figure 10 · p. 16
• Figure 11 · p. 17
• Figure 12 · p. 17
• Table 3 · p. 18
• Table 4 · p. 18
• Figure 13 · p. 19
• Figure 14 · p. 19
• Figure 15 · p. 21
• Figure 16 · p. 30
• Figure 17 · p. 31
• Figure 18 · p. 31
• Figure 19 · p. 32
• Figure 20 · p. 32
• Figure 21 · p. 33
• Figure 22 · p. 33
• Figure 23 · p. 34

Argument and organization

1. 1 Introduction; 2 Softening and Straight-Through Estimation

   Introduces a unified interface and exposes a compositional gradient pitfall.

2. 3 Elementwise Operators; 4 Axiswise Operators

   Builds simple surrogates before transport, simplex, permutahedron and sorting-network families.

3. 5 Benchmarks & Case Study; 6 Conclusion

   Compares practical cost and points to operator coverage and application demonstrations.

Problem, claim and contribution

Fragmented relaxation methods motivate a library abstraction with shared softness controls, rather than an unsupported claim that all discrete programs become optimizable.

Introduction, pp.1-2 · Main paper

Argument and section order

The STE product-rule counterexample comes before the operator catalog, so users encounter a failure of naive composition before the convenience API.

Section 2, p.3 · Main paper

Figures, captions and visual explanation

Three vector fields make the zero-gradient pitfall visible beside the corrective decorator; the benchmark page separates runtime and peak memory with one shared method legend.

Figures 3-4, p.3; Figure 8, p.9 · Main paper

Tables, comparisons and evidence

The headline sort benchmark states a full forward-backward pass, smooth mode, input size and RTX 3060 hardware; the fastest method and most memory-efficient method differ.

Section 5 and Figure 8, p.9 · Main paper

Appendix and reproducibility

The library map explicitly marks unavailable Torch counterparts, while additional benchmarks separate runtime, JIT compilation and memory by regularization mode.

Appendix A, p.14; Appendix E, pp.29,31 · Main paper

Clarity and writing conventions

Soft indices and soft Booleans are explained as changed output domains, and the text distinguishes row-stochastic from bistochastic matrices rather than hiding semantic differences behind drop-in terminology.

Sections 2 and 4.2, pp.2,7-8 · Main paper

What to reuse—and what not to infer

Transferable lesson: Pair an easy API with an early counterexample, then document semantic changes and cost tradeoffs at the operator-family level.

Boundary: Informative surrogate gradients do not establish general optimization convergence or hard-program correctness; the highlighted GPU benchmark concerns SoftJAX sort rather than uniform performance across both libraries.

Apply this lesson using the argument and evidence worksheet.

Inspection record and sources

Checked

2026-10-07

Reading status

Writing-focused manuscript scan + selected visuals inspected

Main text

Writing-oriented scan of the main argument through physical p.9, with focused physical-page text inspection on 1, 2, 3, 4, 5, 6, 7, 8, 9. Framing, section progression, mechanism/proof explanation, displays, comparison conditions and conclusions were inspected. This is not exhaustive technical reading, proof verification or reproduction.

Appendix / supplement

Appendix organization and main-text pointers were checked; focused text inspection of physical pages 14, 20, 29, 31. Remaining appendix proofs and details were not exhaustively read. No separately linked supplement inspected.

Visually inspected PDF pages

3, 9

This is a writing-focused scan of the manuscript and the stated supplementary scope: framing, organization, visual explanation, evidence presentation and wording. It is not a line-by-line proof audit, reproduction of results or exhaustive scientific review.

• Published proceedings PDF, PMLR volume 306, retrieved 2026-10-07 · 35 pages