egnn_jax
