Evaluate migration to flax or haiku for Deep Learning support

Created by: epignatelli