Skip to content

Issues: google/flax

Beta
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

213 Open 784 Closed
213 Open 784 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
Assignee
Filter by who’s assigned
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Issues list

Improve Gemma example
#4740 opened May 7, 2025 by cgarciae
@vfdev-5
Unable to use jax.debug.breakpoint in nnx jitted scope
#4735 opened Apr 29, 2025 by puct9
3
Cannot nnx.eval_shape Modules with BatchNorm
#4734 opened Apr 29, 2025 by LarsKue
1
How to train an NNX netowork using scan and jit?
#4722 opened Apr 20, 2025 by cdelv
8
No SpectralNorm for NNX
#4684 opened Apr 4, 2025 by h-0-0
1
2
Using flax.linen.intercept_methods on an nnx module
#4666 opened Mar 29, 2025 by gardberg
5
Filtering with nnx.PathContains('') returns an empty state
#4660 opened Mar 27, 2025 by NKlug
1
Distribute inference on a TPU
#4644 opened Mar 24, 2025 by zaidalyafeai
1
Problem training model with jax.jvp in forward pass
#4636 opened Mar 18, 2025 by Thomas-Christie
2
Stacking layers with vmap and forwarding with scan result in loss of precision in XLA backend
#4629 opened Mar 16, 2025 by HeavyCrab
2
Gradient Checkpointing causes model to compute junk results (NNX)
#4626 opened Mar 15, 2025 by erfanzar
6
What is the status of nnx and TrainState? Are they compatible?
#4609 opened Mar 10, 2025 by samuela
@cgarciae
5
Bug in MinimalGatedUnit
#4608 opened Mar 9, 2025 by bondquant
1
How do I build docs?
#4607 opened Mar 8, 2025 by canyon289
1
[docs] Functional API example code does not use Functional API functions
#4602 opened Mar 5, 2025 by samuela
1
Flax nnx ConvTranspose Does Not Restore Input Shape When Used with Conv (Unexpected Behavior)
#4593 opened Mar 4, 2025 by Stella-S-Yan
2
Some surprises when adding Gemma2-2b to flax/examples/gemma
#4567 opened Feb 24, 2025 by mdda
1
Improve support for optax.GradientTransformationExtraArgs in NNX
#4545 opened Feb 13, 2025 by NKlug
3
jax.eval_shape modifies nnx.Module state when using nnx.vmap and nnx.scan
#4520 opened Feb 1, 2025 by ytsmiling
2
[linen] Undocumented RNN errors
#4513 opened Jan 29, 2025 by MRiabov
1
Why does LSTMCell keep rngs in its state?
#4509 opened Jan 28, 2025 by JoaoAparicio
3
How to unroll a model end to end?
#4507 opened Jan 27, 2025 by JoaoAparicio
3
nnx.make_causal_mask() usage
#4505 opened Jan 25, 2025 by windmaple
3
How to use (jax) pytrees inside of nnx modules?
#4497 opened Jan 22, 2025 by PhilipVinc
5
Pool functions reduce over batch dimension and not last dimension
#4494 opened Jan 21, 2025 by simonschoelly
1
ProTip! Type g i on any issue or pull request to go back to the issue listing page.