google/flax

Pull Requests Commits

flatten respect nnx.data

cgarciae committed 3 days ago

6dd9baa8

Merge pull request #5102 from samanklesaria:output_sharding

a-googler committed 4 days ago

ef78d658

Merge pull request #5093 from ibbyml:nnx-sdpa-parity

a-googler committed 4 days ago

be1db789

samanklesaria committed 4 days ago

66544200

Copybara import of the project:

a-googler committed 6 days ago

3139e969

Merge pull request #5057 from jorisSchaller:feat/nnx_pooling

a-googler committed 6 days ago

de84eb15

remove fingerprint

Cristian Garcia committed 6 days ago

b8aa99d4

remove abstracted_axes from nnx.jit

Cristian Garcia committed 6 days ago

b3e9b449

[jax:benchmarks] Add tracing/lowering benchmarks for a few flax examples.

danielsuo committed 7 days ago

2204c82d

[flax:examples:nlp_seq] Create a main.py file to run tests with config files to match other examples. #jax-fixit

danielsuo committed 10 days ago

872d50b8

[flax:examples:lm1b_nnx] Update example to work internally. #jax-fixit.

danielsuo committed 10 days ago

881b8ed4

Refactor pooling operation

jorisSchaller committed 11 days ago

004b9324

Avoid passing `concrete` argument to `jax.remat`

Jake VanderPlas committed 11 days ago

ec85cddf

[flax:examples:ppo] Fix some linter / import issues. #jax-fixit

danielsuo committed 13 days ago

59cfd99e

[flax:examples:sst2] Fix notebook error.

danielsuo committed 13 days ago

5a3d1489

Merge pull request #5094 from thijs-vanweezel:filters

a-googler committed 13 days ago

a5af24c2

[flax:examples:sst2] Fix pytype errors.

danielsuo committed 14 days ago

6ac5a783

[flax:examples:gemma] Fixing linter errors.

danielsuo committed 14 days ago

324f4b4b

[flax:examples:vae] Small linter fixes.

danielsuo committed 14 days ago

88d74979

[flax:examples:seq2seq] Create main and default config based on seq2seq.ipynb.

danielsuo committed 14 days ago

9ed59e2f

[flax:examples:wmt] Small linter fixes.

danielsuo committed 14 days ago

c29b6d6b

Merge pull request #5095 from google:jax-minver-0.8.1

a-googler committed 26 days ago

697f4e5c

Add out_sharding argument to call methods for standard layers

samanklesaria committed 26 days ago

1f61d6b7

update jax minver to 0.8.1

cgarciae committed 26 days ago

b424c49e

hint that key is cast to string

thijs-vanweezel committed 26 days ago

a7bb109b

test pathcontains with exact=false

thijs-vanweezel committed 27 days ago

12b2e42f

allow substring matching

thijs-vanweezel committed 27 days ago

3d340558

Add is_causal mask argument and tests

ibbyml committed 27 days ago

9bab1485

Merge pull request #5097 from google:fix-main-4

a-googler committed 27 days ago

bee81ed4

Fixed deprecated .value usage failing CI tests

vfdev-5 committed 27 days ago

00cf44e5

Older