google/flax

Pull Requests Commits

flatten respect nnx.data

cgarciae committed 107 days ago

6dd9baa8

Merge pull request #5102 from samanklesaria:output_sharding

a-googler committed 108 days ago

ef78d658

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

a-googler committed 108 days ago

be1db789

samanklesaria committed 108 days ago

66544200

Copybara import of the project:

a-googler committed 110 days ago

3139e969

Merge pull request #5057 from jorisSchaller:feat/nnx_pooling

a-googler committed 110 days ago

de84eb15

remove fingerprint

Cristian Garcia committed 110 days ago

b8aa99d4

remove abstracted_axes from nnx.jit

Cristian Garcia committed 110 days ago

b3e9b449

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

danielsuo committed 111 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 114 days ago

872d50b8

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

danielsuo committed 114 days ago

881b8ed4

Refactor pooling operation

jorisSchaller committed 114 days ago

004b9324

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

Jake VanderPlas committed 115 days ago

ec85cddf

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

danielsuo committed 117 days ago

59cfd99e

[flax:examples:sst2] Fix notebook error.

danielsuo committed 117 days ago

5a3d1489

Merge pull request #5094 from thijs-vanweezel:filters

a-googler committed 117 days ago

a5af24c2

[flax:examples:sst2] Fix pytype errors.

danielsuo committed 118 days ago

6ac5a783

[flax:examples:gemma] Fixing linter errors.

danielsuo committed 118 days ago

324f4b4b

[flax:examples:vae] Small linter fixes.

danielsuo committed 118 days ago

88d74979

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

danielsuo committed 118 days ago

9ed59e2f

[flax:examples:wmt] Small linter fixes.

danielsuo committed 118 days ago

c29b6d6b

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

a-googler committed 130 days ago

697f4e5c

Add out_sharding argument to call methods for standard layers

samanklesaria committed 130 days ago

1f61d6b7

update jax minver to 0.8.1

cgarciae committed 130 days ago

b424c49e

hint that key is cast to string

thijs-vanweezel committed 130 days ago

a7bb109b

test pathcontains with exact=false

thijs-vanweezel committed 130 days ago

12b2e42f

allow substring matching

thijs-vanweezel committed 130 days ago

3d340558

Add is_causal mask argument and tests

ibbyml committed 131 days ago

9bab1485

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

a-googler committed 131 days ago

bee81ed4

Fixed deprecated .value usage failing CI tests

vfdev-5 committed 131 days ago

00cf44e5

Older