google/flax

Pull Requests Commits

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

a-googler committed 142 days ago

c4f2c5c5

Merge pull request #5007 from pfackeldey:differentiable

a-googler committed 142 days ago

97072989

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

a-googler committed 142 days ago

e3304e60

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

a-googler committed 142 days ago

bc8d18a7

Added a note in nnx.jit about arg donation

vfdev-5 committed 142 days ago

efc7e301

feat: add nnx.set_metadata

pfackeldey committed 142 days ago

Verified a916d6c0

Simplify _normalize_values and merge test cases

mohsinm-dev committed 143 days ago

Verified 48464064

Fix nnx.tabulate with Variable hooks

mohsinm-dev committed 144 days ago

Verified 35c31184

Fix formatting in PR template checklist

rapsealk committed 144 days ago

d230605f

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

rapsealk committed 144 days ago

0e3a856e

Fix indentation

samanklesaria committed 144 days ago

32a64d2e

Fixed attribute visualization in __treescope_repr__

vfdev-5 committed 144 days ago

d0e0df9d

use jax.shard_map

Cristian Garcia committed 146 days ago

fd118b12

Use jax.shard_map in flax

yashk2810 committed 146 days ago

5bf9b35f

Merge pull request #5016 from google:gym-cleanup

a-googler committed 147 days ago

f6ed5ec1

Merge pull request #5006 from google:variable-refactor

a-googler committed 147 days ago

1d3a2eed

Variable refactor

cgarciae committed 147 days ago

33fbbe6b

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

a-googler committed 147 days ago

59cfeed5

Merge pull request #5017 from google:update-smpd-guide

a-googler committed 147 days ago

b38601d1

add Rngs KeylessInitializers

cgarciae committed 147 days ago

d96c9edf

Migrate from `jax.experimental.enable_x64` to `jax.enable_x64`.

Jake VanderPlas committed 148 days ago

1453d295

Merge pull request #4998 from google:remove-kernel-init-from-layers

a-googler committed 148 days ago

bb519438

Merge pull request #5009 from jburnim:jburnim_sphinx_upgrade

a-googler committed 148 days ago

e8e84353

Remove invalid gymnasium dependency in pyproject.toml

IvyZX committed 148 days ago

Verified 492f063d

optimize scan transpose

cgarciae committed 148 days ago

4727490f

Use sphinx 6.2+ for docs, which works with Python 3.13.

jburnim committed 149 days ago

00097aa2

Fix nnx.tabulate metadata serialization; add deterministic repr fallback and test

mohsinm-dev committed 149 days ago

Verified b103d0be

Fix nnx.tabulate crash with Variable value hooks

mohsinm-dev committed 149 days ago

Verified 5ccc845a

Merge pull request #4993 from google:no-setattr-metadata-p2

a-googler committed 150 days ago

75fd8fa1

Merge pull request #5001 from jburnim:jburnim_python3.13

a-googler committed 150 days ago

f2942d72

Newer Older