Entering jaxdb:
(jaxdb) p hidden
Array([[2.2, 2.8]], dtype=float32)
(jaxdb) c

        __
        __