Issue 1:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

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)

If applicable, include full error messages/tracebacks.

True
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
[<ipython-input-7-f0aa3716322d>](https://localhost:8080/#) in <module>()
10 array, unravel_fn = ravel_pytree(pytree)
11 reconstructed_pytree2 = unravel_fn(array)
---> 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)
--> 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):
---> 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()

Description

I thought that reconstructed trees should be the same with both tree_flatten and ravel_pytree. But something is going wrong here. I would appreciate any help.

Issue 2:
Consider the following function which extends jax.scipy.ndimage.map_coordinates to work with arbitrary shaped lattices of coordinates, and an arbitrary number of value channels (including multiple value dimensions).

Because jax.scipy.ndimage.map_coordinates operates on single channel value lattices, this has been implemented by using jax.vmap distributed over each flattened value channel.

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.

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.

Is there any way to verify how sensibly this function is being vmapped?

def map_lattice_coordinates(value_lattice, point_lattice, order=1, mode='nearest', **kargs):
"""
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)) -> result : (5, 5, 10, 10)
"""

# 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))

Example usage (point lattice : (2, 2, 2), value lattice : (3, 3, 2, 2)) -> result : (2, 2, 2, 2)

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)