Dict to NNX state

[thomaspinder]

I am currently integrating GPJax with Numpyro. The output of Numpyro's inference objects is a sequence of samples from the model's parameters, and they are represented as a dictionary. Given I know the structure of my model's PyGraph, is there any functionality already in NNX that would allow me to generate a new State object from a dictionary.

For example, if I have the state:

State({
  'kernel': {
    'lengthscale': VariableState( # 1 (8 B)
      type=PositiveReal,
      value=Array(1., dtype=float64, weak_type=True),
      _tag='positive'
    ),
    'variance': VariableState( # 1 (8 B)
      type=PositiveReal,
      value=Array(1., dtype=float64, weak_type=True),
      _tag='positive'
    )
  },
  'mean_function': {
    'constant': VariableState( # 1 (8 B)
      type=Static,
      value=Array(0., dtype=float64, weak_type=True),
      tag='static',
      _tag='static'
    )
  }
})

then is there a generic way for me to map a dictionary of the form below back into an instance of this State?

{"kernel.lengthscale": jnp.array(1.2), "kernel.variance": jnp.array(0.8), "mean_function.constant": 4.2}

[cgarciae]

Hey @thomaspinder,

You can create a State object simply by passing a dict to its constructor e.g:

state = nnx.State({'a': jnp.array(1.0)})

State exists for historical reasons and might get deleted in the future in favor of just dict. Unless I'm missing something, the problem here is mostly how create a nested structure from a flat structure? Check out nnx.traversals which has the unflatten_mapping helper for this.