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

e97ffe35

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

a-googler committed 116 days ago

dcfabd05

Update tensorflow dependency versions in pyproject.toml

vfdev-5 committed 116 days ago

Verified 395e0331

update to v0.12.8

Cristian Garcia committed 119 days ago

3a491108

optimize jit_partial

Cristian Garcia committed 121 days ago

0a0aca74

Merge pull request #5507 from google:add-dependabot

a-googler committed 121 days ago

0682fe4d

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

a-googler committed 121 days ago

6fec30fa

vfdev-5 committed 121 days ago

5768bb59

Add Dependabot configuration for dependency updates

vfdev-5 committed 122 days ago

Verified 0a682323

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

a-googler committed 122 days ago

e48ce367

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

a-googler committed 122 days ago

4c59d8f3

fix MultiMetric GraphDef not hashable under jit (#5489)

mohsinm-dev committed 128 days ago

Verified 6f745922

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

a-googler committed 133 days ago

5ab9d146

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

vfdev-5 committed 133 days ago

Verified 4ad282fb

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

a-googler committed 136 days ago

77a603be

Added image segmentation example

vfdev-5 committed 136 days ago

094b1a13

Added ViT training example

vfdev-5 committed 136 days ago

a450307f

Fixed gemma example sampler test test_models_from_kaggle

vfdev-5 committed 138 days ago

Verified ec75218c

optimize transform updates

Cristian Garcia committed 148 days ago

0f128b51

add storage backend type in saving telemetry.

lyglst committed 149 days ago

218c7ffc

Merge pull request #5454 from samanklesaria:dataloader_docs

a-googler committed 150 days ago

8485ec3c

Merge pull request #5363 from samanklesaria:rng_guide_update

a-googler committed 154 days ago

d044de84

Add data loaders tutorial

samanklesaria committed 154 days ago

64145c4e

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

a-googler committed 157 days ago

e9aaebe5

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

a-googler committed 157 days ago

d07c450d

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

a-googler committed 157 days ago

2ee10c2b

Merge pull request #5425 from samanklesaria:tensorboard_example

a-googler committed 157 days ago

a9057b88

Cristian Garcia committed 158 days ago

e9717a14

Merge pull request #5434 from google:mha-out_sharding

a-googler committed 158 days ago

2fa7e66e

Merge branch 'main' into mha-out_sharding

vfdev-5 committed 165 days ago

Verified 502012b8

Older