<p dir="auto">Allow functions that return aux information, similarly to grad or jvp.</p>
<p dir="auto">Repro case:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="!pip install -U -q jax jaxlib

import jax, jax.numpy as jnp, functools

@functools.partial(jax.custom_jvp, nondiff_argnums=(0, 2))
def sample(shape, param, seed):
  return jax.random.gamma(key=seed, shape=shape, a=param)

@sample.defjvp
def sample_jvp(shape, seed, primals, tangents):
  param, = primals
  dparam, = tangents
  dparam = jnp.broadcast_to(dparam, shape)
  samples = sample(shape, param, seed)
  return samples, samples * dparam  # dummy jvp for proof of concept

jax.vmap(lambda seed: sample((2,3), 1., seed))(
    jax.random.split(jax.random.PRNGKey(1), 10))"><pre class="notranslate">!p<span class="pl-s1">ip</span> <span class="pl-s1">install</span> <span class="pl-c1">-</span><span class="pl-v">U</span> <span class="pl-c1">-</span><span class="pl-s1">q</span> <span class="pl-s1">jax</span> <span class="pl-s1">jaxlib</span>

<span class="pl-k">import</span> <span class="pl-s1">jax</span>, <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>, <span class="pl-s1">functools</span>

<span class="pl-en">@<span class="pl-s1">functools</span>.<span class="pl-en">partial</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">custom_jvp</span>, <span class="pl-s1">nondiff_argnums</span><span class="pl-c1">=</span>(<span class="pl-c1">0</span>, <span class="pl-c1">2</span>))</span>
<span class="pl-k">def</span> <span class="pl-en">sample</span>(<span class="pl-s1">shape</span>, <span class="pl-s1">param</span>, <span class="pl-s1">seed</span>):
  <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">gamma</span>(<span class="pl-s1">key</span><span class="pl-c1">=</span><span class="pl-s1">seed</span>, <span class="pl-s1">shape</span><span class="pl-c1">=</span><span class="pl-s1">shape</span>, <span class="pl-s1">a</span><span class="pl-c1">=</span><span class="pl-s1">param</span>)

<span class="pl-en">@<span class="pl-s1">sample</span>.<span class="pl-s1">defjvp</span></span>
<span class="pl-k">def</span> <span class="pl-en">sample_jvp</span>(<span class="pl-s1">shape</span>, <span class="pl-s1">seed</span>, <span class="pl-s1">primals</span>, <span class="pl-s1">tangents</span>):
  <span class="pl-s1">param</span>, <span class="pl-c1">=</span> <span class="pl-s1">primals</span>
  <span class="pl-s1">dparam</span>, <span class="pl-c1">=</span> <span class="pl-s1">tangents</span>
  <span class="pl-s1">dparam</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">broadcast_to</span>(<span class="pl-s1">dparam</span>, <span class="pl-s1">shape</span>)
  <span class="pl-s1">samples</span> <span class="pl-c1">=</span> <span class="pl-en">sample</span>(<span class="pl-s1">shape</span>, <span class="pl-s1">param</span>, <span class="pl-s1">seed</span>)
  <span class="pl-k">return</span> <span class="pl-s1">samples</span>, <span class="pl-s1">samples</span> <span class="pl-c1">*</span> <span class="pl-s1">dparam</span>  <span class="pl-c"># dummy jvp for proof of concept</span>

<span class="pl-s1">jax</span>.<span class="pl-en">vmap</span>(<span class="pl-k">lambda</span> <span class="pl-s1">seed</span>: <span class="pl-en">sample</span>((<span class="pl-c1">2</span>,<span class="pl-c1">3</span>), <span class="pl-c1">1.</span>, <span class="pl-s1">seed</span>))(
    <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">split</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">1</span>), <span class="pl-c1">10</span>))</pre></div>
<p dir="auto">=&gt;</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="/usr/local/lib/python3.6/dist-packages/jax/custom_derivatives.py in bind(self, fun, jvp, *args)
    285     _, env_trace_todo = lu.merge_linear_aux(env_trace_todo1, env_trace_todo2)
    286     if env_trace_todo:
--&gt; 287       raise core.UnexpectedTracerError
    288     return map(core.full_lower, outs)
    289 

UnexpectedTracerError: "><pre class="notranslate"><code class="notranslate">/usr/local/lib/python3.6/dist-packages/jax/custom_derivatives.py in bind(self, fun, jvp, *args)
    285     _, env_trace_todo = lu.merge_linear_aux(env_trace_todo1, env_trace_todo2)
    286     if env_trace_todo:
--&gt; 287       raise core.UnexpectedTracerError
    288     return map(core.full_lower, outs)
    289 

UnexpectedTracerError: 
</code></pre></div>