Extending AutoGrad from c++

dfalbel: <p>When creating autograph extensions in C++ by subclassing the <code>torch::autograd::Function</code>…</p>
<p>From the comments here, it seems that only direct <code>Variable</code> arguments are registered in the graph.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/custom_function.h#L28-L35" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/custom_function.h#L28-L35" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/csrc/autograd/custom_function.h#L28-L35</a></h4>
<pre class="onebox"><code class="lang-h"><ol class="start lines" start="28" style="counter-reset: li-counter 27 ;">
<li>/// `forward` can take as many arguments as you want and should return either a</li>
<li>/// variable list or a Variable. Use of any direct Variable arguments will be</li>
<li>/// registered in the graph but no vectors/sets or any other data structures will</li>
<li>/// be traversed. It should take a pointer to `torch::autograd::AutogradContext`</li>
<li>/// as the first argument. Variables can be saved in the `ctx` using `ctx-&gt;save_for_backward`</li>
<li>/// (see `torch::autograd::AutogradContext::save_for_backward`) and other data</li>
<li>/// can be saved in the `ctx-&gt;saved_data` map (see `torch::autograd::AutogradContext::saved_data`)</li>
<li>/// in the form of `&lt;std::string, at::IValue&gt;` pairs.</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Is there any way I could pass a <code>variable_list</code> here? We are building a general interface for libtorch in R, and we only know the number of <code>Variable</code> argument at runtime. I see that python uses a different implementation <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/python_function.cpp#L99" rel="nofollow noopener">here</a> that uses a <code>variable_list</code>.</p>

albanD: <p>I’m afraid this is not possible with the API that we have for torch::autograd::Function.</p>
<p>There is always the solution of creating the Node by hand but this is going to be more complex (and outside of the official public API).</p>

dfalbel: <p>Ok! Thanks!</p>
<p>Would you give us some pointer on how to implement the node by hand?<br>
Would it just work If I implement a subclass of Node and implement it’s methods like:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/custom_function.h#L141-L157" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/custom_function.h#L141-L157" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/csrc/autograd/custom_function.h#L141-L157</a></h4>
<pre class="onebox"><code class="lang-h"><ol class="start lines" start="141" style="counter-reset: li-counter 140 ;">
<li>// CppNode&lt;T&gt; is the Node in the autograd graph that represents the user defined</li>
<li>// backward function for Function&lt;T&gt;. Calls to CppNode::apply are forward to</li>
<li>// T::backward().</li>
<li>template &lt;class T&gt;</li>
<li>struct CppNode : public Node {</li>
<li>
</li>
<li>  variable_list apply(variable_list&amp;&amp; inputs) override;</li>
<li>  AutogradContext ctx_;</li>
<li>  std::vector&lt;bool&gt; is_variable_input_;</li>
<li>  std::vector&lt;VariableInfo&gt; input_info_;</li>
<li>  std::vector&lt;VariableInfo&gt; output_info_;</li>
<li>
</li>
<li>  void release_variables() override;</li>
<li>
</li>
<li>  void set_ctx_grad_fn(const std::shared_ptr&lt;Node&gt; &amp;node);</li>
<li>  void save_variables_to_ctx();</li>
<li>};</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>and</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/python_function.h#L22-L25" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/python_function.h#L22-L25" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/csrc/autograd/python_function.h#L22-L25</a></h4>
<pre class="onebox"><code class="lang-h"><ol class="start lines" start="22" style="counter-reset: li-counter 21 ;">
<li>// A Function which is implemented by a Python object (i.e., a THPFunction).</li>
<li>// Calls to 'apply' are forwarded to the Python method implementation.</li>
<li>struct PyNode : public Node {</li>
<li>  PyNode(THPObjectPtr obj) : obj(obj.release()) {}</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


albanD: <p>I guess the best example would be gather/scatter in this file: <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/functions/comm.cpp">https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/functions/comm.cpp</a></p>
<p>The important bit is that Nodes are only applied. So you need one for the forward and one for the backward.<br>
And you should use the wrap_outputs util from the functions/utils.cpp to properly set up the graph on your outputs before returning them.<br>
Note that if you do inplace or views there is more stuff that you need to do (let me know if you need these).</p>

dfalbel: <p>Thanks! Writing my own node worked. I mostly copied the code for CppNode and torch::Function removing the templates and accepting lambdas.</p>
<p>For other that may get here, here’s what we ended up writing.</p>
<p>See:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.cpp" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.cpp" target="_blank" rel="nofollow noopener">mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.cpp</a></h4>
<pre><code class="lang-cpp">#include "Function.h"
#include &lt;torch/csrc/autograd/functions/accumulate_grad.h&gt;
#include &lt;torch/torch.h&gt;

namespace torch
{
namespace autograd
{

variable_list LanternFunction::apply(
    variable_list args,
    std::function&lt;variable_list(LanternAutogradContext *, variable_list)&gt; forward,
    std::function&lt;variable_list(LanternAutogradContext *, variable_list)&gt; backward)
{

    std::shared_ptr&lt;LanternNode&gt; node(new LanternNode(), deleteNode);
    node-&gt;backward_ = backward;

    const size_t num_inputs = args.size();

</code></pre>

  This file has been truncated. <a href="https://github.com/mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.cpp" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.h" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.h" target="_blank" rel="nofollow noopener">mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.h</a></h4>
<pre><code class="lang-h">#include &lt;iostream&gt;

#define LANTERN_BUILD

#include "lantern/lantern.h"

#include &lt;torch/torch.h&gt;

#include "utils.hpp"

namespace torch
{
namespace autograd
{

// Context to save information during forward that can be accessed in backward
struct LanternAutogradContext
{
    LanternAutogradContext() = default;
    LanternAutogradContext(const LanternAutogradContext &amp;other) = delete;
</code></pre>

  This file has been truncated. <a href="https://github.com/mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.h" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>And here is an example call:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Autograd.cpp#L126-L141" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Autograd.cpp#L126-L141" target="_blank" rel="nofollow noopener">mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Autograd.cpp#L126-L141</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="126" style="counter-reset: li-counter 125 ;">
<li>auto res = LanternFunction::apply(</li>
<li>    {x, y},</li>
<li>    [&amp;](LanternAutogradContext *ctx, variable_list args) {</li>
<li>        ctx-&gt;save_for_backward(args);</li>
<li>        ctx-&gt;saved_data["mul"] = mul;</li>
<li>        return variable_list({args[0] + mul * args[1] + args[0] * args[1]});</li>
<li>    },</li>
<li>    [](LanternAutogradContext *ctx, variable_list grad_output) {</li>
<li>        auto saved = ctx-&gt;get_saved_variables();</li>
<li>        int mul = ctx-&gt;saved_data["mul"].toInt();</li>
<li>        auto var1 = saved[0];</li>
<li>        auto var2 = saved[1];</li>
<li>        variable_list output = {grad_output[0] + grad_output[0] * var2,</li>
<li>                                grad_output[0] * mul + grad_output[0] * var1};</li>
<li>        return output;</li>
<li>    });</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>