Given a state, how to get the values of its arrays?

[Artur A. Galstyan]

As the title says. E.g. I do this:

state, model = eqx.nn.make_with_state(...)(...)

# how to get the arrays within state?

[Artur A. Galstyan]

Oh, my bad. It's just:

state.get(model.state_index) # or whatever is the name of the index