google/flax

Pull Requests Commits

Cristian Garcia committed 5 days ago

a52d55b4

Merge pull request #5446 from samanklesaria:msgpack_restore_fix

a-googler committed 6 days ago

2fe9c869

Merge pull request #5405 from samanklesaria:minigpt

a-googler committed 7 days ago

56b4f38f

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

a-googler committed 7 days ago

35c5ac5d

Merge pull request #5101 from google:rewrite-prefetch_to_device

a-googler committed 8 days ago

c781023d

improve Optimizer

Cristian Garcia committed 8 days ago

63c19603

fix Pytree typing issue

Cristian Garcia committed 8 days ago

36d28aec

Add fix for msgpack serialization of nnx.List structures

samanklesaria committed 9 days ago

b84b69ce

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

a-googler committed 12 days ago

af461c1f

Fix docstrings in graphlib under graph=False

samanklesaria committed 12 days ago

f17dc029

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

vfdev-5 committed 13 days ago

b0559492

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

a-googler committed 13 days ago

7075c6e6

Merge pull request #5431 from samanklesaria:machine_translation

a-googler committed 13 days ago

48141ec0

Reenabled examples tests

vfdev-5 committed 13 days ago

0fb638c2

Merge pull request #5384 from samanklesaria:vars_as_rename

a-googler committed 13 days ago

e3a73c0f

Add minigpt example

samanklesaria committed 13 days ago

2b052ed7

Merge pull request #5408 from google:migration-pytorch

a-googler committed 13 days ago

9f17346a

Rename vars_as to with_vars

samanklesaria committed 13 days ago

d07091fb

Merge pull request #5424 from samanklesaria:example_cleanup

a-googler committed 13 days ago

0ccef558

Merge pull request #5416 from samanklesaria:view_guide_change

a-googler committed 13 days ago

b1f0e220

Merge pull request #5409 from google:gemma-example-using-grain--fix-intermediates

a-googler committed 13 days ago

8671e60a

Updated gemma docs

vfdev-5 committed 13 days ago

b788d794

Disallow `P.UNCONSTRAINED` in `out_shardings` in `jit`. This leads to simplifications in how we lower shardings to stableHLO.

yashk2810 committed 15 days ago

9877bea7

Merge pull request #5422 from SAY-5:fix-doubled-the-typos

a-googler committed 15 days ago

09ed488a

Deprecate Variable hooks in favor of custom Variable types.

a-googler committed 15 days ago

a13993f6

Add deeper explanations to translation example

samanklesaria committed 16 days ago

44d96f73

Add explanation of nnx.view functions

samanklesaria committed 16 days ago

aab0c99f

Update view guide to mention with_ and as_ prefixes

samanklesaria committed 16 days ago

4931b8fe

Reorganize examples

samanklesaria committed 16 days ago

b7d93881

Add machine translation example

samanklesaria committed 16 days ago

2314a04a

Older