CircleCI update of dev docs (2955).

Author: Circle CI

master/_downloads/0228827bdcc2a6735b1faf26c864145d/plot_OT_1D_smooth.zip

@@ -0,0 +1,61 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n# Different gradient computations for regularized optimal transport\n\nThis example illustrates the differences in terms of computation time between the gradient options for the Sinkhorn solver.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# Author: Sonia Mazelet <[email protected]>\n#\n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 1\n\nimport matplotlib.pylab as pl\nimport ot\nfrom ot.backend import torch"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Time comparison of the Sinkhorn solver for different gradient options\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "n_trials = 10\ntimes_autodiff = torch.zeros(n_trials)\ntimes_envelope = torch.zeros(n_trials)\ntimes_last_step = torch.zeros(n_trials)\n\nn_samples_s = 300\nn_samples_t = 300\nn_features = 5\nreg = 0.03\n\n# Time required for the Sinkhorn solver and gradient computations, for different gradient options over multiple Gaussian distributions\nfor i in range(n_trials):\n    x = torch.rand((n_samples_s, n_features))\n    y = torch.rand((n_samples_t, n_features))\n    a = ot.utils.unif(n_samples_s)\n    b = ot.utils.unif(n_samples_t)\n    M = ot.dist(x, y)\n\n    a = torch.tensor(a, requires_grad=True)\n    b = torch.tensor(b, requires_grad=True)\n    M = M.clone().detach().requires_grad_(True)\n\n    # autodiff provides the gradient for all the outputs (plan, value, value_linear)\n    ot.tic()\n    res_autodiff = ot.solve(M, a, b, reg=reg, grad=\"autodiff\")\n    res_autodiff.value.backward()\n    times_autodiff[i] = ot.toq()\n\n    a = a.clone().detach().requires_grad_(True)\n    b = b.clone().detach().requires_grad_(True)\n    M = M.clone().detach().requires_grad_(True)\n\n    # envelope provides the gradient for value\n    ot.tic()\n    res_envelope = ot.solve(M, a, b, reg=reg, grad=\"envelope\")\n    res_envelope.value.backward()\n    times_envelope[i] = ot.toq()\n\n    a = a.clone().detach().requires_grad_(True)\n    b = b.clone().detach().requires_grad_(True)\n    M = M.clone().detach().requires_grad_(True)\n\n    # last_step provides the gradient for all the outputs, but only for the last iteration of the Sinkhorn algorithm\n    ot.tic()\n    res_last_step = ot.solve(M, a, b, reg=reg, grad=\"last_step\")\n    res_last_step.value.backward()\n    times_last_step[i] = ot.toq()\n\npl.figure(1, figsize=(5, 3))\npl.ticklabel_format(axis=\"y\", style=\"sci\", scilimits=(0, 0))\npl.boxplot(\n    ([times_autodiff, times_envelope, times_last_step]),\n    tick_labels=[\"autodiff\", \"envelope\", \"last_step\"],\n    showfliers=False,\n)\npl.ylabel(\"Time (s)\")\npl.show()"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.10.15"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

master/_downloads/06f3dd7c3258690ab730be0a658b5e36/plot_ssw_unif_torch.zip

@@ -0,0 +1,85 @@
+# -*- coding: utf-8 -*-
+"""
+================================================
+Different gradient computations for regularized optimal transport
+================================================
+
+This example illustrates the differences in terms of computation time between the gradient options for the Sinkhorn solver.
+
+"""
+
+# Author: Sonia Mazelet <[email protected]>
+#
+# License: MIT License
+
+# sphinx_gallery_thumbnail_number = 1
+
+import matplotlib.pylab as pl
+import ot
+from ot.backend import torch
+
+
+##############################################################################
+# Time comparison of the Sinkhorn solver for different gradient options
+# -------------
+
+
+# %% parameters
+
+n_trials = 10
+times_autodiff = torch.zeros(n_trials)
+times_envelope = torch.zeros(n_trials)
+times_last_step = torch.zeros(n_trials)
+
+n_samples_s = 300
+n_samples_t = 300
+n_features = 5
+reg = 0.03
+
+# Time required for the Sinkhorn solver and gradient computations, for different gradient options over multiple Gaussian distributions
+for i in range(n_trials):
+    x = torch.rand((n_samples_s, n_features))
+    y = torch.rand((n_samples_t, n_features))
+    a = ot.utils.unif(n_samples_s)
+    b = ot.utils.unif(n_samples_t)
+    M = ot.dist(x, y)
+
+    a = torch.tensor(a, requires_grad=True)
+    b = torch.tensor(b, requires_grad=True)
+    M = M.clone().detach().requires_grad_(True)
+
+    # autodiff provides the gradient for all the outputs (plan, value, value_linear)
+    ot.tic()
+    res_autodiff = ot.solve(M, a, b, reg=reg, grad="autodiff")
+    res_autodiff.value.backward()
+    times_autodiff[i] = ot.toq()
+
+    a = a.clone().detach().requires_grad_(True)
+    b = b.clone().detach().requires_grad_(True)
+    M = M.clone().detach().requires_grad_(True)
+
+    # envelope provides the gradient for value
+    ot.tic()
+    res_envelope = ot.solve(M, a, b, reg=reg, grad="envelope")
+    res_envelope.value.backward()
+    times_envelope[i] = ot.toq()
+
+    a = a.clone().detach().requires_grad_(True)
+    b = b.clone().detach().requires_grad_(True)
+    M = M.clone().detach().requires_grad_(True)
+
+    # last_step provides the gradient for all the outputs, but only for the last iteration of the Sinkhorn algorithm
+    ot.tic()
+    res_last_step = ot.solve(M, a, b, reg=reg, grad="last_step")
+    res_last_step.value.backward()
+    times_last_step[i] = ot.toq()
+
+pl.figure(1, figsize=(5, 3))
+pl.ticklabel_format(axis="y", style="sci", scilimits=(0, 0))
+pl.boxplot(
+    ([times_autodiff, times_envelope, times_last_step]),
+    tick_labels=["autodiff", "envelope", "last_step"],
+    showfliers=False,
+)
+pl.ylabel("Time (s)")
+pl.show()

master/_downloads/07fcc19ba03226cd3d83d4e40ec44385/auto_examples_python.zip

@@ -209,11 +209,13 @@ <h1>Source code for ot.solvers</h1><div class="highlight"><pre>
 <span class="sd">    verbose : bool, optional</span>
 <span class="sd">        Print information in the solver, by default False</span>
 <span class="sd">    grad : str, optional</span>
-<span class="sd">        Type of gradient computation, either or &#39;autodiff&#39; or &#39;envelope&#39;  used only for</span>
+<span class="sd">        Type of gradient computation, either or &#39;autodiff&#39;, &#39;envelope&#39; or &#39;last_step&#39; used only for</span>
 <span class="sd">        Sinkhorn solver. By default &#39;autodiff&#39; provides gradients wrt all</span>
 <span class="sd">        outputs (`plan, value, value_linear`) but with important memory cost.</span>
 <span class="sd">        &#39;envelope&#39; provides gradients only for `value` and and other outputs are</span>
-<span class="sd">        detached. This is useful for memory saving when only the value is needed.</span>
+<span class="sd">        detached. This is useful for memory saving when only the value is needed. &#39;last_step&#39; provides</span>
+<span class="sd">        gradients only for the last iteration of the Sinkhorn solver, but provides gradient for both the OT plan and the objective values.</span>
+<span class="sd">        &#39;detach&#39; does not compute the gradients for the Sinkhorn solver.</span>
 
 <span class="sd">    Returns</span>
 <span class="sd">    -------</span>
@@ -365,7 +367,6 @@ <h1>Source code for ot.solvers</h1><div class="highlight"><pre>
 <span class="sd">        linear regression. NeurIPS.</span>
 
 <span class="sd">    &quot;&quot;&quot;</span>
-
     <span class="c1"># detect backend</span>
     <span class="n">nx</span> <span class="o">=</span> <span class="n">get_backend</span><span class="p">(</span><span class="n">M</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">c</span><span class="p">)</span>
 
@@ -496,7 +497,11 @@ <h1>Source code for ot.solvers</h1><div class="highlight"><pre>
                 <span class="n">potentials</span> <span class="o">=</span> <span class="p">(</span><span class="n">log</span><span class="p">[</span><span class="s2">&quot;u&quot;</span><span class="p">],</span> <span class="n">log</span><span class="p">[</span><span class="s2">&quot;v&quot;</span><span class="p">])</span>
 
             <span class="k">elif</span> <span class="n">reg_type</span><span class="o">.</span><span class="n">lower</span><span class="p">()</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;entropy&quot;</span><span class="p">,</span> <span class="s2">&quot;kl&quot;</span><span class="p">]:</span>
-                <span class="k">if</span> <span class="n">grad</span> <span class="o">==</span> <span class="s2">&quot;envelope&quot;</span><span class="p">:</span>  <span class="c1"># if envelope then detach the input</span>
+                <span class="k">if</span> <span class="n">grad</span> <span class="ow">in</span> <span class="p">[</span>
+                    <span class="s2">&quot;envelope&quot;</span><span class="p">,</span>
+                    <span class="s2">&quot;last_step&quot;</span><span class="p">,</span>
+                    <span class="s2">&quot;detach&quot;</span><span class="p">,</span>
+                <span class="p">]:</span>  <span class="c1"># if envelope, last_step or detach then detach the input</span>
                     <span class="n">M0</span><span class="p">,</span> <span class="n">a0</span><span class="p">,</span> <span class="n">b0</span> <span class="o">=</span> <span class="n">M</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">b</span>
                     <span class="n">M</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">b</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">detach</span><span class="p">(</span><span class="n">M</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">)</span>
 
@@ -505,6 +510,12 @@ <h1>Source code for ot.solvers</h1><div class="highlight"><pre>
                     <span class="n">max_iter</span> <span class="o">=</span> <span class="mi">1000</span>
                 <span class="k">if</span> <span class="n">tol</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                     <span class="n">tol</span> <span class="o">=</span> <span class="mf">1e-9</span>
+                <span class="k">if</span> <span class="n">grad</span> <span class="o">==</span> <span class="s2">&quot;last_step&quot;</span><span class="p">:</span>
+                    <span class="k">if</span> <span class="n">max_iter</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
+                        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
+                            <span class="s2">&quot;The maximum number of iterations must be greater than 0 when using grad=last_step.&quot;</span>
+                        <span class="p">)</span>
+                    <span class="n">max_iter</span> <span class="o">=</span> <span class="n">max_iter</span> <span class="o">-</span> <span class="mi">1</span>
 
                 <span class="n">plan</span><span class="p">,</span> <span class="n">log</span> <span class="o">=</span> <span class="n">sinkhorn_log</span><span class="p">(</span>
                     <span class="n">a</span><span class="p">,</span>
@@ -517,6 +528,22 @@ <h1>Source code for ot.solvers</h1><div class="highlight"><pre>
                     <span class="n">verbose</span><span class="o">=</span><span class="n">verbose</span><span class="p">,</span>
                 <span class="p">)</span>
 
+                <span class="n">potentials</span> <span class="o">=</span> <span class="p">(</span><span class="n">log</span><span class="p">[</span><span class="s2">&quot;log_u&quot;</span><span class="p">],</span> <span class="n">log</span><span class="p">[</span><span class="s2">&quot;log_v&quot;</span><span class="p">])</span>
+
+                <span class="c1"># if last_step, compute the last step of the Sinkhorn algorithm with the non-detached inputs</span>
+                <span class="k">if</span> <span class="n">grad</span> <span class="o">==</span> <span class="s2">&quot;last_step&quot;</span><span class="p">:</span>
+                    <span class="n">loga</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">a0</span><span class="p">)</span>
+                    <span class="n">logb</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">b0</span><span class="p">)</span>
+                    <span class="n">v</span> <span class="o">=</span> <span class="n">logb</span> <span class="o">-</span> <span class="n">nx</span><span class="o">.</span><span class="n">logsumexp</span><span class="p">(</span><span class="o">-</span><span class="n">M0</span> <span class="o">/</span> <span class="n">reg</span> <span class="o">+</span> <span class="n">potentials</span><span class="p">[</span><span class="mi">0</span><span class="p">][:,</span> <span class="kc">None</span><span class="p">],</span> <span class="mi">0</span><span class="p">)</span>
+                    <span class="n">u</span> <span class="o">=</span> <span class="n">loga</span> <span class="o">-</span> <span class="n">nx</span><span class="o">.</span><span class="n">logsumexp</span><span class="p">(</span><span class="o">-</span><span class="n">M0</span> <span class="o">/</span> <span class="n">reg</span> <span class="o">+</span> <span class="n">potentials</span><span class="p">[</span><span class="mi">1</span><span class="p">][</span><span class="kc">None</span><span class="p">,</span> <span class="p">:],</span> <span class="mi">1</span><span class="p">)</span>
+                    <span class="n">plan</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="o">-</span><span class="n">M0</span> <span class="o">/</span> <span class="n">reg</span> <span class="o">+</span> <span class="n">u</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span> <span class="o">+</span> <span class="n">v</span><span class="p">[</span><span class="kc">None</span><span class="p">,</span> <span class="p">:])</span>
+                    <span class="n">potentials</span> <span class="o">=</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span>
+                    <span class="n">log</span><span class="p">[</span><span class="s2">&quot;niter&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">max_iter</span> <span class="o">+</span> <span class="mi">1</span>
+                    <span class="n">log</span><span class="p">[</span><span class="s2">&quot;log_u&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">u</span>
+                    <span class="n">log</span><span class="p">[</span><span class="s2">&quot;log_v&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">v</span>
+                    <span class="n">log</span><span class="p">[</span><span class="s2">&quot;u&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">u</span><span class="p">)</span>
+                    <span class="n">log</span><span class="p">[</span><span class="s2">&quot;v&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">v</span><span class="p">)</span>
+
                 <span class="n">value_linear</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">M</span> <span class="o">*</span> <span class="n">plan</span><span class="p">)</span>
 
                 <span class="k">if</span> <span class="n">reg_type</span><span class="o">.</span><span class="n">lower</span><span class="p">()</span> <span class="o">==</span> <span class="s2">&quot;entropy&quot;</span><span class="p">:</span>
@@ -526,8 +553,6 @@ <h1>Source code for ot.solvers</h1><div class="highlight"><pre>
                         <span class="n">plan</span><span class="p">,</span> <span class="n">a</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span> <span class="o">*</span> <span class="n">b</span><span class="p">[</span><span class="kc">None</span><span class="p">,</span> <span class="p">:]</span>
                     <span class="p">)</span>
 
-                <span class="n">potentials</span> <span class="o">=</span> <span class="p">(</span><span class="n">log</span><span class="p">[</span><span class="s2">&quot;log_u&quot;</span><span class="p">],</span> <span class="n">log</span><span class="p">[</span><span class="s2">&quot;log_v&quot;</span><span class="p">])</span>
-
                 <span class="k">if</span> <span class="n">grad</span> <span class="o">==</span> <span class="s2">&quot;envelope&quot;</span><span class="p">:</span>  <span class="c1"># set the gradient at convergence</span>
                     <span class="n">value</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">set_gradients</span><span class="p">(</span>
                         <span class="n">value</span><span class="p">,</span>