google/flax

Pull Requests Commits

Add pytrees as graph nodes again.

Cristian Garcia committed 1 year ago

9c8f9362

Copybara import of the project:

a-googler committed 1 year ago

e3789de2

Add Sow Config to from_params constructor.

casaro committed 1 year ago

64635457

Fix position and name of Post Attention Norm.

casaro committed 1 year ago

b28822f3

add top_p sampling in gemma example

tianshub committed 1 year ago

a149b6d7

Add option to load checkpoints with transposed Gating Einsum.

casaro committed 1 year ago

ceb7aedd

Merge pull request #4547 from google:nnx-index-pytrees

a-googler committed 1 year ago

bf4c322a

Make RoPE Base Frequency configurable.

casaro committed 1 year ago

324bbe0f

Add configurable Query Pre Attention scalar.

casaro committed 1 year ago

fd32d453

Merge pull request #4584 from IvyZX:linx-misc

a-googler committed 1 year ago

8254dd0a

bridge module with nnx submodules

IvyZX committed 1 year ago

06a50e88

casaro committed 1 year ago

0769411f

Merge pull request #4600 from google:remove-embed-special-case

a-googler committed 1 year ago

50b996d2

remove Embed nan casting

cgarciae committed 1 year ago

0933c2e8

Merge pull request #4592 from jakevdp:fix-shape

a-googler committed 1 year ago

45a8f846

Avoid calls to jnp.shape for non-array inputs.

jakevdp committed 1 year ago

4e300d4c

[bridge] improve sow and add initializers

Cristian Garcia committed 1 year ago

a24d790e

Refactor bridge.Module tests from `wrappers_test.py` to another file.

IvyZX committed 1 year ago

6af8fcb0

Merge pull request #4580 from google:nnx-tabulate-issue

a-googler committed 1 year ago

7974b106

Merge pull request #4579 from google:v0.10.4

a-googler committed 1 year ago

7f175f80

[nnx] fix tabulate

cgarciae committed 1 year ago

14fda9d7

cgarciae committed 1 year ago

5f4f1b84

[bridge] Set _initializing correctly and avoid return RNG states

IvyZX committed 1 year ago

54138501

Merge pull request #4576 from IvyZX:fix-head

a-googler committed 1 year ago

d96be6ce

Fix breakages from newest jax

IvyZX committed 1 year ago

0f9f55d6

Merge pull request #4490 from google:nnx-shard-map

a-googler committed 1 year ago

41bef07c

Merge pull request #4561 from IvyZX:state-dep-warn

a-googler committed 1 year ago

dd6e5955

Fix all state method usages

IvyZX committed 1 year ago

ed73ba43

[nnx] add shard_map

cgarciae committed 1 year ago

10d8e5ca

Fixes spmd to work correctly with xaot compilation by using global mesh's device instead of jax.devices()[0]

marksandler2 committed 1 year ago

1ec5ef29

Older