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 118 days ago

e97ffe35

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

a-googler committed 118 days ago

dcfabd05

Update tensorflow dependency versions in pyproject.toml

vfdev-5 committed 118 days ago

Verified 395e0331

update to v0.12.8

Cristian Garcia committed 121 days ago

3a491108

optimize jit_partial

Cristian Garcia committed 123 days ago

0a0aca74

Merge pull request #5507 from google:add-dependabot

a-googler committed 123 days ago

0682fe4d

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

a-googler committed 123 days ago

6fec30fa

vfdev-5 committed 123 days ago

5768bb59

Add Dependabot configuration for dependency updates

vfdev-5 committed 123 days ago

Verified 0a682323

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

a-googler committed 124 days ago

e48ce367

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

a-googler committed 124 days ago

4c59d8f3

fix MultiMetric GraphDef not hashable under jit (#5489)

mohsinm-dev committed 130 days ago

Verified 6f745922

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

a-googler committed 135 days ago

5ab9d146

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

vfdev-5 committed 135 days ago

Verified 4ad282fb

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

a-googler committed 138 days ago

77a603be

Added image segmentation example

vfdev-5 committed 138 days ago

094b1a13

Added ViT training example

vfdev-5 committed 138 days ago

a450307f

Fixed gemma example sampler test test_models_from_kaggle

vfdev-5 committed 139 days ago

Verified ec75218c

optimize transform updates

Cristian Garcia committed 150 days ago

0f128b51

add storage backend type in saving telemetry.

lyglst committed 151 days ago

218c7ffc

Merge pull request #5454 from samanklesaria:dataloader_docs

a-googler committed 152 days ago

8485ec3c

Merge pull request #5363 from samanklesaria:rng_guide_update

a-googler committed 156 days ago

d044de84

Add data loaders tutorial

samanklesaria committed 156 days ago

64145c4e

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

a-googler committed 159 days ago

e9aaebe5

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

a-googler committed 159 days ago

d07c450d

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

a-googler committed 159 days ago

2ee10c2b

Merge pull request #5425 from samanklesaria:tensorboard_example

a-googler committed 159 days ago

a9057b88

Cristian Garcia committed 160 days ago

e9717a14

Merge pull request #5434 from google:mha-out_sharding

a-googler committed 160 days ago

2fa7e66e

Merge branch 'main' into mha-out_sharding

vfdev-5 committed 167 days ago

Verified 502012b8

Older