Skip to content

Pull requests: google/flax

New pull request New
153 Open 2,595 Closed
153 Open 2,595 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

[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
Remove usages of orbax_utils.save_args_from_target, as this function does nothing (it used to control a checkpointing behavior that has since been optimized away).
#4482 opened Jan 15, 2025 by copybara-service bot Loading…
[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
5
[nnx] add cache_args
#4469 opened Jan 1, 2025 by cgarciae Loading…
1
[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
Update NNX Module class docs in module.py
#4442 opened Dec 16, 2024 by 8bitmp3 Loading…
@8bitmp3
Update NNX state docs in graph.py
#4441 opened Dec 16, 2024 by 8bitmp3 Loading…
@8bitmp3
Lint NNX pop method docs in graph.py
#4440 opened Dec 16, 2024 by 8bitmp3 Loading…
@8bitmp3
Update NNX call method docs in graph.py
#4439 opened Dec 15, 2024 by 8bitmp3 Loading…
@8bitmp3
Fix de-serialisation of numpy fixed-width dtypes.
#4429 opened Dec 11, 2024 by copybara-service bot Loading…
1
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.