google/flax

Pull Requests Commits

[docs] Exposed more functions/classes in state.rst

vfdev-5 committed 75 days ago

2ea9f3b0

Deduplicate `Variable` nodes in `iter_graph` and eliminate recursion.

a-googler committed 75 days ago

bb3c31ac

Merge pull request #5034 from google:fix-reseed

a-googler committed 78 days ago

6770d71a

fix reseed for abstract values

cgarciae committed 78 days ago

0e490a7c

Support for python 3.14

vfdev-5 committed 78 days ago

6622efcf

Merge pull request #5033 from IvyZX:shard-msg

a-googler committed 79 days ago

c0ea12d3

Merge pull request #5031 from google:docs-nnx-jit-arg-donation

a-googler committed 79 days ago

303f8b29

Add flip doc link to eager sharding error message

IvyZX committed 79 days ago

da8ef361

Merge pull request #4987 from google:python-3.13-support

a-googler committed 79 days ago

760be244

Fix tf installation on 3.13 and protobuf version

vfdev-5 committed 79 days ago

ede2bb1b

test against python 3.13

cgarciae committed 79 days ago

44b97a8d

Merge pull request #5008 from mohsinm-dev:fix-nnx-tabulate-variable-hooks

a-googler committed 80 days ago

befe5718

Merge pull request #5023 from rapsealk:docs/readme-nnx-linear

a-googler committed 80 days ago

c4f2c5c5

Merge pull request #5007 from pfackeldey:differentiable

a-googler committed 80 days ago

97072989

Merge pull request #5022 from google:fix-nnx-display-treescope-repr-5004

a-googler committed 80 days ago

e3304e60

Merge pull request #5024 from rapsealk:ci/pr-template-line-feeds

a-googler committed 80 days ago

bc8d18a7

Added a note in nnx.jit about arg donation

vfdev-5 committed 80 days ago

efc7e301

feat: add nnx.set_metadata

pfackeldey committed 80 days ago

Verified a916d6c0

Simplify _normalize_values and merge test cases

mohsinm-dev committed 81 days ago

Verified 48464064

Fix nnx.tabulate with Variable hooks

mohsinm-dev committed 82 days ago

Verified 35c31184

Fix formatting in PR template checklist

rapsealk committed 82 days ago

d230605f

Update README to use fully qualified `nnx.Linear` in example

rapsealk committed 82 days ago

0e3a856e

Fix indentation

samanklesaria committed 82 days ago

32a64d2e

Fixed attribute visualization in __treescope_repr__

vfdev-5 committed 82 days ago

d0e0df9d

use jax.shard_map

Cristian Garcia committed 84 days ago

fd118b12

Use jax.shard_map in flax

yashk2810 committed 84 days ago

5bf9b35f

Merge pull request #5016 from google:gym-cleanup

a-googler committed 85 days ago

f6ed5ec1

Merge pull request #5006 from google:variable-refactor

a-googler committed 85 days ago

1d3a2eed

Variable refactor

cgarciae committed 86 days ago

33fbbe6b

Merge pull request #5015 from google:optimize-scan-transpose

a-googler committed 86 days ago

59cfeed5

Newer Older