How to `vmap` over a batch of trees or dictionaries?

Issue I want to vmap a function over a batch of trees or dictionaries. However, I am not sure what the appropriate way to pass such a batch is (see the example below).

Related #2367 dealt with passing dictionaries to vmaped functions that do not vmap over dictionary entries. #3161 dealt with vmaping over dictionary entries. This issue is different because it deals with vmaping over a batch of dictionaries.

Background I want to vmap a haiku NN execution over a batch of different parameters. Below is an example that deals with the same issue.

Minimal example Case 0 and 1 work as expected, 2 and 3 do not.

from jax import vmap
import jax.numpy as np

def f(x, params):
  return params['a']*x + params['b']

print('\nCase 0:')
print(f(2, dict(a=3, b=4)))

print('\nCase 1:')
print(vmap(f, in_axes=(0, None))(
    np.array([2, 2]), 
    dict(a=3, b=4)
))

print('\nCase 2:')
try:
  print(vmap(f, in_axes=(0, 0))(
      np.array([2, 2]),
      [dict(a=3, b=4), dict(a=3, b=4)]
  ))
except Exception as e:
  print(e)

print('\nCase 3:')
try:
  print(vmap(f, in_axes=(0, 0))(
      np.array([2, 2]),
      np.array([dict(a=3, b=4), dict(a=3, b=4)])
  ))
except Exception as e:
  print(e)

Output

Case 0:
10

Case 1:
[10 10]

Case 2:
vmap got arg 1 of rank 0 but axis to be mapped 0. The tree of ranks is:
((1, [{'a': 0, 'b': 0}, {'a': 0, 'b': 0}]), {})

Case 3:
JAX only supports number and bool dtypes, got dtype object in array

[tomhennigan]

Hi @cisprague , utilities like jax.tree_map(..) often appear in user code, so I would personally (although I am not a JAX core team member) consider this part of the public API. I suspect the docs discussing how pytrees work (e.g. the treedef etc) are what are intended to be “implementation details” which you should only consider if developing extensions (new tree types).

I’d be happy for this function to exist in a JAX library though, to avoid copy pasting, but I’m not sure Haiku is the right place (this function is not unique to Haiku and would benefit users of other JAX libraries). I’m not aware of any libraries with JAX recipes/snippets, but this would be at home there.

[cisprague]

@tomhennigan thanks for the example and Colab notebook! This seems like a reasonable solution. Maybe this could be a new end-user feature in Haiku, since pytrees is intended for JAX internal users.