google/flax

Pull Requests Commits

add Rngs random helpers

cgarciae committed 184 days ago

769cc502

allow update and pop to work on dict

cgarciae committed 185 days ago

b5f67688

Cristian Garcia committed 185 days ago

220910f8

Merge pull request #4773 from sanepunk:main

a-googler committed 185 days ago

dfaaa486

simplify nnx_basics

cgarciae committed 186 days ago

a2329ce1

Refactor code cells in MNIST tutorial for consistency and clarity

sanepunk committed 187 days ago

597ab5b8

Update mnist_tutorial.md

sanepunk committed 188 days ago

Verified a91b3721

Enhance CNN model with BatchNorm and Dropout layers for improved training stability and performance. Update training and evaluation steps to switch model modes appropriately. Clarify documentation on the use of .eval() for inference.

sanepunk committed 188 days ago

25ee4ce5

Copybara import of the project:

IvyZX committed 188 days ago

8dea0de4

Add old property attributes for object->pytree rename.

IvyZX committed 188 days ago

16677735

Merge pull request #4863 from google:nnx-obj-pytree

a-googler committed 189 days ago

38c058b7

Pytree + ArrayRef refactor

cgarciae committed 189 days ago

a5400287

Replace `jax.sharding.use_mesh` with `jax.set_mesh`. `jax.set_mesh` can act as a global setter or a context manager.

yashk2810 committed 189 days ago

3285ba08

Merge pull request #4859 from google:tunix-demo-fixes

a-googler committed 190 days ago

b807186a

Fix string key preservation in replace_by_pure_dict

mohsinm-dev committed 190 days ago

Verified 713745df

miscellaneous improvements

cgarciae committed 190 days ago

5d5c907d

Fix broken philosophy documentation link

mohsinm-dev committed 191 days ago

Verified 5e611eb3

Merge pull request #4795 from IvyZX:doc

a-googler committed 191 days ago

210e4d6d

Merge pull request #4796 from SobhanMP:identity-sequantial

a-googler committed 191 days ago

ac44f694

Add a sharp bits doc

IvyZX committed 191 days ago

4202efac

Update 05_vae.py

nenuadrian committed 191 days ago

Verified 3dfeff9f

Merge pull request #4854 from google:mutable-array-guide

a-googler committed 192 days ago

c9fc00b4

add 0.11.0 migration guide

cgarciae committed 192 days ago

2de53d54

Consistent doctests with JAX repo

jlperla committed 192 days ago

6ef75906

Fix broken contributing documentation link

mohsinm-dev committed 192 days ago

Verified e4094965

jlperla committed 192 days ago

f68ffe50

Merge pull request #4841 from lukeyeh:push-kqruztmzotnw

a-googler committed 193 days ago

c429072b

Fix ToLinen docstring return description

mohsinm-dev committed 194 days ago

cd7c8660

Merge branch 'main' into main

sanepunk committed 195 days ago

Verified 96c565ff

Make `Sequential()` behave like the identity function

SobhanMP committed 195 days ago

10b95ba5

Newer Older