<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<p dir="auto">Split the <code class="notranslate">print(...)</code> statements into different cells in the notebooks so that they're better rendered in .md. If not split, they are usually rendered on separate cells. For example, the <a href="https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#checking-against-numerical-differences" rel="nofollow">following cell</a></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# Set a step size for finite differences calculations
eps = 1e-4

# Check b_grad with scalar finite differences
b_grad_numerical = (loss(W, b + eps / 2.) - loss(W, b - eps / 2.)) / eps
print('b_grad_numerical', b_grad_numerical)
print('b_grad_autodiff', grad(loss, 1)(W, b))

# Check W_grad with finite differences in a random direction
key, subkey = random.split(key)
vec = random.normal(subkey, W.shape)
unitvec = vec / jnp.sqrt(jnp.vdot(vec, vec))
W_grad_numerical = (loss(W + eps / 2. * unitvec, b) - loss(W - eps / 2. * unitvec, b)) / eps
print('W_dirderiv_numerical', W_grad_numerical)
print('W_dirderiv_autodiff', jnp.vdot(grad(loss)(W, b), unitvec))"><pre class="notranslate"><span class="pl-c"># Set a step size for finite differences calculations</span>
<span class="pl-s1">eps</span> <span class="pl-c1">=</span> <span class="pl-c1">1e-4</span>

<span class="pl-c"># Check b_grad with scalar finite differences</span>
<span class="pl-s1">b_grad_numerical</span> <span class="pl-c1">=</span> (<span class="pl-en">loss</span>(<span class="pl-v">W</span>, <span class="pl-s1">b</span> <span class="pl-c1">+</span> <span class="pl-s1">eps</span> <span class="pl-c1">/</span> <span class="pl-c1">2.</span>) <span class="pl-c1">-</span> <span class="pl-en">loss</span>(<span class="pl-v">W</span>, <span class="pl-s1">b</span> <span class="pl-c1">-</span> <span class="pl-s1">eps</span> <span class="pl-c1">/</span> <span class="pl-c1">2.</span>)) <span class="pl-c1">/</span> <span class="pl-s1">eps</span>
<span class="pl-en">print</span>(<span class="pl-s">'b_grad_numerical'</span>, <span class="pl-s1">b_grad_numerical</span>)
<span class="pl-en">print</span>(<span class="pl-s">'b_grad_autodiff'</span>, <span class="pl-en">grad</span>(<span class="pl-s1">loss</span>, <span class="pl-c1">1</span>)(<span class="pl-v">W</span>, <span class="pl-s1">b</span>))

<span class="pl-c"># Check W_grad with finite differences in a random direction</span>
<span class="pl-s1">key</span>, <span class="pl-s1">subkey</span> <span class="pl-c1">=</span> <span class="pl-s1">random</span>.<span class="pl-en">split</span>(<span class="pl-s1">key</span>)
<span class="pl-s1">vec</span> <span class="pl-c1">=</span> <span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">subkey</span>, <span class="pl-v">W</span>.<span class="pl-s1">shape</span>)
<span class="pl-s1">unitvec</span> <span class="pl-c1">=</span> <span class="pl-s1">vec</span> <span class="pl-c1">/</span> <span class="pl-s1">jnp</span>.<span class="pl-en">sqrt</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">vdot</span>(<span class="pl-s1">vec</span>, <span class="pl-s1">vec</span>))
<span class="pl-v">W_grad_numerical</span> <span class="pl-c1">=</span> (<span class="pl-en">loss</span>(<span class="pl-v">W</span> <span class="pl-c1">+</span> <span class="pl-s1">eps</span> <span class="pl-c1">/</span> <span class="pl-c1">2.</span> <span class="pl-c1">*</span> <span class="pl-s1">unitvec</span>, <span class="pl-s1">b</span>) <span class="pl-c1">-</span> <span class="pl-en">loss</span>(<span class="pl-v">W</span> <span class="pl-c1">-</span> <span class="pl-s1">eps</span> <span class="pl-c1">/</span> <span class="pl-c1">2.</span> <span class="pl-c1">*</span> <span class="pl-s1">unitvec</span>, <span class="pl-s1">b</span>)) <span class="pl-c1">/</span> <span class="pl-s1">eps</span>
<span class="pl-en">print</span>(<span class="pl-s">'W_dirderiv_numerical'</span>, <span class="pl-v">W_grad_numerical</span>)
<span class="pl-en">print</span>(<span class="pl-s">'W_dirderiv_autodiff'</span>, <span class="pl-s1">jnp</span>.<span class="pl-en">vdot</span>(<span class="pl-en">grad</span>(<span class="pl-s1">loss</span>)(<span class="pl-v">W</span>, <span class="pl-s1">b</span>), <span class="pl-s1">unitvec</span>))</pre></div>
<p dir="auto">gets rendered like</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="b_grad_numerical"><pre class="notranslate"><code class="notranslate">b_grad_numerical
</code></pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=" -0.29325485
b_grad_autodiff -0.29227245"><pre class="notranslate"><code class="notranslate"> -0.29325485
b_grad_autodiff -0.29227245
</code></pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="W_dirderiv_numerical -0.2002716
W_dirderiv_autodiff -0.19909117"><pre class="notranslate"><code class="notranslate">W_dirderiv_numerical -0.2002716
W_dirderiv_autodiff -0.19909117
</code></pre></div>
<p dir="auto">It is cumbersome to separate the print calls, but without a better fix in the underlying rendering, this is likely the best solution for now...</p>
<p dir="auto">assume we have written a function which work when a, b are single instances</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def f(a,b):
   #some computation  
   return result"><pre class="notranslate"><code class="notranslate">def f(a,b):
   #some computation  
   return result
</code></pre></div>
<p dir="auto">using <strong>vmap</strong> we can repeat same computation for different instances of a and single instance of b as follows</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="v = partial(vmap, in_axes=(0, None))
af = v(f)"><pre class="notranslate"><code class="notranslate">v = partial(vmap, in_axes=(0, None))
af = v(f)
</code></pre></div>
<p dir="auto">now similarly, we can repeat same computation for single instances of a and multiple instances of b as follows</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="v = partial(vmap, in_axes=(None, 0))
bf = v(f)"><pre class="notranslate"><code class="notranslate">v = partial(vmap, in_axes=(None, 0))
bf = v(f)
</code></pre></div>
<p dir="auto">we have different functions <code class="notranslate">af</code>, <code class="notranslate">bf</code>. Is there way to combine this into single function and handle automatically ? If input is batched over <strong>a</strong> use <code class="notranslate">af</code> and if it is batched over <strong>b</strong> use <code class="notranslate">bf</code>.</p>
<p dir="auto">Thanks in advance.</p>