flax

Add jraph dependency for upcoming ogbg_molpcba example.

#1485

Merged