google/flax

Pull Requests Commits

Add count_include_pad argument to flax.linen.pooling.avg_pool

dslisleedh committed 3 years ago

ff97879c

Add count_include_pad argument to flax.linen.pooling.avg_pool

dslisleedh committed 3 years ago

9f5ad668

internal change only

a-googler committed 3 years ago

8687673d

Remove pytype generic workaround

jheek committed 3 years ago

70d845a8

Merge pull request #2388 from cgarciae:simplify-dynamic-context

a-googler committed 3 years ago

fdd1d6fe

Merge pull request #2367 from cgarciae:move-example-to-rtd

a-googler committed 3 years ago

e0baf0d3

Merge pull request #2427 from villeh1:main

a-googler committed 3 years ago

6784fbd7

simplify _DynamicContext

cgarciae committed 3 years ago

9ff36a11

clean examples doc page

cgarciae committed 3 years ago

cb9715dd

add attribute docs to Conv and ConvLocal

cgarciae committed 3 years ago

9a9539b8

PIL does not accept DeviceArray, so needed to use numpy.

villeh1 committed 3 years ago

Verified fb4699f4

Explicitly create GDA directories when saving checkpoints.

IvyZX committed 3 years ago

0451a55b

move examples to RTD + improve examples

cgarciae committed 3 years ago

c1be51ba

Revert change that breaks multi-host serialization with non-replicated data

jheek committed 3 years ago

6781c596

Merge pull request #2371 from cgarciae:optax-partition-function

a-googler committed 3 years ago

b146e16a

Use new name add_decayed_weights instead of additive_weight_decay

mtthss committed 3 years ago

d4bac102

Make sure save checkpoint only happens on process_index 0. This will make save_checkpoint truly multihost-friendly. Also made some refactoring & extra commenting.

IvyZX committed 3 years ago

3508ecf4

add path_aware_map

cgarciae committed 3 years ago

acb7662b

Merge pull request #2316 from cgarciae:improve-tabulate

a-googler committed 3 years ago

14dffa6d

Merge pull request #2409 from lkhphuc:patch-2

a-googler committed 3 years ago

9fe38650

Support multiple inputs in flax lifted vjp/custom_vjp

kho committed 3 years ago

cbbf5984

Explicit reexport initializers from jax

lkhphuc committed 3 years ago

Verified fab2eb99

Plumb spmd_axis_name through transforms.vmap through to JAX vmap

James Lee-Thorp committed 3 years ago

07e513f6

Adds axis_name and axis_index_groups to LayerNorm and GroupNorm.

a-googler committed 3 years ago

2b73efdd

Merge pull request #2407 from yklcs:bump-rich

a-googler committed 3 years ago

c3e5363d

Merge pull request #2405 from andsteing:examples

a-googler committed 3 years ago

5c09c3a3

Ease rich version requirement

yklcs committed 3 years ago

72189153

Updates examples/{imagenet,wmt} requirements.

andsteing committed 3 years ago

f175457a

Removes `flax.optim` & bumps version to 0.6.0

andsteing committed 3 years ago

cda7a4c8

Merge pull request #2392 from banda-larga:patch-1

a-googler committed 3 years ago

47418b32

Newer Older