google/flax

Pull Requests Commits

Migrate away from private jax._src.tree_util._registry

IvyZX committed 82 days ago

a147752a

Merge pull request #5454 from samanklesaria:dataloader_docs

a-googler committed 82 days ago

8485ec3c

Merge pull request #5363 from samanklesaria:rng_guide_update

a-googler committed 86 days ago

d044de84

Add data loaders tutorial

samanklesaria committed 86 days ago

64145c4e

Merge pull request #5441 from samanklesaria:fix-transforms-doctest

a-googler committed 89 days ago

e9aaebe5

Merge pull request #5440 from samanklesaria:fix-graphlib-doctest

a-googler committed 89 days ago

d07c450d

Merge pull request #5442 from samanklesaria:fix-rnglib-doctest

a-googler committed 89 days ago

2ee10c2b

Merge pull request #5425 from samanklesaria:tensorboard_example

a-googler committed 89 days ago

a9057b88

Cristian Garcia committed 90 days ago

e9717a14

Merge pull request #5434 from google:mha-out_sharding

a-googler committed 90 days ago

2fa7e66e

Merge branch 'main' into mha-out_sharding

vfdev-5 committed 98 days ago

Verified 502012b8

Merge pull request #5446 from samanklesaria:msgpack_restore_fix

a-googler committed 101 days ago

2fe9c869

Fix doctests in transforms for graph_updates=False default

samanklesaria committed 101 days ago

bd6bf36e

Update docstring for split_rngs

samanklesaria committed 102 days ago

9a360412

Merge pull request #5405 from samanklesaria:minigpt

a-googler committed 102 days ago

56b4f38f

Merge pull request #5443 from samanklesaria:fix-docstring-graphilib

a-googler committed 102 days ago

35c5ac5d

Merge pull request #5101 from google:rewrite-prefetch_to_device

a-googler committed 103 days ago

c781023d

improve Optimizer

Cristian Garcia committed 103 days ago

63c19603

fix Pytree typing issue

Cristian Garcia committed 103 days ago

36d28aec

Add fix for msgpack serialization of nnx.List structures

samanklesaria committed 104 days ago

b84b69ce

Merge pull request #5098 from google:reenable-pytest-examples

a-googler committed 107 days ago

af461c1f

Fix docstrings in graphlib under graph=False

samanklesaria committed 107 days ago

f17dc029

Fix update doctest for graph_updates=False default

samanklesaria committed 108 days ago

4f9f0acb

Added tests for flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate

vfdev-5 committed 108 days ago

b0559492

Merge pull request #5427 from google:update-gemma-docs

a-googler committed 108 days ago

7075c6e6

Merge pull request #5431 from samanklesaria:machine_translation

a-googler committed 108 days ago

48141ec0

Reenabled examples tests

vfdev-5 committed 108 days ago

0fb638c2

Merge pull request #5384 from samanklesaria:vars_as_rename

a-googler committed 108 days ago

e3a73c0f

Comment out tensorboard magic for sphinx rendering

samanklesaria committed 108 days ago

83e3c21b

Add minigpt example

samanklesaria committed 108 days ago

2b052ed7

Older