segnn_jax
