[Baseline, examples] Migrate to Flax (or other NN libraries)

[sotetsuk]

[lockwo]

How important is Flax to the migration? I've been exploring pgx and it seems really nice and I've started using it, but my go to is equinox (https://github.com/patrick-kidger/equinox, probably the second most popular behind flax). I would be down to work on this issue, as long as you are ok with equinox instead of flax.

[sotetsuk]

Thank you for your interest in pgx! 🙏

Whether it's Flax or not isn't particularly important. I simply considered it as the first candidate because I thought it's the framework used by the largest community.

I think Equinox is a good choice. The issue I had with Haiku was that batch normalization couldn't handle dimensions that were made by vmap, so inevitably, I had to initialize models that included batch dimensions. Equinox doesn't have this issue, so I think it's wonderful.

If you could send a PR, I would greatly appreciate it. If you send a PR, I'll check the following points on my end and then merge:

• A decrease in NN inference speed due to the switch from Haiku to Equinox is acceptable up to about double the time.
• It would be preferable to have similar learning performance with 400 iterations on Go 9x9 (although I think a few percent difference would be acceptable).

I believe this will be a relatively significant change, so from a review standpoint, it would be preferable if the PR could focus on changes to the NN module and keep other changes to a minimum.

[sotetsuk]

At the very least, there will be no members in our team working on this issue by the end of the year, so it's fine for you to proceed at your own pace 👍