google/flax

Pull Requests Commits

examples/gemma: Fix sampler shardings under Explicit mesh

james04nesbitt committed 13 hours ago

96b4d236

Add warmup_steps support to nnx.EMA.

a-googler committed 3 days ago

204a819c

Fix pre-existing Pyrefly type-check errors in Flax and remove migration suppressions

IvyZX committed 7 days ago

54aada1e

Align Python 3.12–3.15 classifiers, GitHub Actions CI, and Pyrefly type checking with JAX

IvyZX committed 9 days ago

71c5b9c5

Use `jax.tree_util.is_tree_node` instead of private `_registry` in NNX

IvyZX committed 10 days ago

b3243f77

Release Flax 0.12.12

Cristian Garcia committed 18 days ago

e937d47c

rollback of breaking change

levskaya committed 18 days ago

62ad9d1f

Merge pull request #5593 from mayuriphad:fix-flatten-dict-sep

a-googler committed 18 days ago

6866a65a

Merge pull request #5592 from Anu27n:fix/nnx-mha-decode-is-causal

a-googler committed 18 days ago

3c3ce3b4

Merge pull request #5582 from jdymitarai:fix/nnx-capture-scan-list

a-googler committed 18 days ago

bf06ba37

Add test for invalid keys with sep

mayuriphad committed 18 days ago

8a285ab7

Update flatten_dict docstring

mayuriphad committed 18 days ago

12e226b4

Fix flatten_dict to raise ValueError on invalid keys with sep

mayuriphad committed 18 days ago

d2bdce7a

Fix NNX MultiHeadAttention is_causal mask during decode.

Anu27n committed 18 days ago

c3721a25

Release Flax 0.12.11

Cristian Garcia committed 18 days ago

899109a6

[nnx] Make `nnx.Sequential` iterable and filter `**kwargs` per layer.

Cristian Garcia committed 18 days ago

7a24e6a7

fix(nnx): skip Sequence modules in capture to prevent mixed key types (fixes #5571)

jdymitarai committed 21 days ago

ce5acb23

Merge pull request #5623 from IvyZX:fix-5577-weightnorm

a-googler committed 24 days ago

479d43d1

Merge pull request #5488 from samanklesaria:multimodal_llm_example

a-googler committed 24 days ago

2e2876d0

Store nnx.WeightNorm scales as Param so they are trainable

IvyZX committed 24 days ago

ae6d2836

Merge remote-tracking branch 'upstream/main' into multimodal_llm_example

cgarciae committed 24 days ago

3ee9be74

Merge pull request #5472 from samanklesaria:audio_example

a-googler committed 24 days ago

7a610fae

Print shapes instead of exact numerics in Embed docstrings

IvyZX committed 24 days ago

c15ac73c

Fix mypy errors with jax 0.12.0

IvyZX committed 24 days ago

82ca84ea

Merge remote-tracking branch 'origin/main' into audio_example

IvyZX committed 24 days ago

c3d172b7

Add `capture_noop` option to `Module.perturb`.

Cristian Garcia committed 32 days ago

41deea36

Don't pass linen_meta_type argument when creating AxisMetadata subclasses.

a-googler committed 36 days ago

6c9494b5

Fix Flax GitHub Actions CI (`pytest` `DeprecationWarning` and `pytype` in `examples/gemma/train.py`).

Cristian Garcia committed 37 days ago

5e60ced0

Replace `# pytype: disable` suppressions with `# pyrefly: ignore`

oprypin committed 38 days ago

81b168d6

Add `nnx_sow_capture_noop` config flag and `capture_noop` option to `Module.sow`.

Cristian Garcia committed 43 days ago

6cde97df

Older