# Do we have differentiable sort in Julia?

**URL:** <https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713>\
**Category:** Machine Learning\
**Tags:** sort, autodiff\
**Created:** [August 11, 2023, 12:35pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713 "2023-08-11T12:35:51Z")\
**Posts on this page:** 15\
**Page:** 1

<div class="post-metadata">

**Author:** ![pat-alt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pat-alt/32/34264_2.png) [@pat-alt](https://discourse.julialang.org/u/pat-alt)\
**Post date:** [August 11, 2023, 12:35pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/1 "2023-08-11T12:35:51Z")

</div>

Python has support for [fast differentiable sorting](https://github.com/google-research/fast-soft-sort/tree/master) which is an implementation of Blondel et al. ([2020](https://arxiv.org/pdf/2002.08871.pdf)). I don’t think we have this in Julia at the moment?

I have specific [use case](https://github.com/JuliaTrustworthyAI/ConformalPrediction.jl/issues/55) for [ConformalPrediction.jl](https://github.com/JuliaTrustworthyAI/ConformalPrediction.jl), but beyond that it seems there should be many applications in ML where this could be useful (just based on reading the abstract of the paper).

Does anyone know of an existing implementation? If not, where would be a good place to kindly ask for it?

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [August 11, 2023, 1:10pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/2 "2023-08-11T13:10:31Z")

</div>

Our package [InferOpt.jl](https://github.com/axelparmentier/InferOpt.jl) implements several SoTA methods for converting a combinatorial solver (anything that can be formulated as a MILP) into a differentiable layer. We have drawn a lot of inspiration from the works of Blondel and colleagues, so I invite you to check out the code, along with our paper [_Learning with Combinatorial Optimization Layers: a Probabilistic Approach_](https://arxiv.org/abs/2207.13513).

The particular “fast differentiable sorting” algorithm you mentioned is not part of our package, because we offer to wrap arbitrary solvers instead of focusing on specific problems. But you can discuss it with my colleague @BatyLeo, since it is very close to our research interests. Perhaps there could even be a library of custom differentiable algorithms to complement the generic wrappers of InferOpt.jl?

---

<div class="post-metadata">

**Author:** ![pat-alt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pat-alt/32/34264_2.png) [@pat-alt](https://discourse.julialang.org/u/pat-alt)\
**Post date:** [August 11, 2023, 1:28pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/3 "2023-08-11T13:28:46Z")

</div>

Ohhh amazing! 👏

From a brief look at it, this indeed looks like it should do the trick for me. Having wrappers would be amazing @BatyLeo, since it’s not immediately obvious to me from the docs how exactly I should build the layer that does “fast differentiable sorting”. But for now this gives me a place to start. Thank you very much!

---

<div class="post-metadata">

**Author:** ![ericphanson](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ericphanson/32/215186_2.png) [@ericphanson](https://discourse.julialang.org/u/ericphanson)\
**Post date:** [August 11, 2023, 4:03pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/4 "2023-08-11T16:03:49Z")

</div>

Really cool! Sorry since this is a bit off-topic, but I was reading the paper and wondered if there is an advantage in e.g. the warcraft example, of training end-to-end as your package allows, as opposed to separately training the ML algorithm to output paths, and then running the CO algorithm as usual? I was hoping to see a comparison in the benchmarks but didn’t find one.

---

<div class="post-metadata">

**Author:** ![GunnarFarneback](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gunnarfarneback/32/1827_2.png) [@GunnarFarneback](https://discourse.julialang.org/u/GunnarFarneback)\
**Post date:** [August 11, 2023, 5:26pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/5 "2023-08-11T17:26:21Z")

</div>

I guess I would find it if I read the linked papers but can someone who already knows briefly explain what it means for sorting to be differentiable? Naively one would expect derivatives to be mostly zero.

---

<div class="post-metadata">

**Author:** ![TomasSkrivan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomasskrivan/32/36280_2.png) [@TomasSkrivan](https://discourse.julialang.org/u/TomasSkrivan)\
**Post date:** [August 11, 2023, 6:23pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/6 "2023-08-11T18:23:29Z")

</div>

If all your values in an array are distinct then there is a small neighbourhood around your array such that doing sorting is just an application of fixed permutation matrix. Thus sorting is locally just a linear map and its derivative is the same linear map. In forward mode, you sort your primary array and apply the same permutation to the tangent array. For arrays which have at least two identical elements, this permutation matrix is not uniquely determined, thus sorting is not differentiable there.

---

<div class="post-metadata">

**Author:** ![BatyLeo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/batyleo/32/35017_2.png) [@BatyLeo](https://discourse.julialang.org/u/BatyLeo)\
**Post date:** [August 11, 2023, 8:25pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/7 "2023-08-11T20:25:27Z")

</div>

Hello,  
As @gdalle mentionned, our package Inferopt.jl provides generic wrappers to turn any combinatorial optimization algorithm into a differentiable function.  
Depending on your needs, it could be enough to directly use these wrappers.  
Feel free to ask more questions if you face any difficulties, we know that the documentation is very minimalistic and lacks details and examples at the moment 🙂 (also do not hesitate to give us feedback so we can improve it for future users !)

However, you won’t find the exact same “fast differentiable sorting” technique as in the paper you cited. By briefly looking at it, it seems quite close to our `Regularized` layers, with some additions specifically taylored for sorting and ranking functions. It does not seem that difficult to implement a specific wrapper, I’ll look more in details when I’m back from holidays.

---

<div class="post-metadata">

**Author:** ![BatyLeo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/batyleo/32/35017_2.png) [@BatyLeo](https://discourse.julialang.org/u/BatyLeo)\
**Post date:** [August 11, 2023, 8:47pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/8 "2023-08-11T20:47:10Z")

</div>

Training the ML model separately would mean that you separately learn to predict edge/cell costs in a supervised way. This first means you assume to have access to the true edge costs in your dataset, which is not the case for our losses, except the SPO+ loss. In the latter case, it’s still better to learn in an integrated way, because a small loss in terms of predicted costs can lead to a large error in term of predicted path, and the other way around. For more details, you can checkout the “Smart Predict then Optimize” paper, which explains this quite well.

---

<div class="post-metadata">

**Author:** ![BatyLeo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/batyleo/32/35017_2.png) [@BatyLeo](https://discourse.julialang.org/u/BatyLeo)\
**Post date:** [August 11, 2023, 8:56pm UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/9 "2023-08-11T20:56:33Z")

</div>

You’re right, combinatorial optimization algorithms like `sort` are by definition piecewise constant, and therefore have zero derivative almost everywhere.

However, in many applications we still want to be able to compute meaningful gradients through combinatorial algorithms, for example when their are combined with machine learning models. In this case, there are some ways to build a regularized version of the combinatorial algorithm such that it’s close to the original one, and is differentiable with useful gradients.

Edit: `sort` is piecewise linear, and `ranking` is piecewise constant

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [August 12, 2023, 6:25am UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/10 "2023-08-12T06:25:59Z")

</div>

> [@GunnarFarneback](#):
>
> can someone who already knows briefly explain what it means for sorting to be differentiable?

Many papers, including ours, study combinatorial solvers that can be formulated as (Mixed Integer) Linear Programs:

\theta \longmapsto \arg\max\_x \theta^\top x \text{ such that } Ax \leq b

The mapping from \theta to a solution is piecewise constant, so the real problem is not the jump points (which we almost surely never hit) but rather the regions ot desperate flatness. Thus, “non-differentiable” is true in the litteral sense (no gradients at some points), but mostly in the practical sense (all gradients are useless anyway). What we really want here is an _approximation_ of the function with more informative gradients.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [August 12, 2023, 6:27am UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/11 "2023-08-12T06:27:16Z")

</div>

> [@TomasSkrivan](#):
>
> sorting is locally just a linear map

Which brings up an interesting question: I’m not sure whether sorting itself can be formulated as a linear program. Finding the sort _permutation_, sure ([algorithms - Sorting as a linear program - Computer Science Stack Exchange](https://cs.stackexchange.com/questions/4805/sorting-as-a-linear-program)), but the actual sorted values?

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [August 12, 2023, 6:30am UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/12 "2023-08-12T06:30:23Z")

</div>

> [@pat-alt](#):
>
> it’s not immediately obvious to me from the docs how exactly I should build the layer that does “fast differentiable sorting”

As Léo said you won’t get the exact same layer from the paper. However this will give you similar _behavior_ (differentiable generation of a sort permutation):

```julia
using InferOpt

ranking(θ::AbstractVector) = invperm(sortperm(θ))
layer = PerturbedAdditive(ranking; ε=1.0, nb_samples=10)

```

The parameter `ε` controls the noise in the approximation, and `nb_samples` tells you how many times you evaluate the original function to get a gradient approximation.  
The object `layer` is a callable whose derivatives are automatically computed using ChainRules.jl

---

<div class="post-metadata">

**Author:** ![BatyLeo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/batyleo/32/35017_2.png) [@BatyLeo](https://discourse.julialang.org/u/BatyLeo)\
**Post date:** [August 23, 2023, 11:11am UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/13 "2023-08-23T11:11:14Z")

</div>

Fast differentiable sorting and ranking is now implemented in this new [PR](https://github.com/axelparmentier/InferOpt.jl/pull/88).

It implements `soft_sort` and `soft_rank` differentiable methods.

You can test it by installing the branch version as follows:

```julia
pkg> add InferOpt#soft-sort

```

Reproducing one of the figures of the paper:

```julia
using InferOpt, Plots, Zygote
plot(x -> sort([0.0, x, 1.0, 2.0]; rev=true)[2], label="Sort")
plot!(x -> soft_sort([0.0, x, 1.0, 2.0]; rev=true)[2], label="Soft sort")
plot!(y -> gradient(x -> soft_sort([0.0, x, 1.0, 2.0]; ε=1.0, rev=true)[2], y)[1], label="Soft sort derivative")

```

![plot](https://global.discourse-cdn.com/julialang/original/3X/8/e/8e0fe78fa8ac879a27d5fe1b37be288f640a0b4c.png)

---

<div class="post-metadata">

**Author:** ![pat-alt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pat-alt/32/34264_2.png) [@pat-alt](https://discourse.julialang.org/u/pat-alt)\
**Post date:** [August 24, 2023, 10:25am UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/14 "2023-08-24T10:25:22Z")

</div>

This is amazing 🔥 thanks @BatyLeo

---

<div class="post-metadata">

**Author:** ![BatyLeo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/batyleo/32/35017_2.png) [@BatyLeo](https://discourse.julialang.org/u/BatyLeo)\
**Post date:** [September 12, 2023, 11:47am UTC](https://discourse.julialang.org/t/do-we-have-differentiable-sort-in-julia/102713/15 "2023-09-12T11:47:04Z")

</div>

The pull request is now merged, and included in the new `v0.6.0` release of `InferOpt` !