jax.debug.print("step {} loss = {}", step, loss, ordered=True)

        __
        __