<p dir="auto">Best shown with a simple reproducer [<a href="https://colab.research.google.com/drive/12kEa87JcSpQBpKM71syZ0MuBE0SrA7IR#scrollTo=2Dd2vZbxIg1Q" rel="nofollow">link</a>].</p>
<p dir="auto">See <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="793759835" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/5509" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/5509/hovercard" href="https://github.com/google/jax/pull/5509">#5509</a>. Things like this work:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax

jax.vmap(lambda x: lax.psum(1, 'i'), axis_name='i')(jnp.arange(3))
jax.jit(jax.vmap(lambda x: lax.psum(1, 'i'), axis_name='i')(jnp.arange(3))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-s1">jax</span>.<span class="pl-en">vmap</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-s1">lax</span>.<span class="pl-en">psum</span>(<span class="pl-c1">1</span>, <span class="pl-s">'i'</span>), <span class="pl-s1">axis_name</span><span class="pl-c1">=</span><span class="pl-s">'i'</span>)(<span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">3</span>))
<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span>(<span class="pl-s1">jax</span>.<span class="pl-en">vmap</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-s1">lax</span>.<span class="pl-en">psum</span>(<span class="pl-c1">1</span>, <span class="pl-s">'i'</span>), <span class="pl-s1">axis_name</span><span class="pl-c1">=</span><span class="pl-s">'i'</span>)(<span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">3</span>))</pre></div>
<p dir="auto">But not this:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="jax.vmap(lambda x: lax.pmin(1, 'i'), axis_name='i')(jnp.arange(3))
jax.jit(jax.vmap(lambda x: lax.pmin(1, 'i'), axis_name='i')(jnp.arange(3))"><pre class="notranslate"><span class="pl-s1">jax</span>.<span class="pl-en">vmap</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-s1">lax</span>.<span class="pl-en">pmin</span>(<span class="pl-c1">1</span>, <span class="pl-s">'i'</span>), <span class="pl-s1">axis_name</span><span class="pl-c1">=</span><span class="pl-s">'i'</span>)(<span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">3</span>))
<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span>(<span class="pl-s1">jax</span>.<span class="pl-en">vmap</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-s1">lax</span>.<span class="pl-en">pmin</span>(<span class="pl-c1">1</span>, <span class="pl-s">'i'</span>), <span class="pl-s1">axis_name</span><span class="pl-c1">=</span><span class="pl-s">'i'</span>)(<span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">3</span>))</pre></div>
<p dir="auto">The first line doesn't work because pmin doesn't have an impl rule, and even if it did the second line wouldn't work because the <code class="notranslate">vmap</code> trace would miss it and then the axis would be unbound when we do the <code class="notranslate">jit</code> lowering.</p>
<p dir="auto">We solve the second issue for <code class="notranslate">psum</code> with <a href="https://github.com/google/jax/blob/84f55f87ecd4b1a69ed7eac48b144b57df9d3f11/jax/_src/lax/parallel.py#L485-L488">a special <code class="notranslate">bind</code> rule</a>. We should do that more generally for collectives.</p>