flax
Optimize Flax ogbg_molpcba example GNN model with custom Pallas Fused LayerNorm kernel, custom VJP for reverse-mode autodiff, and numerical correctness benchmark script.
#5513
Open
Open
Optimize Flax ogbg_molpcba example GNN model with custom Pallas Fused LayerNorm kernel, custom VJP for reverse-mode autodiff, and numerical correctness benchmark script. #5513
copybara-service
wants to merge
1 commit into
main
from
test_936193070
Optimize Flax ogbg_molpcba example GNN model with custom Pallas Fused…
e97ffe35
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Login to write a write a comment.
Login via GitHub