google/flax

Pull Requests Commits

cgarciae committed 353 days ago

17fc6810

Merge pull request #5016 from google:gym-cleanup

a-googler committed 359 days ago

f6ed5ec1

Merge pull request #5006 from google:variable-refactor

a-googler committed 359 days ago

1d3a2eed

Variable refactor

cgarciae committed 359 days ago

33fbbe6b

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

a-googler committed 359 days ago

59cfeed5

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

a-googler committed 359 days ago

b38601d1

add Rngs KeylessInitializers

cgarciae committed 359 days ago

d96c9edf

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

Jake VanderPlas committed 359 days ago

1453d295

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

a-googler committed 359 days ago

bb519438

Merge pull request #5009 from jburnim:jburnim_sphinx_upgrade

a-googler committed 359 days ago

e8e84353

Remove invalid gymnasium dependency in pyproject.toml

IvyZX committed 359 days ago

Verified 492f063d

optimize scan transpose

cgarciae committed 359 days ago

4727490f

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

jburnim committed 360 days ago

00097aa2

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

a-googler committed 361 days ago

75fd8fa1

Merge pull request #5001 from jburnim:jburnim_python3.13

a-googler committed 361 days ago

f2942d72

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

jburnim committed 362 days ago

0a0645c0

Updated recurrent layers

vfdev-5 committed 363 days ago

e1caa822

Removed *_init atttributes from layers

vfdev-5 committed 364 days ago

d81ce1db

Merge pull request #4961 from samanklesaria:issues/4943

a-googler committed 1 year ago

771eadb1

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

a-googler committed 1 year ago

6080725f

disallow setting metadata through settattr

cgarciae committed 1 year ago

e30c9347

improve set/get_metadata

cgarciae committed 1 year ago

c23304db

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

a-googler committed 1 year ago

8025b3dd

Update flax/nnx/nn/normalization.py

vfdev-5 committed 1 year ago

Verified 4721a6e5

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

a-googler committed 1 year ago

21d1289e

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

a-googler committed 1 year ago

8b8b9298

Fixed broken nnx.statelib.diff

vfdev-5 committed 1 year ago

11a42640

use jnp.stack instead of np.stack

vfdev-5 committed 1 year ago

Verified c08e60b3

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

Jake VanderPlas committed 1 year ago

54fa32d9

Improved nnx SpectralNorm implementation

vfdev-5 committed 1 year ago

77a8c4ca

Older