pgx Change examples away from Haiku

A bit later than I meant to, but this addresses https://github.com/sotetsuk/pgx/issues/1059. Pretty straightforward change to equinox style NNs. Some minor speed differences that could be optimized (see https://github.com/patrick-kidger/equinox/issues/928, https://github.com/patrick-kidger/equinox/issues/926), but IMO it's fast enough. This PR also depends on the following PR to equinox to update its batch norm code to be more standard with other jax libraries: https://github.com/patrick-kidger/equinox/pull/948.

Full report comparison haiku and equinox should be available here (note, the results for Haiku were just done on main, also for minatar I used 1xA100 and for AZ I used 4xA100): https://api.wandb.ai/links/rl-exps/fak1nv1n

Feb 10 '25 05:02 lockwo

Awesome! 🤩 Thank you for your PR! I'll check it this weekend 🙏

Feb 12 '25 14:02 sotetsuk

hey guys, I'm glad I found this, as I was struggling to migrate pgx-alphazero code to equinox.. Is this to be merged soon?

@lockwo the wandb link seems to have expired. I'm curious about your go-9x9 run. Do you happen to have any details somewhere else? Thanks

Sep 01 '25 18:09 hyu2000

Oops, I think I cleaned out that WB account recently. I can see about re-running the experiments if needed. For the 9x9 go, there just so happens to be a screenshot of the results in the old equinox PR (https://github.com/patrick-kidger/equinox/pull/948, just look at the batch norm that worked lol since that's the one in use now).

Code wise, it's the exact same as in the PR (no special sauce beyond that)

Sep 02 '25 06:09 lockwo