<p dir="auto">It would be great to have a way of obtaining elements of a mini batch from within <code class="notranslate">vmap</code>, i.e. an equivalent of <code class="notranslate">pmap's</code> <a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.all_gather.html" rel="nofollow">all_gather</a>.</p>
<p dir="auto">Abstracting away the batch dimension seems like one of the core JAX features, but there are a lot of use cases where batch information is necessary. Some that come to mind are normalizing modules, such as BatchNorm, or contrastive estimation which often depends on positive/negative samples across the minibatch.</p>
<p dir="auto">Related to this, are there any plans to unify <code class="notranslate">vmap</code> and <code class="notranslate">pmap</code> and potentially expose the rest of the parallel operators for use with <code class="notranslate">vmap</code>?</p>
<p dir="auto">JAX has a number of semantically different ways to vectorize, parallelize or loop:</p>
<ul dir="auto">
<li><code class="notranslate">vmap</code> (vectorization)</li>
<li><code class="notranslate">pmap</code> (parallelization)</li>
<li><code class="notranslate">lax.map</code> (XLA loop)</li>
<li><code class="notranslate">map</code> (Python builtin loop)</li>
</ul>
<p dir="auto">(and combinations thereof)</p>
<p dir="auto">These are all different ways of implementing the same "mapped" function with various performance tradeoffs. But they cannot be easily swapped back and forth because they use different syntax:</p>
<ul dir="auto">
<li>Signatures:
<ul dir="auto">
<li><code class="notranslate">vmap</code> and <code class="notranslate">pmap</code> are function transformations</li>
<li><code class="notranslate">lax.map</code> acts along the first axis of arrays</li>
<li><code class="notranslate">map</code> acts directly on Python sequences</li>
</ul>
</li>
<li>Functionality:
<ul dir="auto">
<li><code class="notranslate">pmap</code> supports collective operations along mapped axes, but requires that the axis size matches the number of dimensions.</li>
</ul>
</li>
</ul>
<p dir="auto">It would be a boon to user productivity if we could use the exact same API to specify every type of "map", including hybrid versions like <code class="notranslate">soft_pmap</code> (mixed <code class="notranslate">pmap</code>/<code class="notranslate">vmap</code>).</p>
<p dir="auto">This would be valuable even within JAX, because higher level functionality that relies upon mapping (e.g., <code class="notranslate">jacfwd</code>, <code class="notranslate">jacbwd</code>, <code class="notranslate">hvp</code>, <code class="notranslate">numpy.vectorize</code>, <code class="notranslate">soft_pmap</code>) could be parameterized to support any type of <code class="notranslate">map</code>, not just <code class="notranslate">vmap</code>.</p>