google/flax

Pull Requests Commits

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

vfdev-5 committed 1 year ago

2ea9f3b0

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

a-googler committed 1 year ago

bb3c31ac

Merge pull request #5034 from google:fix-reseed

a-googler committed 1 year ago

6770d71a

fix reseed for abstract values

cgarciae committed 1 year ago

0e490a7c

Support for python 3.14

vfdev-5 committed 1 year ago

6622efcf

Merge pull request #5033 from IvyZX:shard-msg

a-googler committed 1 year ago

c0ea12d3

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

a-googler committed 1 year ago

303f8b29

Add flip doc link to eager sharding error message

IvyZX committed 1 year ago

da8ef361

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

a-googler committed 1 year ago

760be244

Fix tf installation on 3.13 and protobuf version

vfdev-5 committed 1 year ago

ede2bb1b

test against python 3.13

cgarciae committed 1 year ago

44b97a8d

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

a-googler committed 1 year ago

befe5718

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

a-googler committed 1 year ago

c4f2c5c5

Merge pull request #5007 from pfackeldey:differentiable

a-googler committed 1 year ago

97072989

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

a-googler committed 1 year ago

e3304e60

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

a-googler committed 1 year ago

bc8d18a7

Added a note in nnx.jit about arg donation

vfdev-5 committed 1 year ago

efc7e301

feat: add nnx.set_metadata

pfackeldey committed 1 year ago

Verified a916d6c0

Simplify _normalize_values and merge test cases

mohsinm-dev committed 1 year ago

Verified 48464064

Fix nnx.tabulate with Variable hooks

mohsinm-dev committed 1 year ago

Verified 35c31184

Fix formatting in PR template checklist

rapsealk committed 1 year ago

d230605f

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

rapsealk committed 1 year ago

0e3a856e

Fix indentation

samanklesaria committed 1 year ago

32a64d2e

Fixed attribute visualization in __treescope_repr__

vfdev-5 committed 1 year ago

d0e0df9d

use jax.shard_map

Cristian Garcia committed 1 year ago

fd118b12

Use jax.shard_map in flax

yashk2810 committed 1 year ago

5bf9b35f

Merge pull request #5016 from google:gym-cleanup

a-googler committed 1 year ago

f6ed5ec1

Merge pull request #5006 from google:variable-refactor

a-googler committed 1 year ago

1d3a2eed

Variable refactor

cgarciae committed 1 year ago

33fbbe6b

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

a-googler committed 1 year ago

59cfeed5

Newer Older