google/flax

Pull Requests Commits

cgarciae committed 60 days ago

17fc6810

Merge pull request #5016 from google:gym-cleanup

a-googler committed 66 days ago

f6ed5ec1

Merge pull request #5006 from google:variable-refactor

a-googler committed 66 days ago

1d3a2eed

Variable refactor

cgarciae committed 66 days ago

33fbbe6b

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

a-googler committed 66 days ago

59cfeed5

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

a-googler committed 66 days ago

b38601d1

add Rngs KeylessInitializers

cgarciae committed 66 days ago

d96c9edf

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

Jake VanderPlas committed 66 days ago

1453d295

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

a-googler committed 66 days ago

bb519438

Merge pull request #5009 from jburnim:jburnim_sphinx_upgrade

a-googler committed 66 days ago

e8e84353

Remove invalid gymnasium dependency in pyproject.toml

IvyZX committed 66 days ago

Verified 492f063d

optimize scan transpose

cgarciae committed 66 days ago

4727490f

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

jburnim committed 67 days ago

00097aa2

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

a-googler committed 68 days ago

75fd8fa1

Merge pull request #5001 from jburnim:jburnim_python3.13

a-googler committed 68 days ago

f2942d72

Avoid install, import, or tests with tensorflow-text under Python 3.13+.

jburnim committed 69 days ago

0a0645c0

Updated recurrent layers

vfdev-5 committed 69 days ago

e1caa822

Removed *_init atttributes from layers

vfdev-5 committed 71 days ago

d81ce1db

Merge pull request #4961 from samanklesaria:issues/4943

a-googler committed 72 days ago

771eadb1

Merge pull request #4985 from google:no-setattr-metadata

a-googler committed 74 days ago

6080725f

disallow setting metadata through settattr

cgarciae committed 74 days ago

e30c9347

improve set/get_metadata

cgarciae committed 74 days ago

c23304db

Merge pull request #4623 from mattbahr:implement-spectral-and-instance-norm

a-googler committed 74 days ago

8025b3dd

Update flax/nnx/nn/normalization.py

vfdev-5 committed 74 days ago

Verified 4721a6e5

Merge pull request #4992 from google:fix-nnx-diff

a-googler committed 74 days ago

21d1289e

Merge pull request #4991 from google:use-jnp-stack-in-stack_forest

a-googler committed 75 days ago

8b8b9298

Fixed broken nnx.statelib.diff

vfdev-5 committed 75 days ago

11a42640

use jnp.stack instead of np.stack

vfdev-5 committed 75 days ago

Verified c08e60b3

Future-proof imports of jax.new_ref / jax.Ref.

Jake VanderPlas committed 76 days ago

54fa32d9

Improved nnx SpectralNorm implementation

vfdev-5 committed 76 days ago

77a8c4ca

Older