CircleCI update of dev docs (3086).

Author: Circle CI

master/_downloads/0228827bdcc2a6735b1faf26c864145d/plot_OT_1D_smooth.zip

@@ -1908,7 +1908,7 @@ <h1>Source code for ot.backend</h1><div class="highlight"><pre>
 <div class="viewcode-block" id="NumpyBackend.trace">
 <a class="viewcode-back" href="../../gen_modules/ot.backend.html#ot.backend.NumpyBackend.trace">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">trace</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">a</span><span class="p">):</span>
-        <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">trace</span><span class="p">(</span><span class="n">a</span><span class="p">)</span></div>
+        <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">einsum</span><span class="p">(</span><span class="s2">&quot;...ii&quot;</span><span class="p">,</span> <span class="n">a</span><span class="p">)</span></div>
 
 
 <div class="viewcode-block" id="NumpyBackend.inv">
@@ -2570,7 +2570,7 @@ <h1>Source code for ot.backend</h1><div class="highlight"><pre>
 <div class="viewcode-block" id="JaxBackend.trace">
 <a class="viewcode-back" href="../../gen_modules/ot.backend.html#ot.backend.JaxBackend.trace">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">trace</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">a</span><span class="p">):</span>
-        <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">trace</span><span class="p">(</span><span class="n">a</span><span class="p">)</span></div>
+        <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">diagonal</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">axis1</span><span class="o">=-</span><span class="mi">2</span><span class="p">,</span> <span class="n">axis2</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span></div>
 
 
 <div class="viewcode-block" id="JaxBackend.inv">
@@ -3352,7 +3352,7 @@ <h1>Source code for ot.backend</h1><div class="highlight"><pre>
 <div class="viewcode-block" id="TorchBackend.trace">
 <a class="viewcode-back" href="../../gen_modules/ot.backend.html#ot.backend.TorchBackend.trace">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">trace</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">a</span><span class="p">):</span>
-        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">trace</span><span class="p">(</span><span class="n">a</span><span class="p">)</span></div>
+        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">diagonal</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">dim1</span><span class="o">=-</span><span class="mi">2</span><span class="p">,</span> <span class="n">dim2</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span></div>
 
 
 <div class="viewcode-block" id="TorchBackend.inv">
@@ -4015,7 +4015,7 @@ <h1>Source code for ot.backend</h1><div class="highlight"><pre>
 <div class="viewcode-block" id="CupyBackend.trace">
 <a class="viewcode-back" href="../../gen_modules/ot.backend.html#ot.backend.CupyBackend.trace">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">trace</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">a</span><span class="p">):</span>
-        <span class="k">return</span> <span class="n">cp</span><span class="o">.</span><span class="n">trace</span><span class="p">(</span><span class="n">a</span><span class="p">)</span></div>
+        <span class="k">return</span> <span class="n">cp</span><span class="o">.</span><span class="n">trace</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">axis1</span><span class="o">=-</span><span class="mi">2</span><span class="p">,</span> <span class="n">axis2</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span></div>
 
 
 <div class="viewcode-block" id="CupyBackend.inv">

master/_downloads/06f3dd7c3258690ab730be0a658b5e36/plot_ssw_unif_torch.zip

@@ -1524,6 +1524,40 @@ <h1>Source code for ot.utils</h1><div class="highlight"><pre>
 
 
 
+<div class="viewcode-block" id="exp_bures">
+<a class="viewcode-back" href="../../gen_modules/ot.utils.html#ot.exp_bures">[docs]</a>
+<span class="k">def</span><span class="w"> </span><span class="nf">exp_bures</span><span class="p">(</span><span class="n">Sigma</span><span class="p">,</span> <span class="n">S</span><span class="p">,</span> <span class="n">nx</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
+<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    Exponential map in Bures-Wasserstein space at Sigma:</span>
+
+<span class="sd">    .. math::</span>
+<span class="sd">        \exp_\Sigma(S) = (I_d+S)\Sigma(I_d+S).</span>
+
+<span class="sd">    Parameters</span>
+<span class="sd">    ----------</span>
+<span class="sd">    Sigma : array-like (d,d)</span>
+<span class="sd">        SPD matrix</span>
+<span class="sd">    S : array-like (d,d)</span>
+<span class="sd">        Symmetric matrix</span>
+<span class="sd">    nx : module, optional</span>
+<span class="sd">        The numerical backend module to use. If not provided, the backend will</span>
+<span class="sd">        be fetched from the input matrices `Sigma, S`.</span>
+
+<span class="sd">    Returns</span>
+<span class="sd">    -------</span>
+<span class="sd">    P : array-like (d,d)</span>
+<span class="sd">        SPD matrix obtained as the exponential map of S at Sigma</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+    <span class="k">if</span> <span class="n">nx</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
+        <span class="n">nx</span> <span class="o">=</span> <span class="n">get_backend</span><span class="p">(</span><span class="n">Sigma</span><span class="p">,</span> <span class="n">S</span><span class="p">)</span>
+    <span class="n">d</span> <span class="o">=</span> <span class="n">S</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
+    <span class="n">Id</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">eye</span><span class="p">(</span><span class="n">d</span><span class="p">,</span> <span class="n">type_as</span><span class="o">=</span><span class="n">S</span><span class="p">)</span>
+    <span class="n">C</span> <span class="o">=</span> <span class="n">Id</span> <span class="o">+</span> <span class="n">S</span>
+
+    <span class="k">return</span> <span class="n">nx</span><span class="o">.</span><span class="n">einsum</span><span class="p">(</span><span class="s2">&quot;ij,jk,kl -&gt; il&quot;</span><span class="p">,</span> <span class="n">C</span><span class="p">,</span> <span class="n">Sigma</span><span class="p">,</span> <span class="n">C</span><span class="p">)</span></div>
+
+
+
 <div class="viewcode-block" id="check_number_threads">
 <a class="viewcode-back" href="../../gen_modules/ot.utils.html#ot.check_number_threads">[docs]</a>
 <span class="k">def</span><span class="w"> </span><span class="nf">check_number_threads</span><span class="p">(</span><span class="n">numThreads</span><span class="p">):</span>