deploy: SimplexLab/TorchJD@7103b6e

Author: ValerianRey

latest/docs/autojac/backward/index.html

@@ -295,14 +295,22 @@
 <h1>backward<a class="headerlink" href="#backward" title="Link to this heading">¶</a></h1>
 <dl class="py function">
 <dt class="sig sig-object py" id="torchjd.autojac.backward">
-<span class="sig-prename descclassname"><span class="pre">torchjd.autojac.</span></span><span class="sig-name descname"><span class="pre">backward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">tensors</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">inputs</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">retain_graph</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">parallel_chunk_size</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference external" href="https://github.com/SimplexLab/TorchJD/blob/main/src/torchjd/autojac/_backward.py#L9-L83"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torchjd.autojac.backward" title="Link to this definition">¶</a></dt>
-<dd><p>Computes the Jacobians of all values in <code class="docutils literal notranslate"><span class="pre">tensors</span></code> with respect to all <code class="docutils literal notranslate"><span class="pre">inputs</span></code> and
-accumulates them in the <code class="docutils literal notranslate"><span class="pre">.jac</span></code> fields of the <code class="docutils literal notranslate"><span class="pre">inputs</span></code>.</p>
+<span class="sig-prename descclassname"><span class="pre">torchjd.autojac.</span></span><span class="sig-name descname"><span class="pre">backward</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">tensors</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">jac_tensors</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">inputs</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">retain_graph</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">parallel_chunk_size</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference external" href="https://github.com/SimplexLab/TorchJD/blob/main/src/torchjd/autojac/_backward.py#L16-L121"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torchjd.autojac.backward" title="Link to this definition">¶</a></dt>
+<dd><p>Computes the Jacobians of <code class="docutils literal notranslate"><span class="pre">tensors</span></code> with respect to <code class="docutils literal notranslate"><span class="pre">inputs</span></code>, left-multiplied by
+<code class="docutils literal notranslate"><span class="pre">jac_tensors</span></code> (or identity if <code class="docutils literal notranslate"><span class="pre">jac_tensors</span></code> is <code class="docutils literal notranslate"><span class="pre">None</span></code>), and accumulates the results in the
+<code class="docutils literal notranslate"><span class="pre">.jac</span></code> fields of the <code class="docutils literal notranslate"><span class="pre">inputs</span></code>.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters<span class="colon">:</span></dt>
 <dd class="field-odd"><ul class="simple">
-<li><p><strong>tensors</strong> (<span class="sphinx_autodoc_typehints-type"><a class="reference external" href="https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence" title="(in Python v3.14)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Sequence</span></code></a>[<a class="reference external" href="https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor" title="(in PyTorch v2.10)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Tensor</span></code></a>] | <a class="reference external" href="https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor" title="(in PyTorch v2.10)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Tensor</span></code></a></span>) – The tensor or tensors to differentiate. Should be non-empty. The Jacobians will
-have one row for each value of each of these tensors.</p></li>
+<li><p><strong>tensors</strong> (<span class="sphinx_autodoc_typehints-type"><a class="reference external" href="https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence" title="(in Python v3.14)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Sequence</span></code></a>[<a class="reference external" href="https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor" title="(in PyTorch v2.10)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Tensor</span></code></a>] | <a class="reference external" href="https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor" title="(in PyTorch v2.10)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Tensor</span></code></a></span>) – The tensor or tensors to differentiate. Should be non-empty.</p></li>
+<li><p><strong>jac_tensors</strong> (<span class="sphinx_autodoc_typehints-type"><a class="reference external" href="https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence" title="(in Python v3.14)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Sequence</span></code></a>[<a class="reference external" href="https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor" title="(in PyTorch v2.10)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Tensor</span></code></a>] | <a class="reference external" href="https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor" title="(in PyTorch v2.10)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Tensor</span></code></a> | <a class="reference external" href="https://docs.python.org/3/library/constants.html#None" title="(in Python v3.14)"><code class="xref py py-obj docutils literal notranslate"><span class="pre">None</span></code></a></span>) – The initial Jacobians to backpropagate, analog to the <cite>grad_tensors</cite>
+parameter of <cite>torch.autograd.backward</cite>. If provided, it must have the same structure as
+<code class="docutils literal notranslate"><span class="pre">tensors</span></code> and each tensor in <code class="docutils literal notranslate"><span class="pre">jac_tensors</span></code> must match the shape of the corresponding
+tensor in <code class="docutils literal notranslate"><span class="pre">tensors</span></code>, with an extra leading dimension representing the number of rows of
+the resulting Jacobian (e.g. the number of losses). All tensors in <code class="docutils literal notranslate"><span class="pre">jac_tensors</span></code> must
+have the same first dimension. If <code class="docutils literal notranslate"><span class="pre">None</span></code>, defaults to the identity matrix. In this case,
+the standard Jacobian of <code class="docutils literal notranslate"><span class="pre">tensors</span></code> is computed, with one row for each value in the
+<code class="docutils literal notranslate"><span class="pre">tensors</span></code>.</p></li>
 <li><p><strong>inputs</strong> (<span class="sphinx_autodoc_typehints-type"><a class="reference external" href="https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable" title="(in Python v3.14)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Iterable</span></code></a>[<a class="reference external" href="https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor" title="(in PyTorch v2.10)"><code class="xref py py-class docutils literal notranslate"><span class="pre">Tensor</span></code></a>] | <a class="reference external" href="https://docs.python.org/3/library/constants.html#None" title="(in Python v3.14)"><code class="xref py py-obj docutils literal notranslate"><span class="pre">None</span></code></a></span>) – The tensors with respect to which the Jacobians must be computed. These must have
 their <code class="docutils literal notranslate"><span class="pre">requires_grad</span></code> flag set to <code class="docutils literal notranslate"><span class="pre">True</span></code>. If not provided, defaults to the leaf tensors
 that were used to compute the <code class="docutils literal notranslate"><span class="pre">tensors</span></code> parameter.</p></li>
@@ -321,7 +329,7 @@ <h1>backward<a class="headerlink" href="#backward" title="Link to this heading">
 </dl>
 <div class="admonition-example admonition">
 <p class="admonition-title">Example</p>
-<p>The following code snippet showcases a simple usage of <code class="docutils literal notranslate"><span class="pre">backward</span></code>.</p>
+<p>This example shows a simple usage of <code class="docutils literal notranslate"><span class="pre">backward</span></code>.</p>
 <div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span><span class="w"> </span><span class="nn">torch</span>
 <span class="gp">&gt;&gt;&gt;</span>
 <span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span><span class="w"> </span><span class="nn">torchjd.autojac</span><span class="w"> </span><span class="kn">import</span> <span class="n">backward</span>
@@ -341,6 +349,32 @@ <h1>backward<a class="headerlink" href="#backward" title="Link to this heading">
 <p>The <code class="docutils literal notranslate"><span class="pre">.jac</span></code> field of <code class="docutils literal notranslate"><span class="pre">param</span></code> now contains the Jacobian of
 <span class="math notranslate nohighlight">\(\begin{bmatrix}y_1 \\ y_2\end{bmatrix}\)</span> with respect to <code class="docutils literal notranslate"><span class="pre">param</span></code>.</p>
 </div>
+<div class="admonition-example admonition">
+<p class="admonition-title">Example</p>
+<p>This is the same example as before, except that we explicitly specify <code class="docutils literal notranslate"><span class="pre">jac_tensors</span></code> as
+the rows of the identity matrix (which is equivalent to using the default <code class="docutils literal notranslate"><span class="pre">None</span></code>).</p>
+<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span><span class="w"> </span><span class="nn">torch</span>
+<span class="gp">&gt;&gt;&gt;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span><span class="w"> </span><span class="nn">torchjd.autojac</span><span class="w"> </span><span class="kn">import</span> <span class="n">backward</span>
+<span class="gp">&gt;&gt;&gt;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">param</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">1.</span><span class="p">,</span> <span class="mf">2.</span><span class="p">],</span> <span class="n">requires_grad</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="c1"># Compute arbitrary quantities that are function of param</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">y1</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="o">-</span><span class="mf">1.</span><span class="p">,</span> <span class="mf">1.</span><span class="p">])</span> <span class="o">@</span> <span class="n">param</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">y2</span> <span class="o">=</span> <span class="p">(</span><span class="n">param</span> <span class="o">**</span> <span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
+<span class="gp">&gt;&gt;&gt;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">J1</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">])</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">J2</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">0.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">])</span>
+<span class="gp">&gt;&gt;&gt;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">backward</span><span class="p">([</span><span class="n">y1</span><span class="p">,</span> <span class="n">y2</span><span class="p">],</span> <span class="n">jac_tensors</span><span class="o">=</span><span class="p">[</span><span class="n">J1</span><span class="p">,</span> <span class="n">J2</span><span class="p">])</span>
+<span class="gp">&gt;&gt;&gt;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">param</span><span class="o">.</span><span class="n">jac</span>
+<span class="go">tensor([[-1.,  1.],</span>
+<span class="go">        [ 2.,  4.]])</span>
+</pre></div>
+</div>
+<p>Instead of using the identity <code class="docutils literal notranslate"><span class="pre">jac_tensors</span></code>, you can backpropagate some Jacobians obtained
+by a call to <a class="reference internal" href="../jac/#torchjd.autojac.jac" title="torchjd.autojac.jac"><code class="xref py py-func docutils literal notranslate"><span class="pre">torchjd.autojac.jac()</span></code></a> on a later part of the computation graph.</p>
+</div>
 <div class="admonition warning">
 <p class="admonition-title">Warning</p>
 <p>To differentiate in parallel, <code class="docutils literal notranslate"><span class="pre">backward</span></code> relies on <code class="docutils literal notranslate"><span class="pre">torch.vmap</span></code>, which has some