<p dir="auto">Is there an easy way to disable jitting for a part of the code only?</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# can't @jax.jit because it contains a non-jittable section
def foo(a):
   b = bar(a) # jittable
   c = scipy.bfgs.minimize(b) # non-jittable
   d = zar(c) # jittable
   return d
"><pre class="notranslate"><span class="pl-c"># can't @jax.jit because it contains a non-jittable section</span>
<span class="pl-k">def</span> <span class="pl-en">foo</span>(<span class="pl-s1">a</span>):
   <span class="pl-s1">b</span> <span class="pl-c1">=</span> <span class="pl-en">bar</span>(<span class="pl-s1">a</span>) <span class="pl-c"># jittable</span>
   <span class="pl-s1">c</span> <span class="pl-c1">=</span> <span class="pl-s1">scipy</span>.<span class="pl-s1">bfgs</span>.<span class="pl-en">minimize</span>(<span class="pl-s1">b</span>) <span class="pl-c"># non-jittable</span>
   <span class="pl-s1">d</span> <span class="pl-c1">=</span> <span class="pl-en">zar</span>(<span class="pl-s1">c</span>) <span class="pl-c"># jittable</span>
   <span class="pl-k">return</span> <span class="pl-s1">d</span></pre></div>
<p dir="auto">This can be manually split into two different jit sections</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="bar = jax.jit(bar)
zar = jax.jit(zar)
def foo(a):
   b = bar(a)
   c = scipy.bfgs.minimize(b) # non-jittable
   d = zar(c)
   return d"><pre class="notranslate"><span class="pl-s1">bar</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">bar</span>)
<span class="pl-s1">zar</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">zar</span>)
<span class="pl-k">def</span> <span class="pl-en">foo</span>(<span class="pl-s1">a</span>):
   <span class="pl-s1">b</span> <span class="pl-c1">=</span> <span class="pl-en">bar</span>(<span class="pl-s1">a</span>)
   <span class="pl-s1">c</span> <span class="pl-c1">=</span> <span class="pl-s1">scipy</span>.<span class="pl-s1">bfgs</span>.<span class="pl-en">minimize</span>(<span class="pl-s1">b</span>) <span class="pl-c"># non-jittable</span>
   <span class="pl-s1">d</span> <span class="pl-c1">=</span> <span class="pl-en">zar</span>(<span class="pl-s1">c</span>)
   <span class="pl-k">return</span> <span class="pl-s1">d</span></pre></div>
<p dir="auto">However I often have an extremely long call-stack that explicit jitting is very tedious (especially through compositions and custom_vjps) and it would be nice to be able to specify something like <code class="notranslate">@jax.nojit</code>. Not sure how the implementation detail would work though, maybe automatically generate two sets of jaxprs/HLOs that sandwich the non-jittable part?</p>
<p dir="auto">(and yes, I've tried tfp's bfgs optimizers and unfortunately it is very deficient compared to that of scipy's bfgs)</p>
<p dir="auto">It would be great to be able to <code class="notranslate">jit</code> functions that make use of custom CPU operation, i.e., implemented with NumPy arrays. This would be a really valuable extension point for integrating JAX with existing codes/algorithms, and would possibly solve the final remaining use-cases for autograd.</p>
<p dir="auto">Right now, you can use custom CPU operations if you don't <code class="notranslate">jit</code>, but that adds a very large amount of dispatch overhead.</p>
<p dir="auto">My understanding is this could be possible by making use of XLA's CustomCall support.</p>