We use state for what init returns, so should probably never use it for what setup returns. That's called tree in code, and this calls it state_tree more places in the docs.
Unsure if we should change headings like Optimisers.update!(tree, model, gradient) -> (tree, model) to the more verbose Optimisers.update!(state_tree, model, gradient) -> (state_tree, model).
Another possibility for docs is to call the whole thing states plural.
Doc-only change.
We use
statefor whatinitreturns, so should probably never use it for whatsetupreturns. That's calledtreein code, and this calls itstate_treemore places in the docs.Unsure if we should change headings like
Optimisers.update!(tree, model, gradient) -> (tree, model)to the more verboseOptimisers.update!(state_tree, model, gradient) -> (state_tree, model).Another possibility for docs is to call the whole thing
statesplural.