WIP: tree vectorizing transformation

[shoyer]

Fixes #1012

The idea behind this transformation is that every argument to a function decorated by @tree_vectorize is virtually flattened and concatenated into 1D vector, e.g., {'x': 1, 'y': array([[2, 3]])} becomes array([1, 2, 3]).

However, under the covers the original arrays are still preserved with their original shapes. This has several advantages:

1. we don't actually need to copy the arrays into a giant vector (avoiding the unnecessary memory copies)
2. all operations happen on arrays with the original shapes (useful for predictable performance on TPUs)
3. we can restore the tree structure on the outputs of the function

The hope is that this should make it considerably easier to write efficient numerical algorithms that work on pytrees (e.g., for optimization, linear equation solving or ODE solving).

Here are a few simple examples:

from jax import tree_vectorize
import jax.numpy as jnp

x = {'a': 1, 'b': jnp.array([2, 3])}
y = 1
add = tree_vectorize(jnp.add)
print(add(x, y))
# {'a': DeviceArray(2, dtype=int32), 'b': DeviceArray([3, 4], dtype=int32)}

@tree_vectorize
def norm(x):
  assert x.shape == (3,)
  return jnp.sqrt(jnp.sum(x ** 2))

print(norm(x))
# 3.7416575

Even though the inputs are always treated as vectors, outputs can have any number of dimensions, e.g., producing nested pytrees-of-pytrees, as shown in this example from the unit tests:

@tree_vectorize
@shapecheck(['n', 'm'], '(n, m)')  # just to show that it works!
def add_outer(x, y):
  return jnp.expand_dims(x, 1) + jnp.expand_dims(y, 0)

tree1 = {'a': 1, 'b': jnp.array([2, 3])}
tree2 = {'c': 10, 'd': jnp.array([20, 30])}
pprint.pprint(add_outer(tree1, tree2))
# {'a': {'c': DeviceArray(11, dtype=int32),
#        'd': DeviceArray([21, 31], dtype=int32)},
#  'b': {'c': DeviceArray([12, 13], dtype=int32),
#        'd': DeviceArray([[22, 32],
#              [23, 33]], dtype=int32)}}

TODO:

• missing ops for cg:
  • jit
  • while_loop
• further missing ops for odeint:
  • concatenate
  • gather
  • scatter
• measure performance improvement for odeint!
• figure out the API, e.g.,
  • explicit API for indicating functions, especially with multiple outputs?
  • explicit API for indicating non-vectorized arguments?
• make things more robust
  • tree_call works with vmap, jvp and vjp (i.e., to implement jacfwd and jacrev)
  • missing Trace methods like process_call, post_process_call
  • support matching against explicit axes with trivial treedefs, e.g., so x + jnp.zeros(...) and explicit broadcasting work.
• more comprehensive tests
  • check exceptions
  • make a reference implementation for a subset of tree_vectorize use cases that works by explicitly unraveling/concatenating

[shoyer]

I think this is very close to working, except we don't handle basis properly yet because it has a different tree structure. (This works in other cases, but not for transformed tree_call functions yet.)

[shoyer]

Gradient support needs JaxprTrace.process_tree_call (which seems straightforward, but tedious).

[shoyer]

@apaszke This is starting to feel pretty repetitive, especially when we consider doing the same thing for every higher order fucntions/transformation in JAX (e.g., partial eval, scan, while_loop, etc). It feels like there must be a better way to leverage existing support for wrapping trees...

[shoyer]

Closing this in favor of #8504