google/flax

Pull Requests Commits

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

danielsuo committed 296 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 299 days ago

872d50b8

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

danielsuo committed 299 days ago

881b8ed4

Refactor pooling operation

jorisSchaller committed 299 days ago

004b9324

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

Jake VanderPlas committed 300 days ago

ec85cddf

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

danielsuo committed 302 days ago

59cfd99e

[flax:examples:sst2] Fix notebook error.

danielsuo committed 302 days ago

5a3d1489

Merge pull request #5094 from thijs-vanweezel:filters

a-googler committed 302 days ago

a5af24c2

[flax:examples:sst2] Fix pytype errors.

danielsuo committed 303 days ago

6ac5a783

[flax:examples:gemma] Fixing linter errors.

danielsuo committed 303 days ago

324f4b4b

[flax:examples:vae] Small linter fixes.

danielsuo committed 303 days ago

88d74979

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

danielsuo committed 303 days ago

9ed59e2f

[flax:examples:wmt] Small linter fixes.

danielsuo committed 303 days ago

c29b6d6b

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

a-googler committed 315 days ago

697f4e5c

Add out_sharding argument to call methods for standard layers

samanklesaria committed 315 days ago

1f61d6b7

update jax minver to 0.8.1

cgarciae committed 315 days ago

b424c49e

hint that key is cast to string

weezeltggv committed 315 days ago

a7bb109b

test pathcontains with exact=false

weezeltggv committed 315 days ago

12b2e42f

allow substring matching

weezeltggv committed 315 days ago

3d340558

Add is_causal mask argument and tests

ibbyml committed 316 days ago

9bab1485

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

a-googler committed 316 days ago

bee81ed4

Fixed deprecated .value usage failing CI tests

vfdev-5 committed 316 days ago

00cf44e5

Merge pull request #5084 from samanklesaria:variable_eq

a-googler committed 317 days ago

09537f49

Merge pull request #5052 from google:deprecate-value-2

a-googler committed 317 days ago

b7e66b74

deprecate .value

cgarciae committed 317 days ago

a2eac321

Merge pull request #5055 from mohsinm-dev:fix-bound-method-auto-unbinding

a-googler committed 317 days ago

d6e8e14f

Merge pull request #5073 from google:fix-nnx-scan-carry-pytree-sow

a-googler committed 317 days ago

544a0505

Merge pull request #5071 from chapman20j:set_mode_info

a-googler committed 317 days ago

0b1d841c

Merge branch 'main' into pr/5055

samanklesaria committed 317 days ago

dece6b34

Merge pull request #5090 from google:fix-main-3

a-googler committed 320 days ago

5109e2c5

Newer Older