flax
Have _graph_flatten respect nnx.data declarations (extension of #5140)
#5159

Merged

Have _graph_flatten respect nnx.data declarations (extension of #5140) #5159

chapman20j merged 2 commits into google:main from samanklesaria:pr/5140

flatten respect nnx.data

750527aa

cgarciae force pushed from 5e977f00 to d73d2424 50 days ago

Fix type error

ebd78685

cgarciae force pushed from d73d2424 to ebd78685 50 days ago

chapman20j approved these changes on 2026-01-05

chapman20j merged 4cc1d1bb into main 50 days ago

Reviewers

chapman20j

Assignees

No one assigned

Labels

None yet

Milestone

No milestone