google/flax

Pull Requests Commits

fix RNN example

cgarciae committed 2 years ago

6c0f800a

cgarciae committed 2 years ago

734729b6

Merge pull request #2872 from chiamp:frozen_dict_feature_flag

a-googler committed 2 years ago

27c37ed1

added feature flag for FrozenDicts

chiamp committed 2 years ago

48f6bda1

Fix or ignore some pytype errors related to jnp.ndarray == jax.Array.

hawkinsp committed 2 years ago

2fbeb3b5

Merge pull request #2881 from chiamp:rmslayernorm

a-googler committed 2 years ago

5f0ac50c

Added RMS normalization layer

chiamp committed 2 years ago

9cff7809

Merge pull request #2873 from levskaya:relaxed_naming2

a-googler committed 2 years ago

98ec8fe2

Refactor and clarify name relaxation policy changes.

levskaya committed 2 years ago

5d8e93db

Merge pull request #2812 from nouiz:imagenet_benchmark

a-googler committed 2 years ago

018725f3

Merge pull request #2879 from cpgaffney1:dev_branch

a-googler committed 2 years ago

5d4040a4

Add additional Orbax link to use_checkpointing.ipynb

cpgaffney1 committed 2 years ago

Verified 21ed0fbd

Add additional Orbax link to use_checkpointing.md

cpgaffney1 committed 2 years ago

Verified 75db8675

Use Sequence rather than Tuple in ConvTranspose type hints.

a-googler committed 2 years ago

36542787

Merge branch 'google:main' into dev_branch

cpgaffney1 committed 2 years ago

Verified 2637428d

Revised docs using Flax contributing guide

cpgaffney1 committed 2 years ago

a96f49ba

Flax layers are parametrizable with custom conv_general_dilated.

a-googler committed 2 years ago

92b0cdd4

Remove fallbacks for old versions of JAX from Flax.

hawkinsp committed 2 years ago

323d6b85

Flax layers are parametrizable with custom dot_general.

a-googler committed 2 years ago

1f6b0949

Replace uses of deprecated JAX sharding APIs with their new names in jax.sharding.

hawkinsp committed 2 years ago

26d46e4b

Merge pull request #2856 from levskaya:point_release4

a-googler committed 2 years ago

524629eb

Flax 0.6.5 point release

levskaya committed 2 years ago

5c6f4267

Merge pull request #2837 from cgarciae:improve-remat-docs

a-googler committed 2 years ago

6ea0aca8

Merge pull request #2816 from jheek:lazy-init

a-googler committed 2 years ago

24224bff

Merge pull request #2811 from arcAman07:main

a-googler committed 2 years ago

c3c4a876

Replace uses of deprecated JAX sharding APIs with their new names in jax.sharding.

hawkinsp committed 2 years ago

398e170b

improve nn.checkpoint docs

cgarciae committed 2 years ago

211c48e0

Merge pull request #2815 from chiamp:initializers_refactor

a-googler committed 2 years ago

06529c9e

Merge pull request #2827 from google:update-jax-dep

a-googler committed 2 years ago

8504e5da

replaced zeros and ones initializers with zeros_init and ones_init builder functions

chiamp committed 2 years ago

2d32f287

Older