google/flax

Pull Requests Commits

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

IvyZX committed 19 hours ago

a1cbcab4

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

IvyZX committed 23 hours ago

b3243f77

Release Flax 0.12.12

Cristian Garcia committed 8 days ago

e937d47c

rollback of breaking change

levskaya committed 8 days ago

62ad9d1f

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

a-googler committed 8 days ago

6866a65a

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

a-googler committed 8 days ago

3c3ce3b4

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

a-googler committed 9 days ago

bf06ba37

Add test for invalid keys with sep

mayuriphad committed 9 days ago

8a285ab7

Update flatten_dict docstring

mayuriphad committed 9 days ago

12e226b4

Fix flatten_dict to raise ValueError on invalid keys with sep

mayuriphad committed 9 days ago

d2bdce7a

Fix NNX MultiHeadAttention is_causal mask during decode.

Anu27n committed 9 days ago

c3721a25

Release Flax 0.12.11

Cristian Garcia committed 9 days ago

899109a6

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

Cristian Garcia committed 9 days ago

7a24e6a7

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

jdymitarai committed 11 days ago

ce5acb23

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

a-googler committed 14 days ago

479d43d1

Merge pull request #5488 from samanklesaria:multimodal_llm_example

a-googler committed 14 days ago

2e2876d0

Store nnx.WeightNorm scales as Param so they are trainable

IvyZX committed 14 days ago

ae6d2836

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

cgarciae committed 14 days ago

3ee9be74

Merge pull request #5472 from samanklesaria:audio_example

a-googler committed 14 days ago

7a610fae

Print shapes instead of exact numerics in Embed docstrings

IvyZX committed 14 days ago

c15ac73c

Fix mypy errors with jax 0.12.0

IvyZX committed 14 days ago

82ca84ea

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

IvyZX committed 14 days ago

c3d172b7

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

Cristian Garcia committed 23 days ago

41deea36

Don't pass linen_meta_type argument when creating AxisMetadata subclasses.

a-googler committed 27 days ago

6c9494b5

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

Cristian Garcia committed 27 days ago

5e60ced0

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

oprypin committed 29 days ago

81b168d6

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

Cristian Garcia committed 34 days ago

6cde97df

Merge pull request #5518 from chenkuanliao:fix-issue-5512

a-googler committed 39 days ago

4c441241

Merge pull request #5580 from samanklesaria:fix_mypy_hijax_error

a-googler committed 39 days ago

78abd68f

Fix mypy hijax error

samanklesaria committed 40 days ago

ce09c08c

Older