Consider make flax an optional dependency

[jjyyxx]

It seems that the sole purpose of the flax dependency is to support conversion between haiku and flax, but flax's dependencies are considerably heavier compared to the rest of haiku. I think it is a good idea to make flax an extra dependency, and avoid importing haiku.experimental.flax eagerly.

[tomhennigan]

Thanks for letting us know! I'll take a look into whether we can do this for the next release.