<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 issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp

pytree = {'theta': jnp.ones((2, 1)), 'alpha': jnp.zeros((2, ))}
leaves, treedef = jax.tree_flatten(pytree)
reconstructed_pytree = jax.tree_unflatten(treedef, leaves)
print(reconstructed_pytree == pytree)

from jax.flatten_util import ravel_pytree
array, unravel_fn = ravel_pytree(pytree)
reconstructed_pytree2 = unravel_fn(array)
print(reconstructed_pytree2 == pytree)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</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">pytree</span> <span class="pl-c1">=</span> {<span class="pl-s">'theta'</span>: <span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>((<span class="pl-c1">2</span>, <span class="pl-c1">1</span>)), <span class="pl-s">'alpha'</span>: <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>((<span class="pl-c1">2</span>, ))}
<span class="pl-s1">leaves</span>, <span class="pl-s1">treedef</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">tree_flatten</span>(<span class="pl-s1">pytree</span>)
<span class="pl-s1">reconstructed_pytree</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">tree_unflatten</span>(<span class="pl-s1">treedef</span>, <span class="pl-s1">leaves</span>)
<span class="pl-en">print</span>(<span class="pl-s1">reconstructed_pytree</span> <span class="pl-c1">==</span> <span class="pl-s1">pytree</span>)

<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">flatten_util</span> <span class="pl-k">import</span> <span class="pl-s1">ravel_pytree</span>
<span class="pl-s1">array</span>, <span class="pl-s1">unravel_fn</span> <span class="pl-c1">=</span> <span class="pl-en">ravel_pytree</span>(<span class="pl-s1">pytree</span>)
<span class="pl-s1">reconstructed_pytree2</span> <span class="pl-c1">=</span> <span class="pl-en">unravel_fn</span>(<span class="pl-s1">array</span>)
<span class="pl-en">print</span>(<span class="pl-s1">reconstructed_pytree2</span> <span class="pl-c1">==</span> <span class="pl-s1">pytree</span>)</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="True
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[&lt;ipython-input-7-f0aa3716322d&gt;](https://localhost:8080/#) in &lt;module&gt;()
     10 array, unravel_fn = ravel_pytree(pytree)
     11 reconstructed_pytree2 = unravel_fn(array)
---&gt; 12 print(reconstructed_pytree2 == pytree)

1 frames
[/usr/lib/python3.7/functools.py](https://localhost:8080/#) in _method(*args, **keywords)
    374             cls_or_self, *rest = args
    375             call_args = (cls_or_self,) + self.args + tuple(rest)
--&gt; 376             return self.func(*call_args, **call_keywords)
    377         _method.__isabstractmethod__ = self.__isabstractmethod__
    378         _method._partialmethod = self

[/usr/local/lib/python3.7/dist-packages/jax/_src/device_array.py](https://localhost:8080/#) in _forward_method(attrname, self, fun, *args)
     39 
     40 def _forward_method(attrname, self, fun, *args):
---&gt; 41   return fun(getattr(self, attrname), *args)
     42 _forward_to_value = partial(_forward_method, &quot;_value&quot;)
     43 

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()"><pre class="notranslate"><code class="notranslate">True
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[&lt;ipython-input-7-f0aa3716322d&gt;](https://localhost:8080/#) in &lt;module&gt;()
     10 array, unravel_fn = ravel_pytree(pytree)
     11 reconstructed_pytree2 = unravel_fn(array)
---&gt; 12 print(reconstructed_pytree2 == pytree)

1 frames
[/usr/lib/python3.7/functools.py](https://localhost:8080/#) in _method(*args, **keywords)
    374             cls_or_self, *rest = args
    375             call_args = (cls_or_self,) + self.args + tuple(rest)
--&gt; 376             return self.func(*call_args, **call_keywords)
    377         _method.__isabstractmethod__ = self.__isabstractmethod__
    378         _method._partialmethod = self

[/usr/local/lib/python3.7/dist-packages/jax/_src/device_array.py](https://localhost:8080/#) in _forward_method(attrname, self, fun, *args)
     39 
     40 def _forward_method(attrname, self, fun, *args):
---&gt; 41   return fun(getattr(self, attrname), *args)
     42 _forward_to_value = partial(_forward_method, "_value")
     43 

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
</code></pre></div>
<h2 dir="auto">Description</h2>
<p dir="auto">I thought that reconstructed trees should be the same with both <code class="notranslate">tree_flatten</code> and <code class="notranslate">ravel_pytree</code>. But something is going wrong here. I would appreciate any help.</p>
<p dir="auto">Consider the following function which extends <code class="notranslate">jax.scipy.ndimage.map_coordinates</code> to work with arbitrary shaped lattices of coordinates, and an arbitrary number of value channels (including multiple value dimensions).</p>
<p dir="auto">Because <code class="notranslate">jax.scipy.ndimage.map_coordinates</code> operates on single channel value lattices, this has been implemented by using <code class="notranslate">jax.vmap</code> distributed over each flattened value channel.</p>
<p dir="auto">Naively this would be quite inefficient, as the computation for every channel would duplicate computing indices and the splines in each dimension to mix the values.</p>
<p dir="auto">Imagine an extrem(ish) case of 2D interpolation between a 10x10 set of volumes 16x16x1024. The value lattice would have size (10, 10, 16, 16, 1024) with a value flattened size of (10, 10, 262144), vmapping over the last dimension.</p>
<p dir="auto">Is there any way to verify how sensibly this function is being vmapped?</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="def map_lattice_coordinates(value_lattice, point_lattice, order=1, mode='nearest', **kargs):
    &quot;&quot;&quot;
    Perform an interpolated lookup of the values in the value lattice using the coordinates in the coordinate lattice.

    Assume a 2D lattice of coordinates in 3-space indexing into a 3D lattice of 10x10 matrices. The result will be a 
    2D lattice of 10x10 matrices sampled by Tri-Linear interpolation. i.e. 

        (point lattice : (5, 5, 3), value lattice : (2, 2, 2, 10, 10)) -&gt; result : (5, 5, 10, 10)
    &quot;&quot;&quot;

    # The point lattice last dimension holds vector coordinates into the leading dimensions of the value lattice
    p_shape    = point_lattice.shape[:-1]  # The result will have this shape for its leading dimensions
    p_channels = point_lattice.shape[-1]

    # Collapse multiple shape dims down to one dim, and transpose result into list of vectors
    point_lattice = jnp.reshape(point_lattice, (-1, p_channels)).T

    # The coordinates in the point lattice will index into the leading dimensions of the value lattice
    v_shape    = value_lattice.shape[:p_channels]
    v_channels = value_lattice.shape[p_channels:]  # The result will have this shape for its trailing dimensions

    # Collapse multiple channel dims to one dim
    value_lattice = jnp.reshape(value_lattice, (*v_shape, -1))

    def interp_channel(channel):
        return jax.scipy.ndimage.map_coordinates(channel, point_lattice, order=order, mode=mode, **kargs)

    # Interpolate each channel independently (in parallel) and stack the results together
    interp_values = jax.vmap(interp_channel, in_axes=-1, out_axes=-1)(value_lattice)

    # Reshape the results to have the spatial dims from the point lattice and the value dims from the value lattice
    return jnp.reshape(interp_values, (*p_shape, *v_channels))"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">map_lattice_coordinates</span>(<span class="pl-s1">value_lattice</span>, <span class="pl-s1">point_lattice</span>, <span class="pl-s1">order</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">mode</span><span class="pl-c1">=</span><span class="pl-s">'nearest'</span>, <span class="pl-c1">**</span><span class="pl-s1">kargs</span>):
    <span class="pl-s">"""</span>
<span class="pl-s">    Perform an interpolated lookup of the values in the value lattice using the coordinates in the coordinate lattice.</span>
<span class="pl-s"></span>
<span class="pl-s">    Assume a 2D lattice of coordinates in 3-space indexing into a 3D lattice of 10x10 matrices. The result will be a </span>
<span class="pl-s">    2D lattice of 10x10 matrices sampled by Tri-Linear interpolation. i.e. </span>
<span class="pl-s"></span>
<span class="pl-s">        (point lattice : (5, 5, 3), value lattice : (2, 2, 2, 10, 10)) -&gt; result : (5, 5, 10, 10)</span>
<span class="pl-s">    """</span>

    <span class="pl-c"># The point lattice last dimension holds vector coordinates into the leading dimensions of the value lattice</span>
    <span class="pl-s1">p_shape</span>    <span class="pl-c1">=</span> <span class="pl-s1">point_lattice</span>.<span class="pl-s1">shape</span>[:<span class="pl-c1">-</span><span class="pl-c1">1</span>]  <span class="pl-c"># The result will have this shape for its leading dimensions</span>
    <span class="pl-s1">p_channels</span> <span class="pl-c1">=</span> <span class="pl-s1">point_lattice</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>]

    <span class="pl-c"># Collapse multiple shape dims down to one dim, and transpose result into list of vectors</span>
    <span class="pl-s1">point_lattice</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">reshape</span>(<span class="pl-s1">point_lattice</span>, (<span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">p_channels</span>)).<span class="pl-v">T</span>

    <span class="pl-c"># The coordinates in the point lattice will index into the leading dimensions of the value lattice</span>
    <span class="pl-s1">v_shape</span>    <span class="pl-c1">=</span> <span class="pl-s1">value_lattice</span>.<span class="pl-s1">shape</span>[:<span class="pl-s1">p_channels</span>]
    <span class="pl-s1">v_channels</span> <span class="pl-c1">=</span> <span class="pl-s1">value_lattice</span>.<span class="pl-s1">shape</span>[<span class="pl-s1">p_channels</span>:]  <span class="pl-c"># The result will have this shape for its trailing dimensions</span>

    <span class="pl-c"># Collapse multiple channel dims to one dim</span>
    <span class="pl-s1">value_lattice</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">reshape</span>(<span class="pl-s1">value_lattice</span>, (<span class="pl-c1">*</span><span class="pl-s1">v_shape</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>))

    <span class="pl-k">def</span> <span class="pl-en">interp_channel</span>(<span class="pl-s1">channel</span>):
        <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">scipy</span>.<span class="pl-s1">ndimage</span>.<span class="pl-en">map_coordinates</span>(<span class="pl-s1">channel</span>, <span class="pl-s1">point_lattice</span>, <span class="pl-s1">order</span><span class="pl-c1">=</span><span class="pl-s1">order</span>, <span class="pl-s1">mode</span><span class="pl-c1">=</span><span class="pl-s1">mode</span>, <span class="pl-c1">**</span><span class="pl-s1">kargs</span>)

    <span class="pl-c"># Interpolate each channel independently (in parallel) and stack the results together</span>
    <span class="pl-s1">interp_values</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">vmap</span>(<span class="pl-s1">interp_channel</span>, <span class="pl-s1">in_axes</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">out_axes</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)(<span class="pl-s1">value_lattice</span>)

    <span class="pl-c"># Reshape the results to have the spatial dims from the point lattice and the value dims from the value lattice</span>
    <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">reshape</span>(<span class="pl-s1">interp_values</span>, (<span class="pl-c1">*</span><span class="pl-s1">p_shape</span>, <span class="pl-c1">*</span><span class="pl-s1">v_channels</span>))</pre></div>
<p dir="auto">Example usage (point lattice : (2, 2, 2), value lattice : (3, 3, 2, 2)) -&gt; result : (2, 2, 2, 2)</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="value_lattice = jnp.array([
    [[[0, 1], [2, 3]], [[1, 2], [3, 4]], [[2, 3], [4, 5]]],
    [[[1, 2], [3, 4]], [[2, 3], [4, 5]], [[3, 4], [5, 6]]],
    [[[2, 3], [4, 5]], [[3, 4], [5, 6]], [[4, 5], [6, 7]]]
])

point_lattice = jnp.array([
    [[0.5, 0.5], [0.5, 1.5]],
    [[1.5, 0.5], [1.5, 1.5]]
])

expected = jnp.array([
    [[[1, 2], [3, 4]], [[2, 3], [4, 5]]],
    [[[2, 3], [4, 5]], [[3, 4], [5, 6]]]
])

assert np.allclose(expected, map_lattice_coordinates(value_lattice, point_lattice), atol=1e-5)"><pre class="notranslate"><span class="pl-s1">value_lattice</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([
    [[[<span class="pl-c1">0</span>, <span class="pl-c1">1</span>], [<span class="pl-c1">2</span>, <span class="pl-c1">3</span>]], [[<span class="pl-c1">1</span>, <span class="pl-c1">2</span>], [<span class="pl-c1">3</span>, <span class="pl-c1">4</span>]], [[<span class="pl-c1">2</span>, <span class="pl-c1">3</span>], [<span class="pl-c1">4</span>, <span class="pl-c1">5</span>]]],
    [[[<span class="pl-c1">1</span>, <span class="pl-c1">2</span>], [<span class="pl-c1">3</span>, <span class="pl-c1">4</span>]], [[<span class="pl-c1">2</span>, <span class="pl-c1">3</span>], [<span class="pl-c1">4</span>, <span class="pl-c1">5</span>]], [[<span class="pl-c1">3</span>, <span class="pl-c1">4</span>], [<span class="pl-c1">5</span>, <span class="pl-c1">6</span>]]],
    [[[<span class="pl-c1">2</span>, <span class="pl-c1">3</span>], [<span class="pl-c1">4</span>, <span class="pl-c1">5</span>]], [[<span class="pl-c1">3</span>, <span class="pl-c1">4</span>], [<span class="pl-c1">5</span>, <span class="pl-c1">6</span>]], [[<span class="pl-c1">4</span>, <span class="pl-c1">5</span>], [<span class="pl-c1">6</span>, <span class="pl-c1">7</span>]]]
])

<span class="pl-s1">point_lattice</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([
    [[<span class="pl-c1">0.5</span>, <span class="pl-c1">0.5</span>], [<span class="pl-c1">0.5</span>, <span class="pl-c1">1.5</span>]],
    [[<span class="pl-c1">1.5</span>, <span class="pl-c1">0.5</span>], [<span class="pl-c1">1.5</span>, <span class="pl-c1">1.5</span>]]
])

<span class="pl-s1">expected</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([
    [[[<span class="pl-c1">1</span>, <span class="pl-c1">2</span>], [<span class="pl-c1">3</span>, <span class="pl-c1">4</span>]], [[<span class="pl-c1">2</span>, <span class="pl-c1">3</span>], [<span class="pl-c1">4</span>, <span class="pl-c1">5</span>]]],
    [[[<span class="pl-c1">2</span>, <span class="pl-c1">3</span>], [<span class="pl-c1">4</span>, <span class="pl-c1">5</span>]], [[<span class="pl-c1">3</span>, <span class="pl-c1">4</span>], [<span class="pl-c1">5</span>, <span class="pl-c1">6</span>]]]
])

<span class="pl-k">assert</span> <span class="pl-s1">np</span>.<span class="pl-en">allclose</span>(<span class="pl-s1">expected</span>, <span class="pl-en">map_lattice_coordinates</span>(<span class="pl-s1">value_lattice</span>, <span class="pl-s1">point_lattice</span>), <span class="pl-s1">atol</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>)</pre></div>