google/flax

Pull Requests Commits

add test_scan_negative_axes

cgarciae committed 2 years ago

221db2f2

fix scan out_axes

cgarciae committed 2 years ago

9aa8ec68

Merge pull request #3529 from chiamp:attention

a-googler committed 2 years ago

50cd169d

Ensure that `_hashable_filter` does not convert strings to a tuple of letters

superbobry committed 2 years ago

3687af03

added sow attention weights

chiamp committed 2 years ago

9d989b0b

Merge pull request #3526 from Micky774:nnx_conv_test_trim

a-googler committed 2 years ago

829e004c

Merge pull request #3531 from chiamp:normalize

a-googler committed 2 years ago

512a6d88

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

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

Older