Skip to content

Pull requests: google/flax

New pull request New
127 Open 2,547 Closed
127 Open 2,547 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 some optimizations to graph.py
#4377 opened Nov 13, 2024 by cgarciae Loading…
Add flax.nnx.value_and_grad docstring
#4372 opened Nov 11, 2024 by 8bitmp3 Loading…
@levskaya @8bitmp3
5
Lint flax.nnx.while_loop docstring pull ready
#4371 opened Nov 11, 2024 by 8bitmp3 Loading…
@8bitmp3
1
Lint flax.nnx.fori_loop docstring
#4370 opened Nov 11, 2024 by 8bitmp3 Loading…
@8bitmp3
Update flax.nnx.jit docstring
#4369 opened Nov 11, 2024 by 8bitmp3 Loading…
@8bitmp3
Update flax.nnx.grad docstring
#4368 opened Nov 11, 2024 by 8bitmp3 Loading…
@8bitmp3
Fix multiple links to Orbax documentation
#4364 opened Nov 8, 2024 by Matt-Hurd Loading…
4 tasks done
1
[nnx] add unsafe_graph_cache decorator
#4356 opened Nov 4, 2024 by cgarciae Loading…
[nnx] Optimizer uses variables
#4347 opened Oct 30, 2024 by cgarciae Loading…
1
Upgrade Flax NNX Gemma Sampling Inference doc
#4325 opened Oct 23, 2024 by 8bitmp3 Loading…
@8bitmp3
18
Avoid depending on JAX internals, which are about to change (see cl/681582933).
#4321 opened Oct 23, 2024 by copybara-service bot Loading…
Fixing numbering in mnist_tutorial.md
#4311 opened Oct 18, 2024 by v1vekkumar Loading…
1 of 4 tasks
@8bitmp3
5
[nnx] flaxlib flatten
#4293 opened Oct 13, 2024 by cgarciae Loading…
Update Flax NNX vs JAX Transformations guide
#4286 opened Oct 10, 2024 by 8bitmp3 Loading…
@8bitmp3
Fix Philosophy links, update copyright year,
#4283 opened Oct 10, 2024 by 8bitmp3 • Draft
@8bitmp3
4
Update Flax NNX Randomness
#4279 opened Oct 10, 2024 by 8bitmp3 Loading…
@8bitmp3
4
Stackless yashful
#4246 opened Oct 3, 2024 by copybara-service bot Loading…
proxy more methods for nnx.Variable
#4234 opened Sep 29, 2024 by njzjz Loading…
1 of 4 tasks
5
Utils for fp8 gated einsum
#4221 opened Sep 24, 2024 by wenscarl Loading…
Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change.
#4219 opened Sep 23, 2024 by copybara-service bot Loading…
Upgrade Flax NNX Filters doc
#4199 opened Sep 16, 2024 by 8bitmp3 Loading…
@cgarciae @8bitmp3
4
Upgrade Flax NNX Bridge API doc
#4198 opened Sep 16, 2024 by 8bitmp3 Loading…
@8bitmp3
12
[WIP] Require transform_metadata when variables have sharding annotation
#4187 opened Sep 10, 2024 by IvyZX • Draft
Add miniGPT-Jax to community Examples.
#4184 opened Sep 10, 2024 by yilunyu Loading…
1 of 4 tasks
1
1
[nnx] add nnx.pytree
#4174 opened Sep 5, 2024 by cgarciae Loading…
ProTip! Follow long discussions with comments:>50.