Is JAX A Good Fit For Geometric Deep Learning?

Short answer: not really, mostly because of the dynamic shapes of batches in Geometric Deep Learning (GDL).