How to resolve pytrees issues when using a Flax optimizer with a Haiku model?

[marcvanzee]

Original question by @j5b:

Does a Flax optim optimizer to work with a Haiku model (and its parameter representation)? I'm running into pytree issues and I'm wondering if there's some incompatibility that I didn't consider.

Here's code reproducing the issue:

import haiku as hk
import jax
from jax import numpy as jnp
from flax import optim

class SomeModule(hk.Module):
  def __call__(self, x):
    w = hk.get_parameter("w", [], init=jnp.zeros)
    return x + w

def loss_fn(params):
  return f.apply(params, None, 1)

f = hk.transform(lambda x: SomeModule()(x))

# Get params from Haiku
params = f.init(None, 1)

# Create Flax optimizer from Haiku parameters.
optimizer = optim.Adam(learning_rate=1.0).create(params)

grad = jax.grad(loss_fn)(optimizer.target)
optimizer.apply_gradient(grad)

This gives the following error:

ValueError: Mismatch custom node data: PyTreeDef(dict[['some_module']], [PyTreeDef(haiku._src.data_structures.FlatMapping[PyTreeDef(dict[['w']], [*])], [*])]) != PyTreeDef(dict[['some_module']], [PyTreeDef(haiku._src.data_structures.FlatMapping[PyTreeDef(dict[['w']], [PyTreeDef(<class 'flax.optim.adam._AdamParamState'>[()], [*,*])])], [*,*])]); value: FlatMapping({
  'some_module': FlatMapping({
                   'w': _AdamParamState(grad_ema=DeviceArray(0., dtype=float32), grad_sq_ema=DeviceArray(0., dtype=float32)),
                 }),
}).

[marcvanzee]

Answer by @jheek:

This is actually a Haiku bug. The FlatMap implementation breaks the assumptions of pytrees.

Here is an example that doesn't use Flax at all:

double_params = jax.tree_map(lambda x: (x, x), params)
treedef = jax.tree_structure(params)
treedef.flatten_up_to(double_params)

flatten_up_to simply doesn't work with the pytree implementation of FlatMap.

You can flatten the representation manually by using flat_params, treedef = jax.tree_flatten(params) and then you use treedef.unflatten(flat_params) before passing to Haiku.

So your code would look like this:

...
def loss(params):
  params = treedef.unflatten(optimizer.target)
  return f.apply(params, None, 1)

...

params = f.init(None, 1)
flat_params, treedef = jax.tree_flatten(params)

optimizer_def = optim.Adam(learning_rate=1.0)
optimizer=optimizer_def.create(params)

[PhilipVinc]

I recently stumbled upon that same issue.
Do you plan on opening an issue over at deepmind/dm-haiku, or can I do it?

[jheek]

I have had contact with the Haiku owners. They are working on a solution.
It's not trivial to fix because their FlatMapping implementation gives a performance boost when you need really low overhead jit calls (currently, at the cost of not fully respecting the tree_util api).

[jheek]

Still it would be good if you open an issue to make clear more people are running into this and so we have something to link to :)