<p dir="auto">I would like to extend the <code class="notranslate">jax.checkpoint</code> aprroach to implement <a href="https://dl.acm.org/doi/10.1145/347837.347846" rel="nofollow">binomial checkpointing</a> for computing the <code class="notranslate">vjp</code> of a time stepping loop the performs computation in step():</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="  def timestep(A,N):
    for i in range(N):
      A = step(A,i)"><pre class="notranslate"><code class="notranslate">  def timestep(A,N):
    for i in range(N):
      A = step(A,i)
</code></pre></div>
<p dir="auto">The<code class="notranslate">remat</code> approach, as I understand it, generates a version of a function that stores inputs of the function and none of the linearization points. For the <code class="notranslate">vjp</code> calculation, it reruns the function calculation thereby regenerating the intermediate linearization points.</p>
<p dir="auto">If I use <code class="notranslate">remat</code> for <code class="notranslate">step(A,i), then for each </code>i<code class="notranslate">, the arguments to </code>step(A,i)<code class="notranslate">will be stored. This is not what I want, because in my test case,</code>A` is very large.</p>
<p dir="auto">An alternate scheme such as binomial checkpointing can use some criterion to decide for which <code class="notranslate">i</code> to store the inputs instead of storing each them for each <code class="notranslate">i</code>. This allows for different tradeoffs between storage and recomputation. To implement such a checkpointing scheme, the following actions will need to be performed:</p>
<ol dir="auto">
<li>Store argument of step(A,i)</li>
<li>Restore arguments of step(Ai,)</li>
<li>Run step(i) without storing/restoring arguments based on step(A,i-1)’s completion.</li>
<li>Compute jvp for step(i) assuming that output of jvp(A,i+1)  and is available</li>
</ol>
<p dir="auto">Does this seem feasible to implement on top of the existing remat approach?<br>
Are there any suggestions on how to go about it?</p>
<p dir="auto">It would be great to have a version of lax.scan used a recursive gradient checkpointing (e.g., "binomial checkpointing") that allows for differentiating through long time series with logarithmic time/space costs.</p>
<p dir="auto">In principle this could be built on top of the experimental <code class="notranslate">remat</code> decorator: <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="527393684" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/1749" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/1749/hovercard" href="https://github.com/google/jax/pull/1749">#1749</a></p>