Perdón si no contesté antes.
La verdad es que tendría que ver todo de cero.
Talvez hay una confusión: yo no soy el que los hizo, ni el que mantiene, ni el responsable de los notebooks.
Sólo puedo responder por el de JAX, que ese sí lo hice, y un poco por el de GNN, que lo revisé.
Podríamos consultar a quien lo dio a ver que dice. Si querés te dirijo a esa persona.
Saludos.