google/flax

Pull Requests Commits

Typo fix in _cached_partial method

vfdev-5 committed 8 days ago

c72bbc47

Merge pull request #5163 from google:fix-CI-numpy-version-lt-240

a-googler committed 8 days ago

9048b4e1

Set numpy<2.4 to fix #5162

vfdev-5 committed 8 days ago

0cbe1923

[pmap] Update `jax_utils.unreplicate` for performant use when `jax.config.jax_pmap_shmap_merge` is `True`.

danielsuo committed 20 days ago

3f2ff40e

Merge pull request #5149 from IvyZX:release

a-googler committed 25 days ago

d1791844

IvyZX committed 25 days ago

6d36623d

Merge pull request #5144 from google:fix-ci-5143

a-googler committed 26 days ago

9213cfdb

Temporary fix for failing CI

vfdev-5 committed 26 days ago

de853173

Merge pull request #5102 from samanklesaria:output_sharding

a-googler committed 32 days ago

ef78d658

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

a-googler committed 32 days ago

be1db789

samanklesaria committed 33 days ago

66544200

Copybara import of the project:

a-googler committed 34 days ago

3139e969

Merge pull request #5057 from jorisSchaller:feat/nnx_pooling

a-googler committed 34 days ago

de84eb15

remove fingerprint

Cristian Garcia committed 35 days ago

b8aa99d4

remove abstracted_axes from nnx.jit

Cristian Garcia committed 35 days ago

b3e9b449

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

danielsuo committed 35 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 39 days ago

872d50b8

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

danielsuo committed 39 days ago

881b8ed4

Refactor pooling operation

jorisSchaller committed 39 days ago

004b9324

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

Jake VanderPlas committed 40 days ago

ec85cddf

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

danielsuo committed 41 days ago

59cfd99e

[flax:examples:sst2] Fix notebook error.

danielsuo committed 41 days ago

5a3d1489

Merge pull request #5094 from thijs-vanweezel:filters

a-googler committed 42 days ago

a5af24c2

[flax:examples:sst2] Fix pytype errors.

danielsuo committed 42 days ago

6ac5a783

[flax:examples:gemma] Fixing linter errors.

danielsuo committed 42 days ago

324f4b4b

[flax:examples:vae] Small linter fixes.

danielsuo committed 43 days ago

88d74979

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

danielsuo committed 43 days ago

9ed59e2f

[flax:examples:wmt] Small linter fixes.

danielsuo committed 43 days ago

c29b6d6b

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

a-googler committed 55 days ago

697f4e5c

Add out_sharding argument to call methods for standard layers

samanklesaria committed 55 days ago

1f61d6b7

Older