patrick-kidger / equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Apache License 2.0
2.13k stars 143 forks source link

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

Closed Artur-Galstyan closed 5 months ago

Artur-Galstyan commented 5 months ago

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


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

# how to get the arrays within state?
Artur-Galstyan commented 5 months ago

Oh, my bad. It's just:

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