NNX migration prep (1/N): pure_nnx flag and init_state_fn scaffolding#3427

copybara-service[bot] merged 1 commit intomainAI-Hypercomputer/maxtext:mainfrom

feat/pure_nnx_flag_and_init_state_fnAI-Hypercomputer/maxtext:feat/pure_nnx_flag_and_init_state_fnCopy head branch name to clipboard