Skip to content

Issues: google/flax

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

210 Open 768 Closed
210 Open 768 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

Problem training model with jax.jvp in forward pass
#4636 opened Mar 18, 2025 by Thomas-Christie
Add a nnx.identity "activation"?
#4635 opened Mar 18, 2025 by jlperla
4
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
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
nnx.vmap with nnx.split_rngs does not create multiple Dropout layers
#4589 opened Mar 1, 2025 by johnnycrab
2
Several tests failing with AttributeError: module 'jax.api_util' has no attribute 'debug_info'
#4585 opened Feb 28, 2025 by GaetanLepage
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
Restoring a checkpoint with Orbax CheckpointManager fails
#4521 opened Feb 1, 2025 by bednarikjan
1
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
Update dependencies for flaxlib
#4491 opened Jan 19, 2025 by Bot-wxt1221
Failed build: CI - with JAX nightly build failed
#4488 opened Jan 18, 2025 by github-actions bot
6
Improve error message when user mistakenly holds a jax Array in an nnx.Module
#4480 opened Jan 13, 2025 by RaghuSpaceRajan
1
2
Questions on Using nnx.value_and_grad for Loss Calculation and Model Decoupling in Flax NNX
#4476 opened Jan 10, 2025 by Tomato-toast
4
ProTip! Mix and match filters to narrow down what you’re looking for.