wrap jax utils into a debugging module

Created by: epignatelli