google/flax

Pull Requests Commits

Optimize Flax ogbg_molpcba example GNN model with custom Pallas Fused LayerNorm kernel, custom VJP for reverse-mode autodiff, and numerical correctness benchmark script.

a-googler committed 46 days ago

e97ffe35

Merge pull request #5511 from google:vfdev-5-patch-1

a-googler committed 46 days ago

dcfabd05

Update tensorflow dependency versions in pyproject.toml

vfdev-5 committed 46 days ago

Verified 395e0331

update to v0.12.8

Cristian Garcia committed 49 days ago

3a491108

optimize jit_partial

Cristian Garcia committed 51 days ago

0a0aca74

Merge pull request #5507 from google:add-dependabot

a-googler committed 51 days ago

0682fe4d

Merge pull request #5506 from google:fix-ci-nest-asyncio

a-googler committed 51 days ago

6fec30fa

vfdev-5 committed 51 days ago

5768bb59

Add Dependabot configuration for dependency updates

vfdev-5 committed 52 days ago

Verified 0a682323

Merge pull request #5455 from google:docs-examples--vit

a-googler committed 52 days ago

e48ce367

Merge pull request #5492 from mohsinm-dev:fix-multimetric-hashable-graphdef

a-googler committed 52 days ago

4c59d8f3

fix MultiMetric GraphDef not hashable under jit (#5489)

mohsinm-dev committed 58 days ago

Verified 6f745922

Merge pull request #5477 from google:fix-gemma-sampler-test-kaggle

a-googler committed 63 days ago

5ab9d146

Merge branch 'main' into docs-examples--vit

vfdev-5 committed 63 days ago

Verified 4ad282fb

Merge pull request #5463 from google:docs-examples--unetr

a-googler committed 66 days ago

77a603be

Added image segmentation example

vfdev-5 committed 66 days ago

094b1a13

Added ViT training example

vfdev-5 committed 66 days ago

a450307f

Fixed gemma example sampler test test_models_from_kaggle

vfdev-5 committed 68 days ago

Verified ec75218c

optimize transform updates

Cristian Garcia committed 78 days ago

0f128b51

add storage backend type in saving telemetry.

lyglst committed 79 days ago

218c7ffc

Merge pull request #5454 from samanklesaria:dataloader_docs

a-googler committed 80 days ago

8485ec3c

Merge pull request #5363 from samanklesaria:rng_guide_update

a-googler committed 84 days ago

d044de84

Add data loaders tutorial

samanklesaria committed 84 days ago

64145c4e

Merge pull request #5441 from samanklesaria:fix-transforms-doctest

a-googler committed 87 days ago

e9aaebe5

Merge pull request #5440 from samanklesaria:fix-graphlib-doctest

a-googler committed 87 days ago

d07c450d

Merge pull request #5442 from samanklesaria:fix-rnglib-doctest

a-googler committed 87 days ago

2ee10c2b

Merge pull request #5425 from samanklesaria:tensorboard_example

a-googler committed 87 days ago

a9057b88

Cristian Garcia committed 88 days ago

e9717a14

Merge pull request #5434 from google:mha-out_sharding

a-googler committed 88 days ago

2fa7e66e

Merge branch 'main' into mha-out_sharding

vfdev-5 committed 95 days ago

Verified 502012b8

Older