<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="jax/interpreters/batching.py in _dimsize(dim, aval, x)
    294       return reduce(set.union, map(_dimsize, dim, aval, x))
    295     elif type(dim) is int:
--&gt; 296       return reduce(set.union, map(partial(_dimsize, dim), aval, x))
    297     elif dim is None:
    298       return set()

TypeError: reduce() of empty sequence with no initial value"><pre class="notranslate"><code class="notranslate">jax/interpreters/batching.py in _dimsize(dim, aval, x)
    294       return reduce(set.union, map(_dimsize, dim, aval, x))
    295     elif type(dim) is int:
--&gt; 296       return reduce(set.union, map(partial(_dimsize, dim), aval, x))
    297     elif dim is None:
    298       return set()

TypeError: reduce() of empty sequence with no initial value
</code></pre></div>
<p dir="auto">Let</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="def split(pytree, masks):
    left  = jax.tree_map(lambda p, m:    p if m else None, pytree, masks)
    right = jax.tree_map(lambda p, m: None if m else p   , pytree, masks)
    return left, right"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">split</span>(<span class="pl-s1">pytree</span>, <span class="pl-s1">masks</span>):
    <span class="pl-s1">left</span>  <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">tree_map</span>(<span class="pl-k">lambda</span> <span class="pl-s1">p</span>, <span class="pl-s1">m</span>:    <span class="pl-s1">p</span> <span class="pl-k">if</span> <span class="pl-s1">m</span> <span class="pl-k">else</span> <span class="pl-c1">None</span>, <span class="pl-s1">pytree</span>, <span class="pl-s1">masks</span>)
    <span class="pl-s1">right</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">tree_map</span>(<span class="pl-k">lambda</span> <span class="pl-s1">p</span>, <span class="pl-s1">m</span>: <span class="pl-c1">None</span> <span class="pl-k">if</span> <span class="pl-s1">m</span> <span class="pl-k">else</span> <span class="pl-s1">p</span>   , <span class="pl-s1">pytree</span>, <span class="pl-s1">masks</span>)
    <span class="pl-k">return</span> <span class="pl-s1">left</span>, <span class="pl-s1">right</span></pre></div>
<p dir="auto">acting on <code class="notranslate">pytree</code>, <code class="notranslate">masks</code> which are PyTrees of identical structure; the latter assumed to have Boolean leaves.</p>
<p dir="auto">I would like to define a function <code class="notranslate">merge</code> for which <code class="notranslate">merge(split(...))</code> is the identity.</p>
<p dir="auto">(More generally, we might consider taking some tuple of PyTrees with common structure, and merging them all together, replacing the <code class="notranslate">None</code>s in each PyTree with the leaves (or subtrees) of the others.)</p>
<p dir="auto">As far as I can see, there is no reasonable way to define such a <code class="notranslate">merge</code> at the moment, so I'm opening this issue in search of a solution. (Which I suspect would involve exposing additional functionality through JAX, so I've tagged this as a feature request.)</p>
<hr>
<p dir="auto"><strong>Appendix:</strong></p>
<p dir="auto">My current solution is to take inspiration from <code class="notranslate">vmap</code> and treat <code class="notranslate">None</code>s as PyTree leaves:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="_sentinel = object()

def _merge(*args):
    for arg in args:
        if arg is not _sentinel:
            return arg
    return _sentinel

def merge(*pytrees: PyTree):
    pytrees = [jax._src.tree_util._replace_nones(_sentinel, pytree) for pytree in pytrees]
    return jax.tree_map(_merge, *pytrees)"><pre class="notranslate"><span class="pl-s1">_sentinel</span> <span class="pl-c1">=</span> <span class="pl-en">object</span>()

<span class="pl-k">def</span> <span class="pl-en">_merge</span>(<span class="pl-c1">*</span><span class="pl-s1">args</span>):
    <span class="pl-k">for</span> <span class="pl-s1">arg</span> <span class="pl-c1">in</span> <span class="pl-s1">args</span>:
        <span class="pl-k">if</span> <span class="pl-s1">arg</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-s1">_sentinel</span>:
            <span class="pl-k">return</span> <span class="pl-s1">arg</span>
    <span class="pl-k">return</span> <span class="pl-s1">_sentinel</span>

<span class="pl-k">def</span> <span class="pl-en">merge</span>(<span class="pl-c1">*</span><span class="pl-s1">pytrees</span>: <span class="pl-v">PyTree</span>):
    <span class="pl-s1">pytrees</span> <span class="pl-c1">=</span> [<span class="pl-s1">jax</span>.<span class="pl-s1">_src</span>.<span class="pl-s1">tree_util</span>.<span class="pl-en">_replace_nones</span>(<span class="pl-s1">_sentinel</span>, <span class="pl-s1">pytree</span>) <span class="pl-k">for</span> <span class="pl-s1">pytree</span> <span class="pl-c1">in</span> <span class="pl-s1">pytrees</span>]
    <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-en">tree_map</span>(<span class="pl-s1">_merge</span>, <span class="pl-c1">*</span><span class="pl-s1">pytrees</span>)</pre></div>
<p dir="auto">However of course this relies on private functionality. Possible solutions might be to:</p>
<ul dir="auto">
<li>Expose <code class="notranslate">jax._src.tree_util._replace_nones</code> as public functionality;
<ul dir="auto">
<li>I note that its implementation also seems to rely on private functionality, so it's not possible to duplicate as an end user.</li>
</ul>
</li>
<li>Or add an optional argument to <code class="notranslate">tree_map</code> that treats <code class="notranslate">None</code>s as leaves.
<ul dir="auto">
<li>c.f. <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="575374187" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/2348" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/2348/hovercard" href="https://github.com/google/jax/issues/2348">#2348</a></li>
</ul>
</li>
</ul>