jax.config.update("jax_debug_nans", True)

        __
        __