flax
replace lazy_init implementation to avoid jax partial_eval
#5551

Merged

replace lazy_init implementation to avoid jax partial_eval #5551

copybara-service merged 2 commits into google:main from mattjj:replace-lazy-init

dce-based solution

00bd0764

dummy-zeros based implementation

f3b0e126

IvyZX approved these changes on 2026-08-10

IvyZX added pull ready

copybara-service merged abef9d81 into main 35 days ago

mattjj deleted the replace-lazy-init branch 35 days ago

Reviewers

IvyZX

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone