google/flax

Pull Requests Commits

Use Sequence rather than Tuple in ConvTranspose type hints.

a-googler committed 3 years ago

36542787

Merge branch 'google:main' into dev_branch

cpgaffney1 committed 3 years ago

Verified 2637428d

Revised docs using Flax contributing guide

cpgaffney1 committed 3 years ago

a96f49ba

Flax layers are parametrizable with custom conv_general_dilated.

a-googler committed 3 years ago

92b0cdd4

Remove fallbacks for old versions of JAX from Flax.

hawkinsp committed 3 years ago

323d6b85

Flax layers are parametrizable with custom dot_general.

a-googler committed 3 years ago

1f6b0949

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

hawkinsp committed 3 years ago

26d46e4b

Merge pull request #2856 from levskaya:point_release4

a-googler committed 3 years ago

524629eb

Flax 0.6.5 point release

levskaya committed 3 years ago

5c6f4267

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

a-googler committed 3 years ago

6ea0aca8

Merge pull request #2816 from jheek:lazy-init

a-googler committed 3 years ago

24224bff

Merge pull request #2811 from arcAman07:main

a-googler committed 3 years ago

c3c4a876

fix temperature_sampler + add test

cgarciae committed 3 years ago

2e1232e1

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

hawkinsp committed 3 years ago

398e170b

improve nn.checkpoint docs

cgarciae committed 3 years ago

211c48e0

Merge pull request #2815 from chiamp:initializers_refactor

a-googler committed 3 years ago

06529c9e

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

a-googler committed 3 years ago

8504e5da

replaced zeros and ones initializers with zeros_init and ones_init builder functions

chiamp committed 3 years ago

2d32f287

Merge pull request #2843 from IvyZX:meta

a-googler committed 3 years ago

c06cba28

Make a separate Python module for Orbax-related utils.

IvyZX committed 3 years ago

694560bb

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

hawkinsp committed 3 years ago

431787b6

Re-box axis metadata in DenseGeneral's init functions

IvyZX committed 3 years ago

ba87b468

Apply suggestions from code review

jheek committed 3 years ago

Verified 6ffdd831

Merge pull request #2830 from chiamp:import_initializers

a-googler committed 3 years ago

aaba5991

Bumps minimal JAX version from 0.3.16 to 0.3.24 and cleans up some code.

marcvanzee committed 3 years ago

83a5b3ba

Clarify we are using the ReZero trick in our ResNet example.

marcvanzee committed 3 years ago

1b5b504b

Bumps minimal JAX version from 0.3.16 to 0.3.24 and cleans up some code.

marcvanzee committed 3 years ago

06323f27

Apply suggestions from code review

jheek committed 3 years ago

Verified f4da05b9

Merge pull request #2833 from levskaya:relaxed_naming

a-googler committed 3 years ago

dfb55c4d

Introduce relaxed naming policy behind a feature flag.

levskaya committed 3 years ago

54e342b4

Newer Older