google/flax

Pull Requests Commits

Use jax.api.device_put_sharded() in place of private JAX apis.

vanderplas committed 5 years ago

94f322a4

Merge pull request #461 from jheek:collections

a-googler committed 5 years ago

604112f9

Merge pull request #462 from mohitreddy1996:remove-reductions

a-googler committed 5 years ago

68343360

RNGFilter -> PRNGSequenceFilter

jheek committed 5 years ago

2ddf972f

Merge pull request #419 from mohitreddy1996:lm1b-unit-test

a-googler committed 5 years ago

26b2ffac

Avoid passing python collections to jax reductions. Also remove it from pytest.ini

mohitreddy1996 committed 5 years ago

2eab0a69

'param' -> 'params'replaces usage of 'param' kind to 'params' collection

jheek committed 5 years ago

81e5c6aa

Remove kinds in favor of collections

jheek committed 5 years ago

dca090d0

Add unittest for LM1B example.

mohitreddy1996 committed 5 years ago

436fa62b

Update README.md

avital committed 5 years ago

Verified 9c295348

Merge pull request #460 from avital:linen-readme

a-googler committed 5 years ago

72641cfe

Add Linen alpha README

avital committed 5 years ago

da975779

Fix LARS optimizer for zero param initialization.

a-googler committed 5 years ago

9ff5feb3

Merge pull request #456 from jheek:fix-frozen-dict

a-googler committed 5 years ago

cb907278

jheek committed 5 years ago

51b34d80

Fix frozen dict

jheek committed 5 years ago

491e9fc0

Merge pull request #453 from mbz:patch-3

a-googler committed 5 years ago

19ea3d02

Merge pull request #451 from jheek:unified-variable-axes

a-googler committed 5 years ago

3609a43f

Exposing more Conv flags from ConvLSTM.

mbz committed 5 years ago

Verified e40f74c4

Merge pull request #422 from mbz:patch-1

a-googler committed 5 years ago

e119c9fc

avital committed 5 years ago

6ebb718d

Merge pull request #444 from vballoli:single_conv

a-googler committed 5 years ago

d5bbe379

unify variable_in_axes and variable_out_axes

jheek committed 5 years ago

eecf36fc

Merge pull request #450 from marcvanzee:get-initial-params

a-googler committed 5 years ago

7a37dd8f

Merge pull request #436 from marcvanzee:errormsg

a-googler committed 5 years ago

efc4ba74

named_call uses layer name rather than layer class name.

levskaya committed 5 years ago

efe9d0a6

Fix module test after change in error message

marcvanzee committed 5 years ago

0f5dcefa

Merge branch 'master' of https://github.com/google/flax into errormsg

marcvanzee committed 5 years ago

bac79a0b

Renames get_param to get_initial_params in seq2seq and mnist Linen examples

marcvanzee committed 5 years ago

c33df578

fix import bug and enable profiling via environment variable.

levskaya committed 5 years ago

0fb2b89d

Older