Merge FlatMappings with DeviceArray as values raises `AttributeError: 'DeviceArray' object has no attribute 'items'`

[google-deepmind]

Merge FlatMappings with DeviceArray as values raises `AttributeError: 'DeviceArray' object has no attribute 'items'`

[tomhennigan]

Hey @sebamenabar , the utilities in data_structures rely on being passed a "two level" mapping (in Haiku the params and state dicts have the following structure: {'module_name': {'parameter_name': parameter_value}}), when you do fast = params['fast'] you are getting one of the inner mappings which cannot be used with functions expecting a two level mapping.

There is an easy solution though 😄 I would suggest using partition and merge to split and combine the params dict while maintaining two level mappings:

params = f.init(..)

# Split fast and slow params to operate on them independently.
predicate = lambda m, n, v: m == 'fast_weights'
fast_params, slow_params = hk.data_structures.partition(predicate, params)

# Recombine fast and slow to (for example) pass to apply.
params = hk.data_structures.merge(fast_params, slow_params)
out = f.apply(params, ..)

Fabio added an example utility function (jax_fn_with_filter) to the tests [0] when he authored these, you may find that useful too.

[0] https://github.com/deepmind/dm-haiku/blob/master/haiku/_src/filtering_test.py

[sebamenabar]

Thanks tom!