Skip to content

Pull requests: google/flax

New pull request New
158 Open 2,603 Closed
158 Open 2,603 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Add nnx.Module.perturb
#4515 opened Jan 30, 2025 by IvyZX Loading…
[nnx] deepseek example
#4512 opened Jan 29, 2025 by cgarciae • Draft
deepseek example
#4511 opened Jan 29, 2025 by copybara-service bot Loading…
make gemma FFW LoRA friendly
#4510 opened Jan 29, 2025 by copybara-service bot Loading…
1
[nnx] add tabulate pull ready
#4493 opened Jan 20, 2025 by cgarciae Loading…
Improve error message for when nnx.Modules use jax or numpy arrays as leaf values
#4492 opened Jan 20, 2025 by RaghuSpaceRajan Loading…
1
4
[nnx] add shard_map
#4490 opened Jan 19, 2025 by cgarciae Loading…
1
[nnx] improve custom_vjp ergonomics
#4489 opened Jan 19, 2025 by cgarciae Loading…
1
[linen] Linesearch (and lbfgs) support for TrainState
#4471 opened Jan 6, 2025 by emiresenov Loading…
3 of 4 tasks
8
[nnx] Add specific model typing for nnx.Optimizer
#4470 opened Jan 1, 2025 by marcelroed Loading…
1 task done
6
[nnx] add cached_partial
#4469 opened Jan 1, 2025 by cgarciae Loading…
7
[nnx] cache flatten
#4463 opened Dec 23, 2024 by cgarciae • Draft
[nnx] simplify unflatten
#4462 opened Dec 23, 2024 by cgarciae • Draft
[nnx] flatten returns FlatState
#4458 opened Dec 21, 2024 by cgarciae • Draft
Update NNX State.state.merge method docs in statelib.py
#4453 opened Dec 17, 2024 by 8bitmp3 Loading…
@8bitmp3
1
Update NNX State.state.filter method docs in statelib.py
#4452 opened Dec 17, 2024 by 8bitmp3 Loading…
@8bitmp3
1
Update NNX State.state.split method docs in statelib.py
#4451 opened Dec 17, 2024 by 8bitmp3 Loading…
@8bitmp3
1
Update NNX State class docs in statelib.py
#4450 opened Dec 17, 2024 by 8bitmp3 Loading…
@8bitmp3
2
Lint NNX BatchStat, Cache, Intermediate class docs in variablelib.py
#4449 opened Dec 17, 2024 by 8bitmp3 Loading…
@8bitmp3
Update NNX Variable class docs in variablelib.py
#4448 opened Dec 17, 2024 by 8bitmp3 Loading…
@8bitmp3
Lint NNX Module iter_* method docs in module.py
#4447 opened Dec 17, 2024 by 8bitmp3 Loading…
@8bitmp3
Update NNX Module eval docs in module.py
#4446 opened Dec 16, 2024 by 8bitmp3 Loading…
@8bitmp3
Update NNX Module train docs in module.py
#4445 opened Dec 16, 2024 by 8bitmp3 Loading…
@8bitmp3
Update NNX Module set_attributes docs in module.py
#4444 opened Dec 16, 2024 by 8bitmp3 Loading…
@8bitmp3
Upgrade NNX Module sow docs in module.py
#4443 opened Dec 16, 2024 by 8bitmp3 Loading…
@8bitmp3
ProTip! Exclude everything labeled bug with -label:bug.