Commits

History for

Commits on Sep 16, 2026

suppress pyrefly errors related to JAX scalar types

Jake VanderPlas
authored and
Flax Authors
committed
4f42ed0

Commits on Aug 6, 2026

add pyrefly suppressions for jax.jit annotation change

Jake VanderPlas
authored and
Flax Authors
committed
c0cb707
[typing] add assertion to avoid pyrefly error

Jake VanderPlas
authored and
Flax Authors
committed
6d224f1

Commits on Jul 2, 2026

Automated Code Change

h-joo
authored and
Flax Authors
committed
18f8f31

Commits on Jun 1, 2026

Fixed gemma example sampler test test_models_from_kaggle
vfdev-5
authored
ec75218

Commits on Apr 23, 2026

Merge pull request #5427 from google:update-gemma-docs
Flax Authors
committed
7075c6e
Rename vars_as to with_vars
samanklesaria
committed
d07091f
Updated gemma docs
vfdev-5
committed
b788d79

Commits on Apr 20, 2026

Addressed comments
vfdev-5
authored
58658ac

Commits on Apr 15, 2026

Fixed intermedites capture
vfdev-5
committed
cdfadd6

Commits on Apr 13, 2026

Merge pull request #4913 from google:gemma-example-using-grain
Flax Authors
committed
9bed435

Commits on Apr 9, 2026

Avoid deprecated pattern of constructing a dtype from an array.

Jake VanderPlas
authored and
Flax Authors
committed
9534e26

Commits on Mar 26, 2026

More Code updates
vfdev-5
committed
5260a1e

Commits on Mar 25, 2026

Added mixed precision: computation and weights dtypes
vfdev-5
committed
cdf3b0c
Use grain in gemma example
vfdev-5
committed
199e653

Commits on Mar 12, 2026

Remove jax/tools/colab_tpu.py.

hawkinsp
authored and
Flax Authors
committed
8ed67b8

Commits on Mar 2, 2026

No public description

danielsuo
authored and
Flax Authors
committed
5486a63

Commits on Feb 12, 2026

Rename sharding_names to out_sharding in NNX Variable metadata

Cristian Garcia
authored and
Flax Authors
committed
a5c7683

Commits on Feb 9, 2026

chore: Migrate gsutil usage to gcloud storage
gurusai-voleti
committed
5d7e13a

Commits on Jan 13, 2026

Update mnist example to use NNX.

danielsuo
authored and
samanklesaria
committed
0ef9abb

Commits on Jan 9, 2026

Copybara import of the project:
Flax Team
committed
a7929b3

Commits on Jan 8, 2026

Merge pull request #5089 from samanklesaria:sharding_metadata
Flax Authors
committed
2863e25

Commits on Jan 6, 2026

improve hijax guide
cgarciae
committed
5e31fe9
[pmap] Avoid degraded performance under the new `jax.pmap`.

danielsuo
authored and
Flax Authors
committed
5d90851

Commits on Jan 5, 2026

Rename sharding_names to sharding_metadata
samanklesaria
committed
3e96343

Commits on Dec 18, 2025

Release 0.12.2
IvyZX
committed
6d36623

Commits on Dec 5, 2025

[flax:examples:nlp_seq] Create a main.py file to run tests with config files to match other examples. #jax-fixit

danielsuo
authored and
Flax Authors
committed
872d50b
[flax:examples:lm1b_nnx] Update example to work internally. #jax-fixit.

danielsuo
authored and
Flax Authors
committed
881b8ed

Commits on Dec 3, 2025

[flax:examples:ppo] Fix some linter / import issues. #jax-fixit

danielsuo
authored and
Flax Authors
committed
59cfd99

Commits on Dec 2, 2025

[flax:examples:sst2] Fix notebook error.

danielsuo
authored and
Flax Authors
committed
5a3d148

Commits on Dec 1, 2025

[flax:examples:sst2] Fix pytype errors.

danielsuo
authored and
Flax Authors
committed
6ac5a78
[flax:examples:gemma] Fixing linter errors.

danielsuo
authored and
Flax Authors
committed
324f4b4
[flax:examples:vae] Small linter fixes.

danielsuo
authored and
Flax Authors
committed
88d7497
[flax:examples:seq2seq] Create main and default config based on seq2seq.ipynb.

danielsuo
authored and
Flax Authors
committed
9ed59e2
[flax:examples:wmt] Small linter fixes.

danielsuo
authored and
Flax Authors
committed
c29b6d6