google/flax

Pull Requests Commits

Update einsum layer for Gemma example

a-googler committed 1 year ago

0eeab5aa

Merge pull request #4499 from google:nnx-improve-module-docs

a-googler committed 1 year ago

e3bcc442

[nnx] improve Module docs

cgarciae committed 1 year ago

bdcc33a5

Remove usages of orbax_utils.save_args_from_target, as this function does nothing (it used to control a checkpointing behavior that has since been optimized away).

cpgaffney1 committed 1 year ago

b5d4ed85

Merge pull request #4486 from google:fix-to-nnx-linen-attributes

a-googler committed 1 year ago

e4418e27

[nnx] fix ToNNX linen_attributes update

cgarciae committed 1 year ago

342a5e83

Merge pull request #4485 from google:nnx-multimetric

a-googler committed 1 year ago

969f3cb2

[nnx] fix MultiMetric typing

cgarciae committed 1 year ago

ca4dc8ca

Merge pull request #4479 from jakevdp:fix-workflows

a-googler committed 1 year ago

1961c125

Merge pull request #4478 from jakevdp:nightly-build

a-googler committed 1 year ago

6abf8bcb

CI: pin actions to specific commits

jakevdp committed 1 year ago

9a362592

CI: add scheduled test against nightly JAX releases

jakevdp committed 1 year ago

44efae32

[nnx] add state summaries for print and display

Cristian Garcia committed 1 year ago

3d1f55a0

Copybara import of the project:

a-googler committed 1 year ago

62688664

Merge pull request #4438 from google:nnx-tabulate

a-googler committed 1 year ago

adbad95e

Allow `nnx.bridge.variables.nnx_attrs_to_linen_vars` take `nnx.VariableState` as argument.

IvyZX committed 1 year ago

9ebdbdc2

[nnx] add tabulate

cgarciae committed 1 year ago

c59bc1a9

Merge pull request #4407 from google:rnn-broadcast-rngs

a-googler committed 1 year ago

5408c40c

Makes flax Modules more compatible with IPython auto-reload.

a-googler committed 1 year ago

e2134afa

Merge pull request #4199 from 8bitmp3:upgrade-nnx-filters

a-googler committed 1 year ago

dd70960a

Merge pull request #4472 from google:nnx-fix-fori

a-googler committed 1 year ago

595e711f

[nnx] test fori_loop output

cgarciae committed 1 year ago

f8164dd4

Merge pull request #4411 from 8bitmp3:update-nnx-merge-docs

a-googler committed 1 year ago

53bde74a

Merge pull request #4325 from 8bitmp3:update-nnx-gemma

a-googler committed 1 year ago

cc29a7b3

Update NNX merge docs in graph.py

8bitmp3 committed 1 year ago

0a40dcdb

Upgrade Flax NNX Gemma

8bitmp3 committed 1 year ago

a46578ec

Introduce a mode for a faster nn.scan mode that avoids extra jax retracing.

levskaya committed 1 year ago

fc38f21e

Merge pull request #4286 from 8bitmp3:update-nnx-jax-transforms

a-googler committed 1 year ago

207966e5

Update Flax NNX vs JAX Transformations guide

8bitmp3 committed 1 year ago

c8bfd2a6

Upgrade NNX Filters guides

8bitmp3 committed 1 year ago

9dd1e4eb

Newer Older