Pull Requests google/flax

add support for hijax Variables in nnx transforms

#5065 opened 2025-11-01 19:19 by cgarciae

Update mnist example to use NNX (clone of #5002) pull ready

#5064 opened 2025-10-31 17:53 by samanklesaria

Add split and key methods to RngStream and Rngs pull ready

#5047 opened 2025-10-23 14:40 by samanklesaria

Box support

#5029 opened 2025-10-14 23:29 by cgarciae

New set mode

#5028 opened 2025-10-14 21:12 by chapman20j

Make Linen and NNX logical rule deduction align

#5003 opened 2025-10-07 01:20 by IvyZX

chore: more specific migration guide description

#4977 opened 2025-09-24 09:44 by emergenz

Add a customizable conv_transpose. pull ready

#4938 opened 2025-09-11 21:52 by copybara-service[bot]

Don't pass linen_meta_type argument when creating AxisMetadata subclasses.

#4933 opened 2025-09-09 11:42 by copybara-service[bot]

Improve error message with class name.

#4932 opened 2025-09-09 08:38 by copybara-service[bot]

adds recursive_map

#4926 opened 2025-09-05 01:08 by copybara-service[bot]

Add intermediate value captures based off hijax.Box

#4925 opened 2025-09-04 02:14 by IvyZX

Updates in gemma example

#4913 opened 2025-08-27 10:45 by vfdev-5

Refactor ArrayCheckpointHandler to use BasePyTreeCheckpointHandler.

#4902 opened 2025-08-22 06:00 by copybara-service[bot]

Add python 3.13 CI tests

#4888 opened 2025-08-13 21:48 by IvyZX

fix_checkpoint helper

#4875 opened 2025-08-06 00:02 by rcrowe-google

Fixing formatting

#4874 opened 2025-08-05 22:13 by rcrowe-google

Adding link to release notes

#4873 opened 2025-08-05 22:07 by rcrowe-google

Now Github release notes

#4872 opened 2025-08-05 22:01 by rcrowe-google

Migrate `rnglib.py` from `.value` to `[...]`

#4840 opened 2025-07-23 02:49 by lukeyeh

Migrate `metrics.py` from `.value` to `[...]`

#4839 opened 2025-07-23 02:49 by lukeyeh

Migrate `module_test.py` from `.value` to `[...]`

#4838 opened 2025-07-23 02:49 by lukeyeh

Migrate `containers_test.py` from `.value` to `[...]`

#4837 opened 2025-07-23 02:49 by lukeyeh

Migrate `conv_test.py` from `.value` to `[...]`

#4836 opened 2025-07-23 02:49 by lukeyeh

Migrate `embed_test.py` from `.value` to `[...]`

#4835 opened 2025-07-23 02:48 by lukeyeh

Migrate `normalization_test.py` from `.value` to `[...]`

#4834 opened 2025-07-23 02:48 by lukeyeh

Migrate `linear_test.py` from `.value` to `[...]`

#4833 opened 2025-07-23 02:48 by lukeyeh

Migrate `sampler_test.py` from `.value` to `[...]`

#4832 opened 2025-07-23 02:48 by lukeyeh

[nnx] update bridge

#4828 opened 2025-07-21 19:32 by cgarciae

Use .format in place of .layout.

#4785 opened 2025-06-19 00:11 by copybara-service[bot]