google/flax

Pull Requests Commits

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

Merge pull request #2831 from cgarciae:improve-map_variables-v2

a-googler committed 3 years ago

f2094bac

Merge pull request #2769 from cgarciae:fix-tabuluate-non-array-values

a-googler committed 3 years ago

8c428faf

Update flax/linen/transforms.py

chiamp committed 3 years ago

Verified 5f67729d

improve the CasualDense example

cgarciae committed 3 years ago

93e5254e

import initializers module directly

chiamp committed 3 years ago

d696c706

Merge pull request #2818 from cgarciae:improve-map-variables-docs

a-googler committed 3 years ago

e54411d2

Merge pull request #2819 from cgarciae:cache-mypy

a-googler committed 3 years ago

cf27f3e8

Update jax dependency

jheek committed 3 years ago

c7fd31f4

Add clarification on mixing concrete and abstract vals

jheek committed 3 years ago

ed861304

Internal change

IvyZX committed 3 years ago

7ce79b3b

cache mypy cache on CI

cgarciae committed 3 years ago

8863549d

Merge pull request #2820 from IvyZX:push

a-googler committed 3 years ago

dc28d8cb

IvyZX committed 3 years ago

bc517e84

improve nn.map_variables docstrings

cgarciae committed 3 years ago

b97757cf

Implement lazy init

jheek committed 3 years ago

65ef2003

Merge pull request #2809 from cgarciae:string-methods

a-googler committed 3 years ago

a3092738

Merge pull request #2803 from cgarciae:repr-issue

a-googler committed 3 years ago

e83711be

Merge pull request #2767 from cgarciae:rename-conv-lstm

a-googler committed 3 years ago

a0a57ebf

Newer Older