Feature request: functionality similar to torch.nn.utils.vector_to_parameters

[timdhondt1]

TLDR: Add jittable functionality similar to Pytorch's torch.nn.utils.vector_to_parameters, because this is useful for Bayesian Neural Networks.

---

I have been spending some time experimenting with Bayesian Neural Networks with Jax and Flax and a big part of BNN's is the fact that one samples parameters from some distribution and subsequently uses a NN with these sampled parameters. The raw parameters are usually sampled as a 1d DeviceArray object and then have to be reformatted into a flax.core.frozen_dict.FrozenDict object, preferably in jittable fashion.

I have some code that is able to do this while being jittable, through some hacky tricks avoiding jax.errors.ConcretizationTypeError-type errors when splitting up the 1d vector into the different parameter groups. However, this code keeps on breaking whenever a new version of jax is released, probably due to under the hood changes in jax.core or lax.numpy. An example is given below

def vector_2_tree(vector, example_tree, group_ranges, group_shapes):
	'''
	Formats a 1d parameter vector into a flax.core.frozen_dict.FrozenDict
	object compatible with a flax.linen.Module

	Args:
		vector: sampled 1d parameter vector
		example_tree: example parameter tree compatible with flax.linen.Module
					  object
		group_ranges: list of indices where 1d vector should be split up into
		               different parameter groups
		group_shapes: list of tuples that describe parameter group shapes

	'''

	# split up sampled vector into parameter groups and reshape them accordingly
	reshaped_vector = tree_util.tree_multimap(lambda x,y: x.reshape(y), jnp.split(vector, group_ranges), group_shapes)
	# put the reshaped vectors into a dict with the proper parameter keys
	flattened_tree = {k: reshaped_vector[i] for i, (k,v) in enumerate(example_tree.items())}
	# renest dict
	return traverse_util.unflatten_dict({tuple(k.split('/')): v for k, v in flattened_tree.items()})

With the newest version of Jax and Flax, jitting this function fails on the first line due to a jax.errors.ConcretizationTypeError in the jnp.split call. It used to work with older versions.

Supporting jittable functionality equivalent to Pytorch's torch.nn.utils.vector_to_parameters in Flax can be useful for Bayesian Deep Learning use cases. This was mentioned as an issue before here but never solved to my knowledge

Cheers,
Tim

[jheek]

This API should do exactly what you want: https://jax.readthedocs.io/en/latest/_autosummary/jax.flatten_util.ravel_pytree.html

[timdhondt1]

@jheek thanks! I had never seen that this function also returned a callable.