google/flax

Pull Requests Commits

Use regular indexing in nn.Embed

romanngg committed 2 years ago

76036460

Add FLOPs calculations to getting started example

romanngg committed 2 years ago

54a42daf

Report forward and backward pass FLOPs of modules and submodules in `linen.Module.tabulate` and `summary.tabulate` (in new `flops` and `vjp_flops` table columns). Pass `compute_flops=True` and/or `compute_vjp_flops=True` to include these columns.

romanngg committed 2 years ago

4c7a45a9

Merge pull request #3374 from levskaya:fp8fix

a-googler committed 2 years ago

f39de8fa

Make fp8 ops use explicit broadcasting.

levskaya committed 2 years ago

2c2bd583

Merge pull request #3373 from levskaya:wsfix

a-googler committed 2 years ago

60f52944

Trailing whitespace fixes.

levskaya committed 2 years ago

8f9d097d

Merge pull request #3371 from chiamp:typing

a-googler committed 2 years ago

c5b568f6

Merge pull request #3322 from wenscarl:shuw/add_fp8_support

a-googler committed 2 years ago

cf7d89a3

fixed scope typing

chiamp committed 2 years ago

d0175ef0

Merge pull request #3367 from chiamp:error

a-googler committed 2 years ago

242f84ca

Merge pull request #3368 from 8bitmp3:add-maxtext

a-googler committed 2 years ago

406116c1

Add MaxText open source LLM to index.rst

8bitmp3 committed 2 years ago

Verified 5d4f3b7d

chiamp committed 2 years ago

2e458817

Merge pull request #3365 from chiamp:bind

a-googler committed 2 years ago

7dba76cd

Merge pull request #3335 from chiamp:spectral_norm

a-googler committed 2 years ago

b05cbfa4

fixed bind-unbind bug

chiamp committed 2 years ago

0d9862e5

Merge pull request #3357 from google:dont-propagate-default-args-in-tabulate

a-googler committed 2 years ago

4d8b4aff

Merge pull request #3359 from chiamp:filter_warning

a-googler committed 2 years ago

fd99fdb2

make DenyList('intermediates') the default for mutable

cgarciae committed 2 years ago

fbc0207b

don't propagate default arguments from .tabulate to .init

cgarciae committed 2 years ago

9c9218d7

Fix Fp8 train state test

wenscarl committed 2 years ago

b1ee11fc

self.params replaces variable_with_axes

wenscarl committed 2 years ago

06f5f280

wenscarl committed 2 years ago

04d71b65

Clean imports and add train_states

wenscarl committed 2 years ago

2bea7a57

Add fp8 custom op and unit test

wenscarl committed 2 years ago

b309c8f2

Merge pull request #3356 from google:disable-pyink

a-googler committed 2 years ago

63420d35

cgarciae committed 2 years ago

b01402e3

cgarciae committed 2 years ago

04e64f9b

added spectral norm

chiamp committed 2 years ago

15f47bf6

Older