google/flax

Pull Requests Commits

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

jburnim committed 82 days ago

0a0645c0

Updated recurrent layers

vfdev-5 committed 83 days ago

e1caa822

Convert linen pytorch guide to nnx

samanklesaria committed 83 days ago

0c06bdc2

Removed *_init atttributes from layers

vfdev-5 committed 84 days ago

d81ce1db

Merge pull request #4961 from samanklesaria:issues/4943

a-googler committed 85 days ago

771eadb1

Update index.rst example to use explicit randomness

samanklesaria committed 86 days ago

e3004690

Add pros and cons of explicit randomness

samanklesaria committed 86 days ago

819ac067

Use fork_rngs with implicit state in randomness guide

samanklesaria committed 86 days ago

1e0b9ea5

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

a-googler committed 87 days ago

6080725f

disallow setting metadata through settattr

cgarciae committed 87 days ago

e30c9347

improve set/get_metadata

cgarciae committed 87 days ago

c23304db

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

a-googler committed 87 days ago

8025b3dd

Update flax/nnx/nn/normalization.py

vfdev-5 committed 87 days ago

Verified 4721a6e5

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

a-googler committed 88 days ago

21d1289e

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

a-googler committed 88 days ago

8b8b9298

Fixed broken nnx.statelib.diff

vfdev-5 committed 88 days ago

11a42640

use jnp.stack instead of np.stack

vfdev-5 committed 88 days ago

Verified c08e60b3

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

Jake VanderPlas committed 89 days ago

54fa32d9

Improved nnx SpectralNorm implementation

vfdev-5 committed 89 days ago

77a8c4ca

Merge pull request #4891 from mohsinm-dev:fix-nnx-tabulate-empty-dict-issue-4889

a-googler committed 89 days ago

f65afdb8

Update tutorial examples to thread explicit RNGs

samanklesaria committed 89 days ago

cdf9ff88

Merge pull request #4948 from samanklesaria:issues/4883

a-googler committed 89 days ago

ebb980a6

clean up jax.Ref-related names

Cristian Garcia committed 89 days ago

4c4bc3b1

Merge pull request #4980 from IvyZX:shardref

a-googler committed 92 days ago

a9dccc06

Clarify `*Norm` layer docstrings: `axis_index_groups` is unused under SPMD jit.

arnoegw committed 93 days ago

ae1594b0

Remove type annotation

samanklesaria committed 93 days ago

8b7c8bd5

Merge pull request #4983 from google:propagate-eager-sharding

a-googler committed 93 days ago

34e1d68e

Merge pull request #4968 from google:get-set-metadata

a-googler committed 93 days ago

b9cc05f4

Merge pull request #4981 from google:jit-wrapped-types

a-googler committed 93 days ago

609a0aba

Cristian Garcia committed 93 days ago

f8c5943c

Newer Older