use State instead of GraphState#5401

copybara-service[bot] wants to merge 1 commit into

maingoogle/flax:mainfrom

test_896115256google/flax:test_896115256Copy head branch name to clipboard