flax

493353c9 - Add input pipeline for ogbg_molpcba example, with unit tests for working with jraph.GraphsTuple.