google/flax

View on GitHub

Flax is a neural network library for JAX that is designed for flexibility.

Stars: 7.3k
Forks: 833
Open beginner issues: 1
Indexed issues: 231
Avg merge: 5h 11m
Merged PRs (30d): 5

Dominant language: Jupyter Notebook
License: Apache-2.0
Last GitHub push: Aug 21, 2026
Latest indexed: Sep 16, 2026
Contributing guide: Contributing guide
Code of conduct: Code of conduct
Beginner labels: No beginner labels indexed

Open All

231 open issues indexed

Linen: Consider raising an error when reading variables from submodule before initialization? Open

Priority: P2 - no schedule Status: pull requests welcome

google/flax#513 · 10 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Improve Error Message: Assigning a list of Modules to `self` incorrectly Open

Priority: P1 - soon

google/flax#524 · 15 comments · 3 reactions · 2 assignees · Indexed Sep 16, 2026
Convert Linen design tests into real unit tests Open

Priority: P2 - no schedule Status: pull requests welcome

google/flax#525 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Updating examples/nlp_seq Open

Priority: P1 - soon

google/flax#568 · 6 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Issues restoring checkpoint of struct.dataclass w/ FrozenDict attr Open

Priority: P2 - no schedule Status: pull requests welcome

google/flax#676 · 7 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
bfloat16 numeric stability w/ log_softmax + cross_entropy as based on Flax examples Open

google/flax#677 · 2 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Add "DotGetter" style variable and submodule accessors? Open

Priority: P2 - no schedule

google/flax#749 · 1 comment · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Improve Error Message: Jitting linen.Module Open

Priority: P2 - no schedule

google/flax#853 · 13 comments · 2 reactions · 0 assignees · Indexed Sep 16, 2026
Write design note: Explain how Linen manages RNGs Open

Priority: P2 - no schedule

google/flax#878 · 3 comments · 3 reactions · 2 assignees · Indexed Sep 16, 2026
Consider deprecating common_utils.shard Open

Priority: P2 - no schedule

google/flax#948 · 2 comments · 3 reactions · 1 assignee · Indexed Sep 16, 2026
Deprecate flax.jax_utils.replicate Open

google/flax#1161 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Return key from `traverse_util.Traversal.iterate()` Open

Status: pull requests welcome

google/flax#1170 · 6 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
AutoRegressive Decoding currently fails if input prompt > 1 Open

Priority: P2 - no schedule

google/flax#1317 · 6 comments · 1 reaction · 2 assignees · Indexed Sep 16, 2026
Improve RTD for initializers Open

Priority: P1 - soon

google/flax#1386 · 9 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Multiple Inheritance -> doesn't recognize as Module throws ValueError: parent must be None, Module or Scope Open

Priority: P2 - no schedule

google/flax#1409 · 4 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Test pickle support Open

Priority: P2 - no schedule Status: pull requests welcome

google/flax#1481 · 1 comment · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Implementing a custom lifted transformation: jax.experimental.ode.odeint Open

Priority: P1 - soon

google/flax#1496 · 9 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Improve documentation of make_causal_mask Open

Priority: P2 - no schedule Status: pull requests welcome

google/flax#1520 · 6 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
[Documentation]: clarify mutability of collections rather than variables Open

Priority: P2 - no schedule

google/flax#1616 · 2 comments · 2 reactions · 1 assignee · Indexed Sep 16, 2026
checkpoints.restore_checkpoint` should raise an error for explicit paths that don't exist Open

Priority: P2 - no schedule

google/flax#1631 · 2 comments · 2 reactions · 0 assignees · Indexed Sep 16, 2026
Improve Error Message: Consider extending AttributeError on Flax Modules Open

Priority: P2 - no schedule

google/flax#1685 · 3 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
PJit example Open

Priority: P2 - no schedule

google/flax#1763 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Scan with split parameters Open

google/flax#1771 · 6 comments · 1 reaction · 1 assignee · Indexed Sep 16, 2026
Deprecate/Remove traverse_util module Open

Priority: P2 - no schedule

google/flax#1839 · 7 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Adding a new transposed convolution function (similar to `torch.nn.ConvTranspose2d()`) Open

Priority: P2 - no schedule Status: blocked

google/flax#1872 · 8 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Create Colab Snippets Open

Priority: P2 - no schedule

google/flax#1912 · 4 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Improve lifted transforms docstrings Open

Priority: P2 - no schedule

google/flax#1977 · 3 comments · 0 reactions · 2 assignees · Indexed Sep 16, 2026
Start using positional-only parameters in Linen Open

Priority: P2 - no schedule

google/flax#2049 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
We broke the 1:1 correspondence with attribute names and variable dict names Open

Priority: P1 - soon

google/flax#2100 · 5 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
HOWTOS tracking issue Open

Priority: P2 - no schedule Status: pull requests welcome

google/flax#2116 · 3 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Improve SEO for Flax Open

Priority: P2 - no schedule

google/flax#2124 · 0 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
LazyRNG can accidentally have key sharing across layers. Open

Priority: P1 - soon

google/flax#2157 · 2 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
better support for RNN layers required Open

Priority: P1 - soon

google/flax#2170 · 2 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Upstream generate_autosummary_content patch to Sphinx Open

google/flax#2196 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Debug differences in `examples/nlp_seq`. Open

Priority: P2 - no schedule

google/flax#2201 · 0 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
IPython autoreload is not supported by flax Open

google/flax#2286 · 2 comments · 1 reaction · 1 assignee · Indexed Sep 16, 2026
`__getattr__` and pytype checking Open

google/flax#2416 · 17 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
FLIP: Generalized metadata Open

FLIP

google/flax#2434 · 1 comment · 0 reactions · 1 assignee · Indexed Sep 16, 2026
nn.scan issues Open

google/flax#2618 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Show how to launch Flax examples with Vertex Open

Priority: P2 - no schedule

google/flax#2619 · 0 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Proposal: New RTD content structuring Open

Priority: P2 - no schedule

google/flax#2627 · 4 comments · 4 reactions · 1 assignee · Indexed Sep 16, 2026
Rename `flax.linen` to `flax.nn`? Open

Priority: P2 - no schedule

google/flax#2677 · 6 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Enable internal type checking for tests Open

Priority: P2 - no schedule

google/flax#2679 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Support loading and fine-tuning TensorFlow models Open

Priority: P2 - no schedule

google/flax#2721 · 1 comment · 2 reactions · 1 assignee · Indexed Sep 16, 2026
Use enterprise test runners for our CI once they are working again Open

Priority: P1 - soon

google/flax#2731 · 1 comment · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Wrap `__post_init__` to check or force that `Module.__post_init__` is called Open

Priority: P2 - no schedule

google/flax#2733 · 2 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Properly bind all submodules Open

Priority: P2 - no schedule

google/flax#2744 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Common Initializers Does Not Work with Bias Open

Priority: P1 - soon

google/flax#2749 · 6 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
`Module.init` does not contain KV-pair of modules without parameters Open

google/flax#2777 · 4 comments · 1 reaction · 0 assignees · Indexed Sep 16, 2026
switching to _src structure Open

Priority: P2 - no schedule

google/flax#2814 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026

google/flax

Get new issues in your inbox