google/flax

Pull Requests Commits

Add pytrees as graph nodes again.

Cristian Garcia committed 298 days ago

9c8f9362

Copybara import of the project:

a-googler committed 298 days ago

e3789de2

Add Sow Config to from_params constructor.

casaro committed 298 days ago

64635457

Fix position and name of Post Attention Norm.

casaro committed 298 days ago

b28822f3

add top_p sampling in gemma example

tianshub committed 298 days ago

a149b6d7

Add option to load checkpoints with transposed Gating Einsum.

casaro committed 298 days ago

ceb7aedd

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

a-googler committed 298 days ago

bf4c322a

Make RoPE Base Frequency configurable.

casaro committed 298 days ago

324bbe0f

Add configurable Query Pre Attention scalar.

casaro committed 298 days ago

fd32d453

Merge pull request #4584 from IvyZX:linx-misc

a-googler committed 299 days ago

8254dd0a

bridge module with nnx submodules

IvyZX committed 299 days ago

06a50e88

casaro committed 300 days ago

0769411f

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

a-googler committed 300 days ago

50b996d2

remove Embed nan casting

cgarciae committed 300 days ago

0933c2e8

Merge pull request #4592 from jakevdp:fix-shape

a-googler committed 300 days ago

45a8f846

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

jakevdp committed 301 days ago

4e300d4c

[bridge] improve sow and add initializers

Cristian Garcia committed 304 days ago

a24d790e

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

IvyZX committed 306 days ago

6af8fcb0

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

a-googler committed 306 days ago

7974b106

Merge pull request #4579 from google:v0.10.4

a-googler committed 306 days ago

7f175f80

[nnx] fix tabulate

cgarciae committed 306 days ago

14fda9d7

cgarciae committed 306 days ago

5f4f1b84

[bridge] Set _initializing correctly and avoid return RNG states

IvyZX committed 307 days ago

54138501

Merge pull request #4576 from IvyZX:fix-head

a-googler committed 307 days ago

d96be6ce

Fix breakages from newest jax

IvyZX committed 307 days ago

0f9f55d6

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

a-googler committed 308 days ago

41bef07c

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

a-googler committed 311 days ago

dd6e5955

Fix all state method usages

IvyZX committed 311 days ago

ed73ba43

[nnx] add shard_map

cgarciae committed 311 days ago

10d8e5ca

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

marksandler2 committed 312 days ago

1ec5ef29

Older