google/flax

Pull Requests Commits

remove deprecated normalize function

chiamp committed 2 years ago

28d41d94

Break apart flax pack function into the multiple scope handling functions to

pschuh committed 2 years ago

230b0d77

Squashed commits, reduced test count

Micky774 committed 2 years ago

18568bbe

Allow reduction axes in GroupNorm to be defined explicitly at layer construction

a-googler committed 2 years ago

d3b2cbb2

allow int key in Module/StateVariablesMapping getitem

cgarciae committed 2 years ago

b4f7bed5

add State.raw_mapping

cgarciae committed 2 years ago

1963668e

add StateVariablesMapping

cgarciae committed 2 years ago

471b956a

Merge pull request #3518 from google:nnx-fix-graph-flatten

a-googler committed 2 years ago

c0a2657a

fix graph flatten/unflatten with Variable sharing

cgarciae committed 2 years ago

a78bdede

Merge pull request #3510 from carlosgmartin:normalization_mask

a-googler committed 2 years ago

c2c3feaa

Merge pull request #3509 from Micky774:nnx_linear_api_tests

a-googler committed 2 years ago

68333d4c

Merge pull request #3516 from google:nnx-variable-tying

a-googler committed 2 years ago

43528790

Add missing mask argument to LayerNorm, RMSNorm, and GroupNorm.

carlosgmartin committed 2 years ago

59c6c30c

Merge pull request #3514 from NeilGirdhar:initbug

a-googler committed 2 years ago

ca9606bc

Merge pull request #3511 from Micky774:nnx_conv_api_tests

a-googler committed 2 years ago

cdc7feba

Merge pull request #3513 from Micky774:nnx_embed_api_test

a-googler committed 2 years ago

e50b1b53

Merge branch 'main' into nnx_linear_api_tests

Micky774 committed 2 years ago

3024a140

allow shared Variables

cgarciae committed 2 years ago

818cd79c

Prevent crash in dataclasses with no-init params

NeilGirdhar committed 2 years ago

2ef87839

Merge pull request #3487 from google:nnx-mha

a-googler committed 2 years ago

a572f6af

Merge branch 'main' into nnx_embed_api_test

Micky774 committed 2 years ago

b74e3a6e

Merge branch 'main' into nnx_conv_api_tests

Micky774 committed 2 years ago

7a235519

Squashed commits

Micky774 committed 2 years ago

b8088321

Merge pull request #3512 from HMUNACHI:geglu

a-googler committed 2 years ago

632e6c0a

fixed minor bug in original commit

HenryNdubuaku committed 2 years ago

0cac9fcd

Add NNX/Linen consistency test for `Embed` layer

Micky774 committed 2 years ago

9ad4d3f7

added geglu activation and tests

HenryNdubuaku committed 2 years ago

c5ca7d71

Add NNX/Linen API consistency test for `Conv` layer

Micky774 committed 2 years ago

0daa8dfb

Merge pull request #3498 from Micky774:prelu_correctness_check

a-googler committed 2 years ago

c3c4cbf7

simplify MultiHeadAttention API

cgarciae committed 2 years ago

0a1f78a3

Newer Older